# -*- coding: utf-8 -*-
"""Classes for Windows Registry access."""

from dfwinreg import definitions
from dfwinreg import key_paths
from dfwinreg import regf
from dfwinreg import virtual


class WinRegistryFileMapping(object):
  """Windows Registry file mapping.

  Attributes:
    key_path_prefix (str): Windows Registry key path prefix.
    unique_key_paths (list[str]): key paths unique to the Windows Registry file.
    windows_path (str): Windows path to the Windows Registry file, such as:
        C:\\Windows\\System32\\config\\SYSTEM
  """

  def __init__(self, key_path_prefix, windows_path, unique_key_paths):
    """Initializes the Windows Registry file mapping.

    Args:
      key_path_prefix (str): Windows Registry key path prefix.
      windows_path (str): Windows path to the Windows Registry file, such as:
          C:\\Windows\\System32\\config\\SYSTEM
      unique_key_paths (list[str]): key paths unique to the Windows Registry
          file.
    """
    super(WinRegistryFileMapping, self).__init__()
    self.key_path_prefix = key_path_prefix
    self.unique_key_paths = unique_key_paths
    self.windows_path = windows_path


class WinRegistry(object):
  """Windows Registry."""

  _REGISTRY_FILE_MAPPINGS = [
      # TODO: Windows 3.1
      # Windows 9x/Me
      WinRegistryFileMapping(
          'HKEY_LOCAL_MACHINE',
          '%SystemRoot%\\SYSTEM.DAT',
          ['\\Config', '\\Enum', '\\Hardware', '\\Network', '\\Software',
           '\\System']),
      WinRegistryFileMapping(
          'HKEY_USERS',
          '%SystemRoot%\\USER.DAT',
          ['\\.DEFAULT\\AppEvents', '\\.DEFAULT\\Control Panel',
           '\\.DEFAULT\\Keyboard Layout', '\\.DEFAULT\\Network',
           '\\.DEFAULT\\Software']),
      # Windows NT
      WinRegistryFileMapping(
          'HKEY_CURRENT_USER',
          '%UserProfile%\\NTUSER.DAT',
          ['\\AppEvents', '\\Console', '\\Control Panel', '\\Environment',
           '\\Keyboard Layout', '\\Software']),
      # Windows Vista and later UsrClass.dat mapping
      WinRegistryFileMapping(
          'HKEY_CURRENT_USER\\Software\\Classes',
          '%UserProfile%\\AppData\\Local\\Microsoft\\Windows\\UsrClass.dat',
          ['\\Local Settings\\Software\\Microsoft\\Windows\\CurrentVersion']),
      # Windows 2000, XP and 2003 UsrClass.dat mapping
      WinRegistryFileMapping(
          'HKEY_CURRENT_USER\\Software\\Classes',
          ('%UserProfile%\\Local Settings\\Application Data\\Microsoft\\'
           'Windows\\UsrClass.dat'),
          ['\\Software\\Microsoft\\Windows\\CurrentVersion']),
      WinRegistryFileMapping(
          'HKEY_LOCAL_MACHINE\\SAM',
          '%SystemRoot%\\System32\\config\\SAM',
          ['\\SAM\\Domains\\Account\\Users']),
      WinRegistryFileMapping(
          'HKEY_LOCAL_MACHINE\\Security',
          '%SystemRoot%\\System32\\config\\SECURITY',
          ['\\Policy\\PolAdtEv']),
      WinRegistryFileMapping(
          'HKEY_LOCAL_MACHINE\\Software',
          '%SystemRoot%\\System32\\config\\SOFTWARE',
          ['\\Microsoft\\Windows\\CurrentVersion\\App Paths']),
      WinRegistryFileMapping(
          'HKEY_LOCAL_MACHINE\\System',
          '%SystemRoot%\\System32\\config\\SYSTEM',
          ['\\MountedDevices', '\\Select', '\\Setup'])]

  _ROOT_KEY_ALIASES = {
      'HKCC': 'HKEY_CURRENT_CONFIG',
      'HKCR': 'HKEY_CLASSES_ROOT',
      'HKCU': 'HKEY_CURRENT_USER',
      'HKLM': 'HKEY_LOCAL_MACHINE',
      'HKU': 'HKEY_USERS'}

  _ROOT_SUB_KEYS = frozenset([
      'HKEY_CLASSES_ROOT',
      'HKEY_CURRENT_CONFIG',
      'HKEY_CURRENT_USER',
      'HKEY_DYN_DATA',
      'HKEY_LOCAL_MACHINE',
      'HKEY_PERFORMANCE_DATA',
      'HKEY_USERS'])

  _LOCAL_MACHINE_SUB_KEYS = frozenset([
      'SAM',
      'Security',
      'Software',
      'System'])

  _SELECT_KEY_PATH = 'HKEY_LOCAL_MACHINE\\System\\Select'

  _USER_PROFILE_LIST_KEY_PATH = (
      'HKEY_LOCAL_MACHINE\\Software\\Microsoft\\Windows NT\\CurrentVersion\\'
      'ProfileList')

  _USER_SOFTWARE_CLASSES_KEY_PATH = 'HKEY_CURRENT_USER\\SOFTWARE\\CLASSES'

  # TODO: add support for HKEY_CLASSES_ROOT
  # TODO: add support for HKEY_CURRENT_CONFIG
  # TODO: add support for HKEY_CURRENT_USER
  # TODO: add support for HKEY_DYN_DATA

  _VIRTUAL_KEYS = [
      ('HKEY_USERS', '_GetUsersVirtualKey')]

  def __init__(self, ascii_codepage='cp1252', registry_file_reader=None):
    """Initializes the Windows Registry.

    Args:
      ascii_codepage (Optional[str]): ASCII string codepage.
      registry_file_reader (Optional[WinRegistryFileReader]): Windows Registry
          file reader.
    """
    super(WinRegistry, self).__init__()
    self._ascii_codepage = ascii_codepage
    self._registry_file_reader = registry_file_reader
    self._registry_files = {}
    self._root_key = None
    self._user_registry_files = {}

  def __del__(self):
    """Cleans up the Windows Registry object."""
    for key_path_prefix_upper, registry_file in self._registry_files.items():
      self._registry_files[key_path_prefix_upper] = None
      if registry_file:
        registry_file.Close()

    for profile_path_upper, registry_file in self._user_registry_files.items():
      self._user_registry_files[profile_path_upper] = None
      if registry_file:
        registry_file.Close()

  def _GetCachedFileByPath(self, key_path_upper):
    """Retrieves a cached Windows Registry file for a key path.

    Args:
      key_path_upper (str): Windows Registry key path, in upper case with
          a resolved root key alias.

    Returns:
      tuple: consist:

        str: key path prefix
        WinRegistryFile: corresponding Windows Registry file or None if not
            available.
    """
    longest_key_path_prefix_upper = ''
    longest_key_path_prefix_length = len(longest_key_path_prefix_upper)
    for key_path_prefix_upper in self._registry_files:
      if key_path_upper.startswith(key_path_prefix_upper):
        key_path_prefix_length = len(key_path_prefix_upper)
        if key_path_prefix_length > longest_key_path_prefix_length:
          longest_key_path_prefix_upper = key_path_prefix_upper
          longest_key_path_prefix_length = key_path_prefix_length

    if not longest_key_path_prefix_upper:
      return None, None

    registry_file = self._registry_files.get(
        longest_key_path_prefix_upper, None)
    return longest_key_path_prefix_upper, registry_file

  def _GetCachedUserFileByPath(self, profile_path_upper):
    """Retrieves a cached user Windows Registry file for a profile path.

    Args:
      profile_path_upper (str): user profile path, in upper case.

    Returns:
      WinRegistryFile: corresponding Windows Registry file or None if not
          available.
    """
    return self._user_registry_files.get(profile_path_upper, None)

  def _GetCandidateFileMappingsByPath(self, key_path_upper):
    """Retrieves candidate Windows Registry file mappings for a specific path.

    Args:
      key_path_upper (str): Windows Registry key path, in upper case with
          a resolved root key alias.

    Returns:
      list[WinRegistryFileMapping]: candidate Windows Registry file mappings.
    """
    if key_path_upper[-1] == '\\':
      key_path_upper = key_path_upper[:-1]

    candidate_mappings = []
    for mapping in self._REGISTRY_FILE_MAPPINGS:
      if key_path_upper.startswith(mapping.key_path_prefix.upper()):
        candidate_mappings.append(mapping)

    return candidate_mappings

  def _GetFileByPath(self, key_path_upper):
    """Retrieves a Windows Registry file for a specific path.

    Args:
      key_path_upper (str): Windows Registry key path, in upper case with
          a resolved root key alias.

    Returns:
      tuple[st, WinRegistryFile]: upper case key path prefix and corresponding
          Windows Registry file or None if not available.
    """
    key_path_prefix, cached_registry_file = self._GetCachedFileByPath(
        key_path_upper)

    # NTUSER.DAT and UsrClass.dat have overlapping keys hence for key paths that
    # start "HKEY_CURRENT_USER\\Software\\Classes" do not assume the Windows
    # Registry file for "HKEY_CURRENT_USER" contains the classes key as well.

    if cached_registry_file and (
        key_path_prefix != 'HKEY_CURRENT_USER' or
        not key_path_upper.startswith('HKEY_CURRENT_USER\\SOFTWARE\\CLASSES')):
      return key_path_prefix, cached_registry_file

    key_path_prefix = ''
    candidate_mappings = self._GetCandidateFileMappingsByPath(key_path_upper)
    matching_mappings = []

    for mapping in candidate_mappings:
      mapping_key_path_prefix = mapping.key_path_prefix.upper()
      if mapping_key_path_prefix in self._registry_files:
        # Skip the Windows Registry file if it is already cached.
        continue

      try:
        registry_file = self._OpenFile(mapping.windows_path)
      except IOError:
        registry_file = None

      if len(key_path_prefix) < len(mapping_key_path_prefix):
        key_path_prefix = mapping_key_path_prefix

      if not registry_file:
        continue

      match = True
      if mapping.unique_key_paths:
        # If all unique key paths are found consider the file to match.
        for unique_key_path in mapping.unique_key_paths:
          registry_key = registry_file.GetKeyByPath(unique_key_path)
          if not registry_key:
            match = False

      if match:
        matching_mappings.append((mapping, registry_file))

    if not matching_mappings:
      return key_path_prefix, cached_registry_file

    key_path_prefix = None
    registry_file = None

    while matching_mappings:
      mapping, matching_registry_file = matching_mappings.pop(0)

      if (not key_path_prefix or
          len(mapping.key_path_prefix) > len(key_path_prefix)):
        if registry_file:
          registry_file.Close()

        key_path_prefix = mapping.key_path_prefix
        registry_file = matching_registry_file

    key_path_prefix_upper = key_path_prefix.upper()
    self.MapFile(key_path_prefix_upper, registry_file)

    return key_path_prefix_upper, registry_file

  def _GetKeyByPathFromFile(self, key_path):
    """Retrieves the key for a specific path from file.

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

    Returns:
      WinRegistryKey: Windows Registry key or None if not available.

    Raises:
      RuntimeError: if the key path prefix does not match the key path.
    """
    key_path_upper = key_path.upper()

    key_path_prefix_upper, registry_file = self._GetFileByPath(key_path_upper)
    if not registry_file:
      return None

    if not key_path_upper.startswith(key_path_prefix_upper):
      raise RuntimeError('Key path prefix mismatch.')

    key_path_suffix = key_path[len(key_path_prefix_upper):]
    relative_key_path = key_path_suffix or definitions.KEY_PATH_SEPARATOR

    return registry_file.GetKeyByPath(relative_key_path)

  def _GetRootVirtualKey(self):
    """Retrieves the root key.

    Returns:
      VirtualWinRegistryKey: Windows Registry root key.

    Raises:
      RuntimeError: if there are multiple matching mappings and
          the correct mapping cannot be resolved.
    """
    if not self._root_key:
      self._root_key = virtual.VirtualWinRegistryKey('')

      local_machine_key = None
      for sub_key_name in self._ROOT_SUB_KEYS:
        sub_registry_key = self.GetKeyByPath(sub_key_name)
        if not sub_registry_key:
          sub_registry_key = virtual.VirtualWinRegistryKey(
              sub_key_name, registry=self, relative_key_path=sub_key_name)

          if sub_key_name == 'HKEY_LOCAL_MACHINE':
            local_machine_key = sub_registry_key

        self._root_key.AddSubkey(sub_key_name, sub_registry_key)

      if local_machine_key:
        for sub_key_name in self._LOCAL_MACHINE_SUB_KEYS:
          sub_key_path = definitions.KEY_PATH_SEPARATOR.join([
              'HKEY_LOCAL_MACHINE', sub_key_name])
          sub_registry_key = self.GetKeyByPath(sub_key_path)
          if not sub_registry_key:
            sub_registry_key = virtual.VirtualWinRegistryKey(
                sub_key_name, registry=self, relative_key_path=sub_key_name)

          local_machine_key.AddSubkey(sub_key_name, sub_registry_key)

    return self._root_key

  def _GetUserFileByPath(self, path):
    """Retrieves an user Windows Registry file for a specific path.

    Args:
      path (str): path of the user Windows Registry file.

    Returns:
      WinRegistryFile: corresponding Windows Registry file or None if not
          available.
    """
    path_upper = path.upper()
    registry_file = self._GetCachedUserFileByPath(path_upper)
    if not registry_file:
      try:
        registry_file = self._OpenFile(path)
      except IOError:
        registry_file = None

    return registry_file

  def _GetUsersVirtualKey(self, key_path_suffix):
    """Virtual key callback to determine the users sub keys.

    Args:
      key_path_suffix (str): users Windows Registry key path suffix with
          leading path separator.

    Returns:
      WinRegistryKey: the users Windows Registry key or None if not available.

    Raises:
      RuntimeError: if the key path suffix is not supported.
    """
    if key_path_suffix[0] != definitions.KEY_PATH_SEPARATOR:
      raise RuntimeError(f'Unsupported key path suffix: {key_path_suffix:s}')

    user_key_name, _, key_path_suffix = key_path_suffix[1:].partition(
        definitions.KEY_PATH_SEPARATOR)
    user_key_name = user_key_name.upper()

    # HKEY_USERS\.DEFAULT is an alias for HKEY_USERS\S-1-5-18 which is
    # the Local System account.
    if user_key_name == '.DEFAULT':
      search_key_name = 'S-1-5-18'
    else:
      search_key_name = user_key_name

    if search_key_name.endswith('_CLASSES'):
      search_key_name = search_key_name[:-8]

    user_profile_list_key = self.GetKeyByPath(self._USER_PROFILE_LIST_KEY_PATH)
    if not user_profile_list_key:
      return None

    for user_profile_key in user_profile_list_key.GetSubkeys():
      if search_key_name == user_profile_key.name:
        profile_path_value = user_profile_key.GetValueByName('ProfileImagePath')
        if not profile_path_value:
          break

        profile_path = profile_path_value.GetDataAsObject()
        if not profile_path:
          break

        # HKEY_USERS\%SID%_CLASSES maps to UsrClass.dat
        if user_key_name.endswith('_CLASSES'):
          profile_path = definitions.KEY_PATH_SEPARATOR.join([
              profile_path, 'AppData', 'Local', 'Microsoft', 'Windows',
              'UsrClass.dat'])
        else:
          profile_path = definitions.KEY_PATH_SEPARATOR.join([
              profile_path, 'NTUSER.DAT'])

        registry_file = self._GetUserFileByPath(profile_path)
        if not registry_file:
          break

        key_path_prefix = definitions.KEY_PATH_SEPARATOR.join([
            'HKEY_USERS', user_key_name])
        key_path = definitions.KEY_PATH_SEPARATOR.join([
            key_path_prefix, key_path_suffix])

        registry_file.SetKeyPathPrefix(key_path_prefix)
        return registry_file.GetKeyByPath(key_path)

    return None

  def _GetVirtualKeyByPath(self, key_path):
    """Retrieves the virtual key for a specific path.

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

    Returns:
      VirtualWinRegistryKey: virtual Windows Registry key or None if not
          available.
    """
    key_path_upper = key_path.upper()

    for virtual_key_path, virtual_key_callback in self._VIRTUAL_KEYS:
      virtual_key_path_upper = virtual_key_path.upper()
      if key_path_upper.startswith(virtual_key_path_upper):
        key_path_suffix = key_path[len(virtual_key_path):]

        callback_function = getattr(self, virtual_key_callback)
        virtual_key = callback_function(key_path_suffix)
        if virtual_key:
          return virtual_key

    return None

  def _OpenFile(self, path):
    """Opens a Windows Registry file.

    Args:
      path (str): path of the Windows Registry file.

    Returns:
      WinRegistryFile: Windows Registry file or None if not available.
    """
    if not self._registry_file_reader:
      return None

    return self._registry_file_reader.Open(
        path, ascii_codepage=self._ascii_codepage)

  def GetKeyByPath(self, key_path):
    """Retrieves the key for a specific path.

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

    Returns:
      WinRegistryKey: Windows Registry key or None if not available.

    Raises:
      RuntimeError: if the root key is not supported or the key path prefix
          does not match the key path.
    """
    root_key_path, _, key_path = key_path.partition(
        definitions.KEY_PATH_SEPARATOR)

    # Resolve a root key alias.
    root_key_path = root_key_path.upper()
    root_key_path = self._ROOT_KEY_ALIASES.get(root_key_path, root_key_path)

    if root_key_path not in self._ROOT_SUB_KEYS:
      raise RuntimeError(f'Unsupported root subkey: {root_key_path:s}')

    key_path = definitions.KEY_PATH_SEPARATOR.join([root_key_path, key_path])

    registry_key = self._GetKeyByPathFromFile(key_path)
    if not registry_key:
      registry_key = self._GetVirtualKeyByPath(key_path)

    return registry_key

  def GetRegistryFileMapping(self, registry_file):
    """Determines the Registry file mapping based on the content of the file.

    Args:
      registry_file (WinRegistyFile): Windows Registry file.

    Returns:
      str: key path prefix or an empty string.

    Raises:
      RuntimeError: if there are multiple matching mappings and
          the correct mapping cannot be resolved.
    """
    if not registry_file:
      return ''

    candidate_mappings = []
    for mapping in self._REGISTRY_FILE_MAPPINGS:
      if not mapping.unique_key_paths:
        continue

      # If all unique key paths are found consider the file to match.
      match = True
      for key_path in mapping.unique_key_paths:
        registry_key = registry_file.GetKeyByPath(key_path)
        if not registry_key:
          match = False

      if match:
        candidate_mappings.append(mapping)

    if not candidate_mappings:
      return ''

    if len(candidate_mappings) == 1:
      return candidate_mappings[0].key_path_prefix

    key_path_prefixes = frozenset([
        mapping.key_path_prefix for mapping in candidate_mappings])

    expected_key_path_prefixes = frozenset([
        'HKEY_CURRENT_USER',
        'HKEY_CURRENT_USER\\Software\\Classes'])

    if key_path_prefixes == expected_key_path_prefixes:
      return 'HKEY_CURRENT_USER'

    raise RuntimeError('Unable to resolve Windows Registry file mapping.')

  def GetRootKey(self):
    """Retrieves the Windows Registry root key.

    Returns:
      VirtualWinRegistryKey: Windows Registry root key.
    """
    return self._GetRootVirtualKey()

  def MapFile(self, key_path_prefix, registry_file):
    """Maps the Windows Registry file to a specific key path prefix.

    Args:
      key_path_prefix (str): key path prefix.
      registry_file (WinRegistryFile): Windows Registry file.
    """
    key_path_prefix_upper = key_path_prefix.upper()

    self._registry_files[key_path_prefix_upper] = registry_file
    registry_file.SetKeyPathPrefix(key_path_prefix)

    # If HKEY_CURRENT_USER set HKEY_CURRENT_USER\\Software\\Classes as a virtual
    # key in the file.
    if key_path_prefix_upper == 'HKEY_CURRENT_USER' and isinstance(
        registry_file, regf.REGFWinRegistryFile):
      key_path_prefix_upper, usrclass_registry_file = self._GetFileByPath(
          self._USER_SOFTWARE_CLASSES_KEY_PATH)
      if (key_path_prefix_upper == self._USER_SOFTWARE_CLASSES_KEY_PATH and
          usrclass_registry_file is not None):
        registry_key = usrclass_registry_file.GetKeyByPath(
            '\\Software\\Classes')
        if registry_key and isinstance(registry_key, regf.REGFWinRegistryKey):
          # pylint: disable=protected-access
          pyregf_key = registry_key._pyregf_key
          registry_file.AddVirtualKey('\\Software\\Classes', pyregf_key)

    # If HKEY_LOCAL_MACHINE\\System set HKEY_LOCAL_MACHINE\\System\\
    # CurrentControlSet as a virtual key in the file.
    elif key_path_prefix_upper == 'HKEY_LOCAL_MACHINE\\SYSTEM' and isinstance(
        registry_file, regf.REGFWinRegistryFile):
      registry_file.AddCurrentControlSetKey()

  def MapUserFile(self, profile_path, registry_file):
    """Maps the user Windows Registry file to a specific profile path.

    Args:
      profile_path (str): profile path.
      registry_file (WinRegistryFile): user Windows Registry file.
    """
    self._user_registry_files[profile_path.upper()] = registry_file

  def OpenAndMapFile(self, path):
    """Opens Windows Registry file and maps it to its key path prefix.

    Args:
      path (str): path of the Windows Registry file.
    """
    registry_file = self._OpenFile(path)
    key_path_prefix = self.GetRegistryFileMapping(registry_file)
    self.MapFile(key_path_prefix, registry_file)

  def SplitKeyPath(self, key_path):
    """Splits the key path into path segments.

    Args:
      key_path (str): key path.

    Returns:
      list[str]: key path segments without the root path segment, which is an
          empty string.
    """
    return key_paths.SplitKeyPath(key_path)
