#!/usr/bin/env python
# Copyright 2012 Google Inc. All Rights Reserved.

"""Utils exporting data from AFF4 to the rest of the world."""

import os
import Queue
import stat
import time

import logging

from grr.lib import aff4
from grr.lib import client_index
from grr.lib import rdfvalue
from grr.lib import serialize
from grr.lib import threadpool
from grr.lib import utils
from grr.lib.aff4_objects import aff4_grr
from grr.lib.flows.general import file_finder
from grr.lib.rdfvalues import client as rdf_client
from grr.lib.rdfvalues import flows as rdf_flows


BUFFER_SIZE = 16 * 1024 * 1024


def GetAllClients(token=None):
  """Return a list of all client urns."""
  index = aff4.FACTORY.Create(
      client_index.MAIN_INDEX, aff4_type="ClientIndex",
      mode="rw", object_exists=True, token=token)

  return index.LookupClients(["."])


class IterateAllClientUrns(object):
  """Class to iterate over all URNs."""

  THREAD_POOL_NAME = "ClientUrnIter"
  QUEUE_TIMEOUT = 30

  def __init__(self, func=None, max_threads=10, token=None):
    """Iterate over all clients in a threadpool.

    Args:
      func: A function to call with each client urn.
      max_threads: Number of threads to use.
      token: Auth token.

    Raises:
      RuntimeError: If function not specified.
    """
    self.thread_pool = threadpool.ThreadPool.Factory(self.THREAD_POOL_NAME,
                                                     max_threads)
    self.thread_pool.Start()
    self.token = token
    self.func = func
    self.broken_subjects = []  # Entries that are broken or fail to run.

    self.out_queue = Queue.Queue()

  def GetInput(self):
    """Yield client urns."""
    clients = GetAllClients(token=self.token)
    logging.debug("Got %d clients", len(clients))
    return clients

  def Run(self):
    """Run the iteration."""
    count = 0
    for count, input_data in enumerate(self.GetInput()):
      if count % 2000 == 0:
        logging.debug("%d processed.", count)
      args = (input_data, self.out_queue, self.token)
      self.thread_pool.AddTask(target=self.IterFunction, args=args,
                               name=self.THREAD_POOL_NAME)

    while count >= 0:
      try:
        # We only use the timeout to wait if we got to the end of the Queue but
        # didn't process everything yet.
        out = self.out_queue.get(timeout=self.QUEUE_TIMEOUT, block=True)
        if out:
          yield out
          count -= 1
      except Queue.Empty:
        break

    # Join and stop to clean up the threadpool.
    self.thread_pool.Stop()

  def IterFunction(self, *args):
    """Function to run on each input. This can be overridden."""
    self.func(*args)


class IterateAllClients(IterateAllClientUrns):
  """Class to iterate over all GRR Client objects."""

  def __init__(self, max_age, client_chunksize=25, **kwargs):
    """Iterate over all clients in a threadpool.

    Args:
      max_age: Maximum age in seconds of clients to check.
      client_chunksize: A function to call with each client urn.
      **kwargs: Arguments passed to init.
    """
    super(IterateAllClients, self).__init__(**kwargs)
    self.client_chunksize = client_chunksize
    self.max_age = max_age

  def GetInput(self):
    """Yield client urns."""
    client_list = GetAllClients(token=self.token)
    logging.debug("Got %d clients", len(client_list))
    for client_group in utils.Grouper(client_list, self.client_chunksize):
      for fd in aff4.FACTORY.MultiOpen(client_group, mode="r",
                                       aff4_type="VFSGRRClient",
                                       token=self.token):
        if isinstance(fd, aff4_grr.VFSGRRClient):
          # Skip if older than max_age
          oldest_time = (time.time() - self.max_age) * 1e6
        if fd.Get(aff4.VFSGRRClient.SchemaCls.PING) >= oldest_time:
          yield fd


