"""Helper to create filters based on forensic artifact definitions."""

import os

from artifacts import definitions as artifact_types

from dfvfs.helpers import file_system_searcher as dfvfs_file_system_searcher

from dfwinreg import registry_searcher as dfwinreg_registry_searcher

from plaso.engine import artifacts_trie
from plaso.engine import logger
from plaso.engine import path_helper


class ArtifactDefinitionsFiltersHelper:
  """Helper to create collection filters based on artifact definitions.

  Builds collection filters from forensic artifact definitions.

  For more information about Forensic Artifacts see:
  https://github.com/ForensicArtifacts/artifacts/blob/main/docs/Artifacts%20definition%20format%20and%20style%20guide.asciidoc

  Attributes:
    artifacts_trie (ArtifactsTrie): Trie structure for storing artifact
        definition paths.
    file_system_artifact_names (set[str]): names of artifacts definitions that
        generated file system find specifications.
    file_system_find_specs (list[dfvfs.FindSpec]): file system find
        specifications of paths to include in the collection.
    registry_artifact_names (set[str]): names of artifacts definitions that
        generated Windows Registry find specifications.
    registry_find_specs (list[dfwinreg.FindSpec]): Windows Registry find
        specifications.
    registry_find_specs_artifact_names (list[str]): Windows Registry artifact
        names corresponding to the find specifications.
  """

  _COMPATIBLE_REGISTRY_KEY_PATH_PREFIXES = frozenset([
      'HKEY_CURRENT_USER',
      'HKEY_LOCAL_MACHINE\\SYSTEM',
      'HKEY_LOCAL_MACHINE\\SOFTWARE',
      'HKEY_LOCAL_MACHINE\\SAM',
      'HKEY_LOCAL_MACHINE\\SECURITY',
      'HKEY_USERS'])

  def __init__(self, artifacts_registry):
    """Initializes an artifact definitions filters helper.

    Args:
      artifacts_registry (artifacts.ArtifactDefinitionsRegistry): artifact
          definitions registry.
    """
    super().__init__()
    self._artifacts_registry = artifacts_registry

    self.artifacts_trie = artifacts_trie.ArtifactsTrie()
    self.file_system_artifact_names = set()
    self.file_system_find_specs = []
    self.registry_artifact_names = set()
    self.registry_find_specs = []
    self.registry_find_specs_artifact_names = []

  def _BuildFindSpecsFromArtifact(
      self, definition, environment_variables, user_accounts,
      enable_artifacts_map=False,
      original_registry_artifact_filter_names=None):
    """Builds find specifications from an artifact definition.

    Args:
      definition (artifacts.ArtifactDefinition): artifact definition.
      environment_variables (list[EnvironmentVariableArtifact]): environment
          variables.
      user_accounts (list[UserAccountArtifact]): user accounts.
      enable_artifacts_map (Optional[bool]): True if the artifacts path map
          should be generated. Defaults to False.
      original_registry_artifact_filter_names (Optional[set[str]]): Set of
          original registry filter names, used in case Windows Registry hive
          files are being requested as a result of a previous filter.

    Returns:
      list[dfvfs.FindSpec|dfwinreg.FindSpec]: dfVFS or dfWinReg find
          specifications.
    """
    find_specs = []
    for source in definition.sources:
      if source.type_indicator == artifact_types.TYPE_INDICATOR_FILE:
        for path_entry in set(source.paths):
          specifications = self._BuildFindSpecsFromFileSourcePath(
              definition.name, path_entry, source.separator,
              environment_variables, user_accounts,
              enable_artifacts_map=enable_artifacts_map,
              original_registry_artifact_filter_names=(
                  original_registry_artifact_filter_names))
          find_specs.extend(specifications)
          self.file_system_artifact_names.add(definition.name)

      elif (source.type_indicator ==
            artifact_types.TYPE_INDICATOR_WINDOWS_REGISTRY_KEY):
        for key_path in set(source.keys):
          if ArtifactDefinitionsFiltersHelper.CheckKeyCompatibility(key_path):
            specifications = self._BuildFindSpecsFromRegistrySourceKey(key_path)
            find_specs.extend(specifications)
            self.registry_artifact_names.add(definition.name)

      elif (source.type_indicator ==
            artifact_types.TYPE_INDICATOR_WINDOWS_REGISTRY_VALUE):
        # TODO: Handle Registry Values Once Supported in dfwinreg.
        # https://github.com/log2timeline/dfwinreg/issues/98

        # Use set-comprehension to create a set of the source key paths.
        key_paths = {key_value['key'] for key_value in source.key_value_pairs}
        key_paths_string = ', '.join(key_paths)

        logger.warning((
            f'Windows Registry values are not supported, extracting keys: '
            f'"{key_paths_string!s}"'))

        for key_path in key_paths:
          if ArtifactDefinitionsFiltersHelper.CheckKeyCompatibility(key_path):
            specifications = self._BuildFindSpecsFromRegistrySourceKey(key_path)
            find_specs.extend(specifications)
            self.registry_artifact_names.add(definition.name)

      elif (source.type_indicator ==
            artifact_types.TYPE_INDICATOR_ARTIFACT_GROUP):
        for name in source.names:
          specifications = self._BuildFindSpecsFromGroupName(
              name, environment_variables, user_accounts,
              enable_artifacts_map=enable_artifacts_map,
              original_registry_artifact_filter_names=(
                  original_registry_artifact_filter_names))
          find_specs.extend(specifications)

      else:
        logger.warning((
            f'Unsupported artifact definition source type: '
            f'"{source.type_indicator:s}"'))

    return find_specs

  def _BuildFindSpecsFromGroupName(
      self, group_name, environment_variables, user_accounts,
      enable_artifacts_map=False,
      original_registry_artifact_filter_names=None):
    """Builds find specifications from a artifact group name.

    Args:
      group_name (str): artifact group name.
      environment_variables (list[EnvironmentVariableArtifact]): environment
          variables.
      user_accounts (list[UserAccountArtifact]): user accounts.
      enable_artifacts_map (Optional[bool]): True if the artifacts path map
          should be generated. Defaults to False.
      original_registry_artifact_filter_names (Optional[set[str]]): Set of
          original registry filter names, used in case registry hive files
          are being requested as a result of a previous filter.

    Returns:
      list[dfwinreg.FindSpec|dfvfs.FindSpec]: find specifications or None if no
          artifact with the given name can be retrieved.
    """
    definition = self._artifacts_registry.GetDefinitionByName(group_name)
    if not definition:
      definition = self._artifacts_registry.GetDefinitionByAlias(group_name)
    if not definition:
      return None

    return self._BuildFindSpecsFromArtifact(
        definition, environment_variables, user_accounts,
        enable_artifacts_map=enable_artifacts_map,
        original_registry_artifact_filter_names=(
            original_registry_artifact_filter_names))

  def _BuildFindSpecsFromRegistrySourceKey(self, key_path):
    """Build find specifications from a Windows Registry source type.

    Args:
      key_path (str): Windows Registry key path defined by the source.

    Returns:
      list[dfwinreg.FindSpec]: find specifications for the Windows Registry
          source type.
    """
    find_specs = []
    for key_path_glob in path_helper.PathHelper.ExpandGlobStars(key_path, '\\'):
      logger.debug(f'building find spec from key path glob: {key_path_glob:s}')

      key_path_glob_upper = key_path_glob.upper()
      if key_path_glob_upper.startswith(
          'HKEY_LOCAL_MACHINE\\SYSTEM\\CURRENTCONTROLSET'):
        # Rewrite CurrentControlSet to ControlSet* for Windows NT.
        key_path_glob = ''.join([
            'HKEY_LOCAL_MACHINE\\System\\ControlSet*', key_path_glob[43:]])

      elif key_path_glob_upper.startswith('HKEY_USERS\\%%USERS.SID%%'):
        # Escaping charachter excluded from string index.
        key_path_glob = ''.join(['HKEY_CURRENT_USER', key_path_glob[24:]])

      find_spec = dfwinreg_registry_searcher.FindSpec(
          key_path_glob=key_path_glob)
      find_specs.append(find_spec)

    return find_specs

  def _BuildFindSpecsFromFileSourcePath(
      self, artifact_name, source_path, path_separator,
      environment_variables, user_accounts,
      enable_artifacts_map=False,
      original_registry_artifact_filter_names=None):
    """Builds find specifications from a file source type.

    Args:
      artifact_name (str): artifact name.
      source_path (str): file system path defined by the source.
      path_separator (str): file system path segment separator.
      environment_variables (list[EnvironmentVariableArtifact]): environment
          variables.
      user_accounts (list[UserAccountArtifact]): user accounts.
      enable_artifacts_map (Optional[bool]): True if the artifacts path map
          should be generated. Defaults to False.
      original_registry_artifact_filter_names (Optional[set[str]]): Set of
          original registry filter names, used in case Windows Registry hive
          files are being requested as a result of a previous filter.

    Returns:
      list[dfvfs.FindSpec]: find specifications for the file source type.
    """
    find_specs = []
    for path_glob in path_helper.PathHelper.ExpandGlobStars(
        source_path, path_separator):
      logger.debug(f'building find spec from path glob: {path_glob:s}')

      for path in path_helper.PathHelper.ExpandUsersVariablePath(
          path_glob, path_separator, user_accounts):
        logger.debug(f'building find spec from path: {path:s}')

        expanded_path = self._ExpandPathVariables(
            path, environment_variables, path_separator)
        if expanded_path is None:
          continue

        find_spec = self._CreateFindSpec(expanded_path, path_separator)
        if find_spec is None:
          continue

        find_specs.append(find_spec)

        if enable_artifacts_map:
          self._AddToArtifactsTrie(
              artifact_name, expanded_path,
              original_registry_artifact_filter_names,
              path_separator)

    return find_specs

  def _AddToArtifactsTrie(
      self, artifact_name, path,
      original_registry_artifact_filter_names,
      path_separator):
    """Adds a path to the artifacts trie.

    Args:
      artifact_name (str): artifact name.
      path (str): file system path.
      original_registry_artifact_filter_names (Optional[set[str]]): Set of
          original registry filter names.
      path_separator (str): path separator.
    """
    normalized_path = path.replace(path_separator, os.sep)
    self.artifacts_trie.AddPath(artifact_name, normalized_path, os.sep)
    if original_registry_artifact_filter_names:
      for name in original_registry_artifact_filter_names:
        self.artifacts_trie.AddPath(name, normalized_path, os.sep)

  def _ExpandPathVariables(self, path, environment_variables, path_separator):
    """Expands Windows paths and validates the result.

    Args:
      path (str): file system path with environment variables
      environment_variables (list[EnvironmentVariableArtifact]):
          environment variables.
      path_separator (str): file system path segment separator.

    Returns:
      str: expanded path, or None if the path is invalid.
    """
    if '%' in path:
      path = path_helper.PathHelper.ExpandWindowsPath(
          path, environment_variables)
      logger.debug(f'building find spec from expanded path: {path:s}')

    if not path.startswith(path_separator):
      logger.warning((
          f'The path filter must be defined as an absolute path: '
          f'"{path:s}"'))
      return None
    return path

  def _CreateFindSpec(self, path, path_separator):
    """Creates a dfVFS find specification.

    Args:
      path (str): Path to match.
      path_separator (str): file system path segment separator.

    Returns:
      dfvfs.FindSpec: a find specification or None if one cannot be created.
    """
    try:
      find_spec = dfvfs_file_system_searcher.FindSpec(
          case_sensitive=False, location_glob=path,
          location_separator=path_separator)
      return find_spec
    except ValueError as exception:
      logger.error((
          f'Unable to build find specification for path: "{path:s}" with '
          f'error: {exception!s}'))
      return None

  def BuildFindSpecs(
      self, artifact_filter_names, environment_variables=None,
      user_accounts=None, enable_artifacts_map=False,
      original_registry_artifact_filter_names=None):
    """Builds find specifications from artifact definitions.

    Args:
      artifact_filter_names (list[str]): names of artifact definitions that are
          used for filtering file system and Windows Registry key paths.
      environment_variables (list[EnvironmentVariableArtifact]): environment
          variables.
      user_accounts (Optional[list[UserAccountArtifact]]): user accounts.
      enable_artifacts_map (Optional[bool]): True if the artifacts path map
          should be generated. Defaults to False.
      original_registry_artifact_filter_names (Optional[set[str]]): Set of
          original registry filter names, used in case registry hive files
          are being requested as a result of a previous filter.
    """
    find_specs = {}
    for name in artifact_filter_names:
      definition = self._artifacts_registry.GetDefinitionByName(name)
      if not definition:
        definition = self._artifacts_registry.GetDefinitionByAlias(name)
      if not definition:
        logger.debug(f'undefined artifact definition: {name:s}')
        continue

      logger.debug(f'building find spec from artifact definition: {name:s}')
      artifact_find_specs = self._BuildFindSpecsFromArtifact(
          definition, environment_variables, user_accounts,
          enable_artifacts_map=enable_artifacts_map,
          original_registry_artifact_filter_names=(
              original_registry_artifact_filter_names))
      find_spec_list = find_specs.setdefault(name, [])
      find_spec_list.extend(artifact_find_specs)

    for name, find_spec_values in find_specs.items():
      for find_spec in find_spec_values:
        if isinstance(find_spec, dfvfs_file_system_searcher.FindSpec):
          self.file_system_find_specs.append(find_spec)

        elif isinstance(find_spec, dfwinreg_registry_searcher.FindSpec):
          self.registry_find_specs.append(find_spec)
          # Artifact names ordered similar to registry find specs.
          self.registry_find_specs_artifact_names.append(name)
        else:
          type_string = type(find_spec)
          logger.warning(
              f'Unsupported find specification type: {type_string!s}')

  @classmethod
  def CheckKeyCompatibility(cls, key_path):
    """Checks if a Windows Registry key path is supported by dfWinReg.

    Args:
      key_path (str): path of the Windows Registry key.

    Returns:
      bool: True if key is compatible or False if not.
    """
    key_path_upper = key_path.upper()
    for key_path_prefix in cls._COMPATIBLE_REGISTRY_KEY_PATH_PREFIXES:
      if key_path_upper.startswith(key_path_prefix):
        return True

    logger.warning(f'Key path: "{key_path:s}" is currently not supported')
    return False
