# -*- coding: utf-8 -*-
"""Collects osquery from the command line and the local file system."""

import json
import os
from typing import Optional, List, Callable

from dftimewolf.lib import module
from dftimewolf.lib.containers import containers
from dftimewolf.lib.modules import manager as modules_manager
from dftimewolf.lib import cache
from dftimewolf.lib import telemetry
from dftimewolf.lib.containers import manager as container_manager


_ALL_PLATFORMS = ['darwin', 'freebsd', 'linux', 'windows']


class OsqueryCollector(module.BaseModule):
  """Osquery collector that creates OsqueryQuery containers.

  Attributes:
    osqueries (List[containers.OsqueryQuery]): list of osquery containers.
    configuration_content (str): the JSON configuration content.
    file_collection_columns (List[str]): The list of file collection
        columns.
  """

  def __init__(self,
               name: str,
               container_manager_: container_manager.ContainerManager,
               cache_: cache.DFTWCache,
               telemetry_: telemetry.BaseTelemetry,
               publish_message_callback: Callable[[str, str, bool], None]):
    """Initializes a local file system collector.

    Args:
      name: The modules runtime name.
      container_manager_: A common container manager object.
      cache_: A common DFTWCache object.
      telemetry_: A common telemetry collector object.
      publish_message_callback: A callback to send modules messages to.
    """
    super().__init__(name=name,
                     cache_=cache_,
                     container_manager_=container_manager_,
                     telemetry_=telemetry_,
                     publish_message_callback=publish_message_callback)
    self.osqueries: List[containers.OsqueryQuery] = []
    self.configuration_content: str = ''
    self.file_collection_columns: List[str] = []

  def _ValidateOsquery(self, query: str) -> bool:
    """Validate Osquery query.

    Args:
      query (str): osquery query.

    Returns:
      True if the query appears to be valid, False otherwise
    """
    # TODO(sydp): add more checks.
    return (
        query.strip().upper().startswith('SELECT ')
        and query.strip().endswith(';'))

  def _ParsePlatforms(self, platforms: str) -> List[str]:
    """Parse and normalise the platforms value from an osquery pack.

    Arguments:
      platforms: the platforms value from an osquery pack.

    Returns:
      a list of operating system platforms.  Valid values in the list are
      'darwin', 'freebsd', 'linux', 'windows'
    """
    if not platforms:
      return []

    unique_platforms = set()
    for platform in platforms.split(','):
      platform = platform.strip()
      if platform in ('all', 'any'):
        unique_platforms.update(_ALL_PLATFORMS)
      elif platform == 'posix':
        unique_platforms.update(['darwin', 'freebsd', 'linux'])
      elif platform in _ALL_PLATFORMS:
        unique_platforms.add(platform)
      else:
        self.ModuleError(
            f'Unexpected value {platform} in platform value.', critical=True)

    return list(unique_platforms)

  def _LoadOsqueryPackToState(self, path: str) -> None:
    """Loads osquery from an osquery pack file and creates Osquery containers.

    Args:
      path: the path to the JSON file.
    """
    with open(path, mode='r') as fd:
      global_platform = []

      query_pack = json.load(fd)

    # A 'global' platform value can be set at the root level
    if 'platform' in query_pack:
      global_platform = self._ParsePlatforms(query_pack.get('platform'))

    for num, (name, entry) in enumerate(query_pack.get('queries', {}).items()):
      query = entry['query']
      if not self._ValidateOsquery(query):
        self.ModuleError(
            f'Entry {num} in query pack {path} does not appear to be valid.')
        continue

      if 'platform' in entry:
        platform = self._ParsePlatforms(entry.get('platform'))
      else:
        platform = global_platform

      self.osqueries.append(
          containers.OsqueryQuery(
              query=query,
              name=name,
              description=entry.get('description', ''),
              platforms=platform,
              configuration_content=self.configuration_content,
              file_collection_columns=self.file_collection_columns))

  def _LoadTextFileToState(self, path: str) -> None:
    """Loads osquery from a text file and creates Osquery containers.

    Args:
      path: the path to the text file.
    """
    with open(path, mode='r') as fd:
      for line_number, line in enumerate(fd.readlines()):
        if self._ValidateOsquery(line):
          self.osqueries.append(
              containers.OsqueryQuery(
                  query=line,
                  name='',
                  description='',
                  platforms=None,
                  configuration_content=self.configuration_content,
                  file_collection_columns=self.file_collection_columns))
        else:
          self.ModuleError(
              f'Osquery on line {line_number} of {path} '
              'does not appear to be valid.', critical=True)

  # pylint: disable=arguments-differ
  def SetUp(
      self,
      query: str,
      paths: str,
      local_configuration_path: str = '',
      configuration_content: str = '',
      file_collection_columns: Optional[str] = None
  ) -> None:
    """Sets up the osquery to collect.

    Supported files are:
    * text files that contain one Osquery
    * json files containing an osquery pack. See https://osquery.readthedocs.io
        /en/stable/deployment/configuration/#query-packs for details and
        https://github.com/osquery/osquery/tree/master/packs for examples.

    The GRR osquery flow can also be set up to use a custom osquery
    configuration on invocation (see
    https://osquery.readthedocs.io/en/stable/deployment/configuration/)
    as a temporary file on the GRR client where the content can come from
    a file, using local_cofiguration_path, on the user's local machine or a
    string value, using configuration_content.

    GRR can also collect files based on the results of an Osquery flow using the
    file_collection_columns argument.

    Args:
      query: osquery query.
      paths: osquery filepaths.
      configuration_content: the configuration content, in JSON format.
      local_configuration_path: the path to a local osquery configuration file.
      file_collection_columns: The comma-seaparated list of file collection
          columns names.
    """
    if not query and not paths:
      self.ModuleError('Both query and paths cannot be empty.', critical=True)

    if local_configuration_path and configuration_content:
      self.ModuleError(
          'Only one configuration argument can be set.', critical=True)

    if local_configuration_path:
      with open(local_configuration_path, mode='r') as fd:
        configuration_content = fd.read()

    if configuration_content:
      try:
        content = json.loads(configuration_content)
      except json.JSONDecodeError:
        self.ModuleError(
            'Osquery configuration does not contain valid JSON.',
            critical=True)
      self.configuration_content = json.dumps(content)

    if file_collection_columns:
      self.file_collection_columns = [
          col.strip() for col in file_collection_columns.split(',')]

    if query:
      if self._ValidateOsquery(query):
        self.osqueries.append(containers.OsqueryQuery(
            query=query,
            configuration_content=self.configuration_content,
            file_collection_columns=self.file_collection_columns))
      else:
        self.ModuleError(
            'Osquery parameter not set or does not appear to be valid.',
            critical=True)

    if paths:
      split_paths = [path.strip() for path in paths.split(',')]

      for path in split_paths:
        if not os.path.exists(path):
          self.logger.warning(f'Path {path} does not exist.')
          continue
        if os.path.splitext(path)[1] == '.json':
          self._LoadOsqueryPackToState(path)
        else:
          self._LoadTextFileToState(path)

    if not self.osqueries:
      self.ModuleError(
          message='No valid osquery collected.', critical=True)

  def Process(self) -> None:
    """Collects osquery from the command line and local file system."""
    for osquery in self.osqueries:
      self.StoreContainer(osquery)


modules_manager.ModulesManager.RegisterModule(OsqueryCollector)