def DownloadFile(file_obj, target_path, buffer_size=BUFFER_SIZE):
  """Download an aff4 file to the local filesystem overwriting it if it exists.

  Args:
    file_obj: An aff4 object that supports the file interface (Read, Seek)
    target_path: Full path of file to write to.
    buffer_size: Read in chunks this size.
  """
  logging.info(u"Downloading: %s to: %s", file_obj.urn, target_path)

  target_file = open(target_path, "w")
  file_obj.Seek(0)
  count = 0

  data_buffer = file_obj.Read(buffer_size)
  while data_buffer:
    target_file.write(data_buffer)
    data_buffer = file_obj.Read(buffer_size)
    count += 1
    if not count % 3:
      logging.debug(u"Downloading: %s: %s done", file_obj.urn,
                    utils.FormatNumberAsString(count * buffer_size))
  target_file.close()


def RecursiveDownload(dir_obj, target_dir, max_depth=10, depth=1,
                      overwrite=False, max_threads=10):
  """Recursively downloads a file entry to the target path.

  Args:
    dir_obj: An aff4 object that contains children.
    target_dir: Full path of the directory to write to.
    max_depth: Depth to download to. 1 means just the directory itself.
    depth: Current depth of recursion.
    overwrite: Should we overwrite files that exist.
    max_threads: Use this many threads to do the downloads.
  """
  if (not isinstance(dir_obj, aff4.AFF4Volume) or
      isinstance(dir_obj, aff4.HashImage)):
    return

  # Reuse the same threadpool as we call recursively.
  thread_pool = threadpool.ThreadPool.Factory("Downloader", max_threads)
  thread_pool.Start()

  for sub_file_entry in dir_obj.OpenChildren():
    path_elements = [target_dir]
    sub_target_dir = u"/".join(path_elements)
    try:
      # Any file-like object with data in AFF4 should inherit AFF4Stream.
      if isinstance(sub_file_entry, aff4.AFF4Stream):
        args = (sub_file_entry.urn, sub_target_dir, sub_file_entry.token,
                overwrite)
        thread_pool.AddTask(target=CopyAFF4ToLocal, args=args,
                            name="Downloader")
      elif "Container" in sub_file_entry.behaviours:
        if depth >= max_depth:  # Don't go any deeper.
          continue
        try:
          os.makedirs(sub_target_dir)
        except OSError:
          pass
        RecursiveDownload(sub_file_entry, sub_target_dir, overwrite=overwrite,
                          depth=depth + 1)
    except IOError:
      logging.exception("Unable to download %s", sub_file_entry.urn)
    finally:
      sub_file_entry.Close()

  # Join and stop the threadpool.
  if depth <= 1:
    thread_pool.Stop()


def DownloadCollection(coll_path, target_path, token=None, overwrite=False,
                       dump_client_info=False, flatten=False,
                       max_threads=10):
  """Iterate through a Collection object downloading all files.

  Args:
    coll_path: Path to an AFF4 collection.
    target_path: Base directory to write to.
    token: Token for access.
    overwrite: If True, overwrite existing files.
    dump_client_info: If True, this will detect client paths, and dump a yaml
      version of the client object to the root path. This is useful for seeing
      the hostname/users of the machine the client id refers to.
    flatten: If True, produce a "files" flat folder with links to all the found
             files.
    max_threads: Use this many threads to do the downloads.
  """
  completed_clients = set()
  try:
    coll = aff4.FACTORY.Open(coll_path, aff4_type="RDFValueCollection",
                             token=token)
  except IOError:
    logging.error("%s is not a valid collection. Typo? "
                  "Are you sure something was written to it?", coll_path)
    return
  thread_pool = threadpool.ThreadPool.Factory("Downloader", max_threads)
  thread_pool.Start()

  logging.info("Expecting to download %s files", coll.size)

  # Collections can include anything they want, but we only handle RDFURN and
  # StatEntry entries in this function.
  for grr_message in coll:
    source = None
    # If a raw message, work out the type.
    if isinstance(grr_message, rdf_flows.GrrMessage):
      source = grr_message.source
      grr_message = grr_message.payload

    # Collections can contain AFF4ObjectSummary objects which encapsulate
    # RDFURNs and StatEntrys.
    if isinstance(grr_message, rdf_client.AFF4ObjectSummary):
      urn = grr_message.urn
    elif isinstance(grr_message, rdfvalue.RDFURN):
      urn = grr_message
    elif isinstance(grr_message, rdf_client.StatEntry):
      urn = rdfvalue.RDFURN(grr_message.aff4path)
    elif isinstance(grr_message, file_finder.FileFinderResult):
      urn = rdfvalue.RDFURN(grr_message.stat_entry.aff4path)
    elif isinstance(grr_message, rdfvalue.RDFBytes):
      try:
        os.makedirs(target_path)
      except OSError:
        pass
      try:
        # We just dump out bytes and carry on.
        client_id = source.Split()[0]
        with open(os.path.join(target_path, client_id), "wb") as fd:
          fd.write(str(grr_message))
      except AttributeError:
        pass
      continue
    else:
      continue

    # Handle dumping client info, but only once per client.
    client_id = urn.Split()[0]
    re_match = aff4.AFF4Object.VFSGRRClient.CLIENT_ID_RE.match(client_id)
    if dump_client_info and re_match and client_id not in completed_clients:
      args = (rdf_client.ClientURN(client_id), target_path, token, overwrite)
      thread_pool.AddTask(target=DumpClientYaml, args=args,
                          name="ClientYamlDownloader")
      completed_clients.add(client_id)

    # Now queue downloading the actual files.
    args = (urn, target_path, token, overwrite)
    if flatten:
      target = CopyAndSymlinkAFF4ToLocal
    else:
      target = CopyAFF4ToLocal
    thread_pool.AddTask(target=target, args=args, name="Downloader")

  # Join and stop the threadpool.
  thread_pool.Stop()


