#!/usr/bin/env python
"""A library for check-specific tests."""
import collections
import os
import StringIO


import yaml

from grr import config
from grr.lib import registry
from grr.lib import type_info
from grr.lib.rdfvalues import anomaly as rdf_anomaly
from grr.lib.rdfvalues import client as rdf_client
from grr.lib.rdfvalues import paths as rdf_paths
from grr.parsers import linux_service_parser
from grr.server.checks import checks
from grr.server.checks import filters
from grr.server.checks import hints
from grr.test_lib import test_lib


class HostCheckTest(test_lib.GRRBaseTest):
  """The base class for host check tests."""
  __metaclass__ = registry.MetaclassRegistry

  loaded_checks = None

  @classmethod
  def LoadCheck(cls, cfg_file, *check_ids):
    """Loads checks from a file once per Test class.

    LoadCheck will read a file containing a check configuration and instantiate
    the checks from it. Specific checks can be selected by providing the check
    ids that should be loaded from the file.

    Checks are stored as a class attribute to prevent re-loading as each test
    method is set up.

    Args:
      cfg_file: A path to the file that should be read.
      *check_ids: A list of check ids that should be loaded from the file.

    Returns:
      The loaded check objects.
    """
    if HostCheckTest.loaded_checks is None:
      HostCheckTest.loaded_checks = {}

    cfg = os.path.join(config.CONFIG["Test.srcdir"], "grr", "checks", cfg_file)
    if check_ids:
      key = "%s:%s" % (cfg, ",".join(check_ids))
      if key in HostCheckTest.loaded_checks:
        return HostCheckTest.loaded_checks[key]
      loaded = []
      for chk_id in check_ids:
        loaded.append(checks.LoadCheckFromFile(cfg, chk_id))
      HostCheckTest.loaded_checks[key] = loaded
      return loaded
    else:
      key = "%s:*" % cfg_file
      if key in HostCheckTest.loaded_checks:
        return HostCheckTest.loaded_checks[key]
      else:
        result = checks.LoadChecksFromFiles([cfg])
        HostCheckTest.loaded_checks[key] = result
        return result

  def TestDataPath(self, file_name):
    """Generates a full path to the test data."""
    path = os.path.join(config.CONFIG["Test.data_dir"], file_name)
    if not os.path.isfile(path):
      raise test_lib.Error("Missing test data: %s" % file_name)
    return path

  @staticmethod
  def SetKnowledgeBase(fqdn="test.example.com", host_os="Linux",
                       host_data=None):
    """Generates a KnowledgeBase entry in the host_data used by checks."""
    if not host_data:
      host_data = {}
    host_data["KnowledgeBase"] = rdf_client.KnowledgeBase(fqdn=fqdn, os=host_os)
    return host_data

  def SetArtifactData(self, anomaly=None, parsed=None, raw=None, results=None):
    """Adds data in the format required by host_data."""
    if not results:
      results = {"ANOMALY": [], "PARSER": [], "RAW": []}
    results["ANOMALY"].extend(anomaly or [])
    results["PARSER"].extend(parsed or [])
    results["RAW"].extend(raw or [])
    return results

  def AddData(self, parser, *args, **kwargs):
    """Initialize the parser and add parsed data to host_data."""
    return [parser().Parse(*args, **kwargs)]

  def AddListener(self, ip, port, family="INET", sock_type="SOCK_STREAM"):
    """Create a network connection."""
    conn = rdf_client.NetworkConnection()
    conn.state = "LISTEN"
    conn.family = family
    conn.type = sock_type
    conn.local_address = rdf_client.NetworkEndpoint(ip=ip, port=port)
    return conn

  def RunChecks(self, host_data, labels=None, restrict_checks=None):
    """Runs the registered checks against the provided host data.

    Args:
      host_data: A dictionary of artifact_names and results. Results are, in
        turn, a dictionary of {'ANOMALY': [], 'PARSED': [], 'RAW': []} items.
      labels: Additional labels attached to the host.
      restrict_checks: A list specifying a subset of check_ids to run.

    Returns:
      An iterator of check results.
    """
    return {
        r.check_id: r
        for r in checks.CheckHost(
            host_data, labels=labels, restrict_checks=restrict_checks)
    }

  def GetCheckErrors(self, check_spec):
    """Collect errors generated by host checking tools."""
    errors = []
    try:
      check = checks.Check(**check_spec)
      check.Validate()
    except (checks.Error, filters.Error, hints.Error, type_info.Error) as e:
      errors.append(str(e))
    except Exception as e:
      # TODO(user): More granular exception handling.
      errors.append("Unknown error %s: %s" % (type(e), e))
    return errors

  def CreateStat(self, path, uid=0, gid=0, mode=0o0100640):
    """Given path, uid, gid and file mode, this returns a StatEntry."""
    pathspec = rdf_paths.PathSpec(path=path, pathtype="OS")
    return rdf_client.StatEntry(
        pathspec=pathspec, st_uid=uid, st_gid=gid, st_mode=mode)

  def _AddToHostData(self, host_data, artifact, data, parser):
    """Parse raw data collected for an artifact into the host_data table."""
    if type(data) != dict:
      raise test_lib.Error("Data for %s is not of type dictionary." % artifact)

    rdfs = []
    stats = []
    files = []
    for path, lines in data.items():
      stat = self.CreateStat(path)
      stats.append(stat)
      file_obj = StringIO.StringIO(lines)
      files.append(file_obj)
      if not parser.process_together:
        rdfs.extend(list(parser.Parse(stat, file_obj, None)))
    if parser.process_together:
      rdfs = list(parser.ParseMultiple(stats, files, None))
    host_data[artifact] = self.SetArtifactData(
        anomaly=[a for a in rdfs if isinstance(a, rdf_anomaly.Anomaly)],
        parsed=[r for r in rdfs if not isinstance(r, rdf_anomaly.Anomaly)],
        raw=stats,
        results=host_data.get(artifact))
    return host_data

  def GenResults(self, artifact_list, sources_list, parser_list=None):
    """Given a list of artifacts, sources and parsers, will RunChecks on them.

    Sample: (["Artifact1", "Artifact2"], [artifact1_data, artifact2_data],
             [config_file.artifact1_Parser(), config_file.artifact2_Parser()]

    artifact1_Parser().parse will run on artifact1_data & parsed results, along
    with raw and anomalies will be inserted into the host_data["Artifact1"].

    artifact2_Parser().parse will run on artifact2_data & parsed results, along
    with raw and anomalies will be inserted into the host_data["Artifact2"].

    Once artifacts added to host_data, loaded checks will be run against it.

    Args:
      artifact_list: list of artifacts to add to host_data for running checks
      sources_list: list of dictionaries containing file names and file data.
        If parser_list is empty then sources_list must contain a list of lists
        containing StatEntry or lists of other raw artifact data.
      parser_list: list of parsers to apply to file data from sources_list.
        This can be empty if no parser is to be applied.

    Returns:
      CheckResult containing any findings in sources_list against loaded checks.

    Raises:
      Error: When there are issues with the passed input.
    """
    if parser_list is None:
      parser_list = [None] * len(artifact_list)

    # make sure all vars are lists
    if any(
        type(lst) != list
        for lst in [artifact_list, sources_list, parser_list]):
      raise test_lib.Error("All inputs are not lists.")
    # make sure all lists are of equal length
    if any(
        len(lst) != len(artifact_list) for lst in [sources_list, parser_list]):
      raise test_lib.Error("All lists are not of the same length.")

    host_data = self.SetKnowledgeBase()
    for artifact, sources, parser in zip(artifact_list, sources_list,
                                         parser_list):
      if parser is None:
        host_data[artifact] = self.SetArtifactData(
            raw=sources, results=host_data.get(artifact))
      else:
        host_data = self._AddToHostData(host_data, artifact, sources, parser)
    return self.RunChecks(host_data)

  def GenProcessData(self, processes, **kwargs):
    """Create some process-based host data."""
    host_data = self.SetKnowledgeBase(**kwargs)
    data = []
    for (name, pid, cmdline) in processes:
      data.append(rdf_client.Process(name=name, pid=pid, cmdline=cmdline))
    # ListProcessesGrr is a flow artifact, thus it needs stored as raw.
    host_data["ListProcessesGrr"] = self.SetArtifactData(raw=data)
    return host_data

  @classmethod
  def _GenFileData(cls, paths, data, stats=None, files=None, modes=None):
    """Generate a tuple of list of stats and list of file contents."""
    if stats is None:
      stats = []
    if files is None:
      files = []
    if modes is None:
      modes = {}
    modes.setdefault("st_uid", 0)
    modes.setdefault("st_gid", 0)
    modes.setdefault("st_mode", 0o0100644)
    for path in paths:
      p = rdf_paths.PathSpec(path=path, pathtype="OS")
      stats.append(rdf_client.StatEntry(pathspec=p, **modes))
    for val in data:
      files.append(StringIO.StringIO(val))
    return stats, files

  @classmethod
  def GenStatFileData(cls, data, modes=None):
    """Gen a tuple of list of stats and list of file contents from a dict."""
    paths = []
    contents = []
    for path, content in data.iteritems():
      paths.append(path)
      contents.append(content)
    return cls._GenFileData(paths, contents, modes=modes)

  def GenFileData(self, artifact, data, parser=None, modes=None, include=None):
    """Create a set of host_data results based on file parser results.

    Creates a host data result populated with a knowledge base, then processes
    each piece of host data as if it were file results using a parser. Specific
    stat values can be provided to the parser, if required, so that permissions,
    ownership and file types can be defined.

    Args:
      artifact: The artifact name that generated data will be mapped to.
      data: A dictionary of pathnames and data strings. The data strings are
        converted into file objects for the parser.
      parser: The FileParser that processes the data (and stats)
      modes: A dictionary of pathnames and stat values. Stat values are a dict
        of {st_uid: int, st_gid: int, st_mode: oct} entries.
      include: A list of pathnames to include in processing. If not provided,
        all paths are parsed.

    Returns:
      the host_data map populated with a knowledge base and artifact data.

    Raises:
      Error: When the handed parser was not initialized.
    """
    host_data = self.SetKnowledgeBase()
    if not parser:
      raise test_lib.Error("Test method requires an initialized parser.")
    if not modes:
      modes = {}
    kb = host_data["KnowledgeBase"]
    files = []
    rdfs = []
    stats = []
    for path in data:
      if include and path not in include:
        continue
      file_modes = modes.get(path, {})
      stats, files = self._GenFileData(
          [path], [data[path]], stats=stats, files=files, modes=file_modes)
    if parser.process_together:
      rdfs = list(parser.ParseMultiple(stats, files, kb))
    else:
      for stat, file_obj in zip(stats, files):
        # Make sure the parser result is iterable.
        rslt = parser.Parse(stat, file_obj, kb) or []
        rdfs.extend(rslt)
    anomaly = [a for a in rdfs if isinstance(a, rdf_anomaly.Anomaly)]
    parsed = [r for r in rdfs if not isinstance(r, rdf_anomaly.Anomaly)]
    host_data[artifact] = self.SetArtifactData(
        parsed=parsed, anomaly=anomaly, raw=stats)
    return host_data

  def GenSysVInitData(self, links):
    """Create some Sys V init host data."""
    return self.GenFileData(
        artifact="LinuxServices",
        data={x: ""
              for x in links},
        parser=linux_service_parser.LinuxSysVInitParser(),
        modes={x: {
            "st_mode": 41471
        }
               for x in links})

  # The assert methods

  def assertRanChecks(self, check_ids, results):
    """Tests that the specified checks were run."""
    self.assertTrue(set(check_ids).issubset(set(results.keys())))

  def assertChecksNotRun(self, check_ids, results):
    """Tests that the specified checks were not run."""
    self.assertFalse(set(check_ids).intersection(set(results.keys())))

  def assertResultEqual(self, rslt1, rslt2):
    """Tests whether two check results are identical."""
    # Build a map of anomaly symptoms to findings.
    if rslt1.check_id != rslt2.check_id:
      self.fail("Check IDs differ: %s vs %s" % (rslt1.check_id, rslt2.check_id))

    # Quick check to see if anomaly counts are the same and they have the same
    # ordering, using symptoms as a measure.
    rslt1_anoms = {}
    for a in rslt1.anomaly:
      anoms = rslt1_anoms.setdefault(a.symptom, [])
      anoms.extend(a.finding)
    rslt2_anoms = {}
    for a in rslt2.anomaly:
      anoms = rslt2_anoms.setdefault(a.symptom, [])
      anoms.extend(a.finding)

    self.assertItemsEqual(rslt1_anoms, rslt2_anoms,
                          "Results have different anomaly items.:\n%s\n%s" %
                          (rslt1_anoms.keys(), rslt2_anoms.keys()))

    # Now check that the anomalies are the same, modulo newlines.
    for symptom, findings in rslt1_anoms.iteritems():
      rslt1_found = [f.strip() for f in findings]
      rslt2_found = [f.strip() for f in rslt2_anoms[symptom]]
      self.assertItemsEqual(rslt1_found, rslt2_found)

  def assertIsCheckIdResult(self, rslt, expected):
    """Tests if a check has the expected check_id."""
    self.assertIsInstance(rslt, checks.CheckResult)
    self.assertEqual(expected, rslt.check_id)

  def assertValidCheck(self, check_spec):
    """Tests if a check definition generates structural errors."""
    errors = self.GetCheckErrors(check_spec)
    if errors:
      self.fail("\n".join(errors))

  def assertValidCheckFile(self, path):
    """Tests whether a check definition has a valid configuration."""
    # Figure out the relative path of the check files.
    prefix = os.path.commonprefix(config.CONFIG["Checks.config_dir"])
    relpath = os.path.relpath(path, prefix)
    # If the config can't load fail immediately.
    try:
      configs = checks.LoadConfigsFromFile(path)
    except yaml.error.YAMLError as e:
      self.fail("File %s could not be parsed: %s\n" % (relpath, e))
    # Otherwise, check all the configs and pass/fail at the end.
    errors = collections.OrderedDict()
    for check_id, check_spec in configs.iteritems():
      check_errors = self.GetCheckErrors(check_spec)
      if check_errors:
        msg = errors.setdefault(relpath, ["check_id: %s" % check_id])
        msg.append(check_errors)
    if errors:
      message = ""
      for k, v in errors.iteritems():
        message += "File %s errors:\n" % k
        message += "  %s\n" % v[0]
        for err in v[1]:
          message += "    %s\n" % err
      self.fail(message)

  def _HasSymptom(self, anomalies, sym):
    """Tests if one or more anomalies contain the expected symptom string."""
    if sym is None:
      return True
    rslts = {rslt.symptom: rslt for rslt in anomalies}
    rslt = rslts.get(sym)
    # Anomalies evaluate false if there are no finding strings.
    self.assertTrue(rslt is not None,
                    "Didn't get expected symptom string '%s' in '%s'" %
                    (sym, ",".join(rslts)))

  def _GetFindings(self, anomalies, sym):
    """Generate a set of findings from anomalies that match the symptom."""
    result = set()
    for anomaly in anomalies:
      if anomaly.symptom == sym:
        result.update(set(anomaly.finding))
    return result

  def _MatchFindings(self, expected, found):
    """Check that every expected finding is a substring of a found finding."""
    matched_so_far = set()
    for finding_str in expected:
      no_match = True
      for found_str in found:
        if finding_str in found_str:
          matched_so_far.add(found_str)
          no_match = False
          break
      if no_match:
        return False
    # If we got here, all expected's match at least one item.
    # Now check if every item in found was matched at least once.
    # If so, everything is as expected, If not, Badness.
    if not matched_so_far.symmetric_difference(found):
      return True

  def assertCheckDetectedAnom(self, check_id, results, sym=None, findings=None):
    """Assert a check was performed and specific anomalies were found.

    Results may contain multiple anomalies. The check will hold true if any
    one of them matches. As some results can contain multiple anomalies we
    will need to make sure the right anomalies are selected.

    If an symptom is provided, look for anomalies that matches the
    expression string and use those. Otherwise, all anomalies in the
    check should be used.

    If finding strings are provided, the check tests if the substring is present
    in the findings of the anomalies that are selected for testing. If the
    finding results can have variable ordering, use a substring that will remain
    constant for each finding.

    Args:
      check_id: The check_id as a string.
      results: A dictionary of check results, mapped to check_ids
      sym: An symptom string. This is the "title" of an advisory.
      findings: A list of finding strings that should be present in the findings
        of the selected anomaly.

    Returns:
      True if tests have succeeded and no further processing is required.
    """
    chk = results.get(check_id)
    self.assertTrue(chk is not None, "check %s did not run" % check_id)
    # Checks return true if there were anomalies.
    self.assertTrue(chk, "check %s did not generate anomalies" % check_id)
    # If sym or results are passed as args, look for anomalies with these
    # values.
    self._HasSymptom(chk.anomaly, sym)
    if findings is None:
      # We are not expecting to match on findings, so skip checking them.
      return True
    findings = set(findings)
    found = self._GetFindings(chk.anomaly, sym)
    if self._MatchFindings(findings, found):
      # Everything matches, and nothing unexpected, so all is good.
      return True
    # If we have made it here, we have the expected symptom but
    # the findings didn't match up.
    others = "\n".join([str(a) for a in chk.anomaly])
    self.fail("Findings don't match for symptom '%s':\nExpected:\n  %s\n"
              "Got:\n  %s\nFrom:\n%s" % (sym, ", ".join(findings),
                                         ", ".join(found), others))

  def assertCheckUndetected(self, check_id, results):
    """Assert a check_id was performed, and resulted in no anomalies."""
    if not isinstance(results, collections.Mapping):
      self.fail("Invalid arg, %s should be dict-like.\n" % type(results))
    if check_id not in results:
      self.fail("Check %s was not performed.\n" % check_id)
    # A check result will evaluate as True if it contains an anomaly.
    if results.get(check_id):
      self.fail("Check %s unexpectedly produced an anomaly.\nGot: %s\n" %
                (check_id, results.get(check_id).anomaly))

  def assertChecksUndetected(self, check_ids, results):
    """Assert multiple check_ids were performed & they produced no anomalies."""
    for check_id in check_ids:
      self.assertCheckUndetected(check_id, results)