def CopyAFF4ToLocal(aff4_urn, target_dir, token=None, overwrite=False):
  """Copy an AFF4 object that supports a read interface to local filesystem.

  Args:
    aff4_urn: URN of thing to copy.
    target_dir: Directory to copy the file to.
    token: Auth token.
    overwrite: If True overwrite the file if it exists.

  Returns:
    If aff4_urn points to a file, returns path to the downloaded file.
    Otherwise returns None.

  By default file will only be overwritten if file size differs.
  """
  try:
    fd = aff4.FACTORY.Open(aff4_urn, token=token)
    filepath = os.path.join(target_dir, fd.urn.Path()[1:])

    # If urn points to a directory, just create it.
    if isinstance(fd, aff4.VFSDirectory):
      try:
        os.makedirs(filepath)
      except OSError:
        pass

      return None
    # If urn points to a file, download it.
    elif isinstance(fd, aff4.AFF4Stream):
      if not os.path.isfile(filepath):
        try:
          # Ensure directory exists.
          os.makedirs(os.path.dirname(filepath))
        except OSError:
          pass
        DownloadFile(fd, filepath)
      elif (os.stat(filepath)[stat.ST_SIZE] != fd.Get(fd.Schema.SIZE) or
            overwrite):
        # We should overwrite because user said, or file sizes differ.
        DownloadFile(fd, filepath)
      else:
        logging.info("File %s exists, skipping", filepath)

      return filepath
    else:
      raise RuntimeError("Opened urn is neither a downloaded file nor a "
                         "directory: %s" % aff4_urn)

  except IOError as e:
    logging.exception("Failed to read %s due to %s", aff4_urn, e)
    raise


def CopyAndSymlinkAFF4ToLocal(aff4_urn, target_dir, token=None,
                              overwrite=False):
  path = CopyAFF4ToLocal(aff4_urn, target_dir, token=token,
                         overwrite=overwrite)
  if path:
    files_output_dir = os.path.join(target_dir, "files")
    try:
      os.makedirs(files_output_dir)
    except OSError:
      pass

    unique_name = "_".join(aff4_urn.Split())
    symlink_path = os.path.join(files_output_dir, unique_name)
    try:
      os.symlink(path, symlink_path)
    except OSError:
      logging.exception("Can't create symlink to a file: %s -> %s",
                        symlink_path, path)


def DumpClientYaml(client_urn, target_dir, token=None, overwrite=False):
  """Dump a yaml file containing client info."""
  fd = aff4.FACTORY.Open(client_urn, "VFSGRRClient", token=token)
  dirpath = os.path.join(target_dir, fd.urn.Split()[0])
  try:
    # Due to threading this can actually be created by another thread.
    os.makedirs(dirpath)
  except OSError:
    pass
  filepath = os.path.join(dirpath, "client_info.yaml")
  if not os.path.isfile(filepath) or overwrite:
    with open(filepath, "w") as out_file:
      out_file.write(serialize.YamlDumper(fd))
