# This file is Copyright 2019 Volatility Foundation and licensed under the Volatility Software License 1.0
# which is available at https://www.volatilityfoundation.org/license/vsl-v1.0
#
import contextlib
import logging
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union

from volatility3.framework import constants, exceptions, interfaces, objects
from volatility3.framework.configuration import requirements
from volatility3.framework.layers import linear
from volatility3.framework.symbols import intermed
from volatility3.framework.symbols.windows import extensions

vollog = logging.getLogger(__name__)


class RegistryException(exceptions.LayerException):
    """Base Registry Exception class for catching Registry layer errors."""


class RegistryFormatException(RegistryException):
    """Thrown when an error occurs with the underlying Registry file format."""


class RegistryInvalidIndex(RegistryException):
    """Thrown when an index that doesn't exist or can't be found occurs."""


class RegistryHive(linear.LinearlyMappedLayer):
    def __init__(
        self,
        context: interfaces.context.ContextInterface,
        config_path: str,
        name: str,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        super().__init__(
            context=context, config_path=config_path, name=name, metadata=metadata
        )

        self._base_layer = self.config["base_layer"]
        self._hive_offset = self.config["hive_offset"]
        self._table_name = self.config["nt_symbols"]
        self._page_size = 1 << 12

        self._reg_table_name = intermed.IntermediateSymbolTable.create(
            context, self._config_path, "windows", "registry"
        )

        cmhive = self.context.object(
            self._table_name + constants.BANG + "_CMHIVE",
            self._base_layer,
            self._hive_offset,
        )
        self._cmhive_name = cmhive.get_name()
        self.hive = cmhive.Hive

        # TODO: Check the checksum
        if self.hive.Signature != 0xBEE0BEE0:
            raise RegistryFormatException(
                self.name,
                f"Registry hive at {self._hive_offset} does not have a valid signature",
            )

        # Win10 17063 introduced the Registry process to map most hives.  Check
        # if it exists and update RegistryHive._base_layer
        try:
            registry_proc = self._find_registry_process()
            if registry_proc:
                self._base_layer = registry_proc.add_process_layer()
        except ValueError:
            vollog.log(
                constants.LOGLEVEL_VVVV,
                "Error walking process list, results may not be valid.",
            )

        self._base_block = self.hive.BaseBlock.dereference()

        self._minaddr = 0
        try:
            self._hive_maxaddr_non_volatile = self.hive.Storage[0].Length
            self._hive_maxaddr_volatile = self.hive.Storage[1].Length
            self._maxaddr = 0x80000000 | self._hive_maxaddr_volatile
            vollog.log(
                constants.LOGLEVEL_VVVV,
                f"Setting hive {self.name} max address to {hex(self._maxaddr)}",
            )
        except exceptions.InvalidAddressException:
            self._hive_maxaddr_non_volatile = 0x7FFFFFFF
            self._hive_maxaddr_volatile = 0x7FFFFFFF
            self._maxaddr = 0x80000000 | self._hive_maxaddr_volatile
            vollog.log(
                constants.LOGLEVEL_VVVV,
                f"Exception when setting hive {self.name} max address, using {hex(self._maxaddr)}",
            )

    def _find_registry_process(self) -> Optional[extensions.EPROCESS]:
        """Walk the active process list and return the Registry process if it exists. Duplicates
        PsList.list_processes() since pulling in the plugin causes problems.

        Returns:
            The Registry EPROCESS object if it exists, or None
        """

        kernel = self.context.modules.get(self.config["kernel_module_name"])

        if not kernel or not kernel.offset:
            raise ValueError(
                "Intel layer does not have an associated kernel virtual offset, failing"
            )

        ps_aph_offset = kernel.get_symbol("PsActiveProcessHead").address
        list_entry = kernel.object(object_type="_LIST_ENTRY", offset=ps_aph_offset)
        reloff = kernel.get_type("_EPROCESS").relative_child_offset(
            "ActiveProcessLinks"
        )
        eproc = kernel.object(
            object_type="_EPROCESS",
            offset=list_entry.vol.offset - reloff,
            absolute=True,
        )

        for proc in eproc.ActiveProcessLinks:
            proc_name = proc.ImageFileName.cast(
                "string", max_length=proc.ImageFileName.vol.count, errors="replace"
            )
            if proc_name == "Registry" and proc.InheritedFromUniqueProcessId == 4:
                return proc

        return None

    def _get_hive_maxaddr(self, volatile):
        return (
            self._hive_maxaddr_volatile if volatile else self._hive_maxaddr_non_volatile
        )

    def get_name(self) -> str:
        return self._cmhive_name or "[NONAME]"

    @property
    def hive_offset(self) -> int:
        return self._hive_offset

    @property
    def address_mask(self) -> int:
        """Return a mask that allows for the volatile bit to be set."""
        return super().address_mask | 0x80000000

    @property
    def root_cell_offset(self) -> int:
        """Returns the offset for the root cell in this hive."""
        with contextlib.suppress(exceptions.InvalidAddressException):
            if (
                self._base_block.Signature.cast(
                    "string", max_length=4, encoding="latin-1"
                )
                == "regf"
            ):
                return self._base_block.RootCell
        return 0x20

    def get_cell(self, cell_offset: int) -> "objects.StructType":
        """Returns the appropriate Cell value for a cell offset."""
        # This would be an _HCELL containing CELL_DATA, but to save time we skip the size of the HCELL
        cell = self._context.object(
            object_type=self._table_name + constants.BANG + "_CELL_DATA",
            offset=cell_offset + 4,
            layer_name=self.name,
        )
        return cell

    def get_node(self, cell_offset: int) -> "objects.StructType":
        """Returns the appropriate Node, interpreted from the Cell based on its
        Signature."""
        cell = self.get_cell(cell_offset)
        try:
            signature = cell.cast("string", max_length=2, encoding="latin-1")
        except (RegistryException, exceptions.InvalidAddressException):
            vollog.debug(
                f"Failed to get cell signature for cell (0x{cell.vol.offset:x})"
            )
            return cell
        if signature == "nk":
            return cell.u.KeyNode
        elif signature == "sk":
            return cell.u.KeySecurity
        elif signature == "vk":
            return cell.u.KeyValue
        elif signature == "db":
            # Big Data
            return cell.u.ValueData
        elif signature == "lf" or signature == "lh" or signature == "ri":
            # Fast Leaf, Hash Leaf, Index Root
            return cell.u.KeyIndex
        else:
            # It doesn't matter that we use KeyNode, we're just after the first two bytes
            vollog.debug(
                f"Unknown Signature {signature} (0x{cell.u.KeyNode.Signature:x}) at offset {cell_offset}"
            )
            return cell

    def get_key(
        self, key: str, return_list: bool = False
    ) -> Union[List[objects.StructType], objects.StructType]:
        """Gets a specific registry key by key path.

        return_list specifies whether the return result will be a single
        node (default) or a list of nodes from root to the current node
        (if return_list is true).

        Raises RegistryFormatException if an invalid structure is encountered
        Raises KeyError if the key is not found
        """
        root_node = self.get_node(self.root_cell_offset)
        if not root_node.vol.type_name.endswith(constants.BANG + "_CM_KEY_NODE"):
            raise RegistryFormatException(
                self.name,
                f"Encountered {root_node.vol.type_name} instead of _CM_KEY_NODE",
            )
        node_key = [root_node]
        if key.endswith("\\"):
            key = key[:-1]
        key_array = key.split("\\")
        found_key: List[str] = []
        while key_array and node_key:
            subkeys = node_key[-1].get_subkeys()
            for subkey in subkeys:
                # registry keys are not case sensitive so compare likewise
                # https://learn.microsoft.com/en-us/windows/win32/sysinfo/structure-of-the-registry
                if subkey.get_name().casefold() == key_array[0].casefold():
                    node_key = node_key + [subkey]
                    found_key, key_array = found_key + [key_array[0]], key_array[1:]
                    break
            else:
                node_key = []
        if not node_key:
            raise KeyError(
                "Key {} not found under {}".format(key_array[0], "\\".join(found_key))
            )
        if return_list:
            return node_key
        return node_key[-1]

    def visit_nodes(
        self,
        visitor: Callable[[objects.StructType], None],
        node: Optional[objects.StructType] = None,
    ) -> None:
        """Applies a callable (visitor) to all nodes within the registry tree
        from a given node."""
        if not node:
            node = self.get_node(self.root_cell_offset)
        visitor(node)
        for node in node.get_subkeys():
            self.visit_nodes(visitor, node)

    @staticmethod
    def _mask(value: int, high_bit: int, low_bit: int) -> int:
        """Returns the bits of a value between highbit and lowbit inclusive."""
        high_mask = (2 ** (high_bit + 1)) - 1
        low_mask = (2**low_bit) - 1
        mask = high_mask ^ low_mask
        # print(high_bit, low_bit, bin(mask), bin(value))
        return value & mask

    @classmethod
    def get_requirements(cls) -> List[interfaces.configuration.RequirementInterface]:
        return [
            requirements.IntRequirement(
                name="hive_offset",
                description="Offset within the base layer at which the hive lives",
                default=0,
                optional=False,
            ),
            requirements.SymbolTableRequirement(
                name="nt_symbols", description="Windows kernel symbols"
            ),
            requirements.TranslationLayerRequirement(
                name="base_layer",
                description="Layer in which the registry hive lives",
                optional=False,
            ),
        ]

    def _translate(self, offset: int) -> int:
        """Translates a single cell index to a cell memory offset and the
        suboffset within it."""

        # Ignore the volatile bit when determining maxaddr validity
        volatile = self._mask(offset, 31, 31) >> 31
        if offset & 0x7FFFFFFF > self._get_hive_maxaddr(volatile):
            vollog.log(
                constants.LOGLEVEL_VVV,
                "Layer {} couldn't translate offset {}, greater than {} in {} store of {}".format(
                    self.name,
                    hex(offset & 0x7FFFFFFF),
                    hex(self._get_hive_maxaddr(volatile)),
                    "volatile" if volatile else "non-volatile",
                    self.get_name(),
                ),
            )
            raise RegistryInvalidIndex(
                self.name, "Mapping request for value greater than maxaddr"
            )

        storage = self.hive.Storage[volatile]
        dir_index = self._mask(offset, 30, 21) >> 21
        table_index = self._mask(offset, 20, 12) >> 12
        suboffset = self._mask(offset, 11, 0) >> 0

        table = storage.Map.Directory[dir_index]
        entry = table.Table[table_index]
        return entry.get_block_offset() + suboffset

    def mapping(
        self, offset: int, length: int, ignore_errors: bool = False
    ) -> Iterable[Tuple[int, int, int, int, str]]:
        if length < 0:
            raise ValueError("Mapping length of RegistryHive must be positive or zero")

        # Return the translated offset without checking bounds within the HBIN.  The check runs into
        # issues when pages are swapped on large HBINs, and did not seem to find any errors on single page
        # HBINs while dramatically slowing performance.
        response = []
        current_offset = offset
        remaining_length = length
        chunk_size = self._page_size - (offset & (self._page_size - 1))
        while remaining_length > 0:
            chunk_size = min(chunk_size, remaining_length, self._page_size)
            try:
                translated_offset = self._translate(current_offset)
                response.append(
                    (
                        current_offset,
                        chunk_size,
                        translated_offset,
                        chunk_size,
                        self._base_layer,
                    )
                )
            except exceptions.LayerException:
                if not ignore_errors:
                    raise
            current_offset += chunk_size
            remaining_length -= chunk_size
            chunk_size = self._page_size
        return response

    @property
    def dependencies(self) -> List[str]:
        """Returns a list of layer names that this layer translates onto."""
        return [self.config["base_layer"]]

    def is_valid(self, offset: int, length: int = 1) -> bool:
        """Returns a boolean based on whether the offset is valid or not."""
        with contextlib.suppress(exceptions.InvalidAddressException):
            # Pass this to the lower layers for now
            return all(
                self.context.layers[layer].is_valid(offset, length)
                for (_, _, offset, length, layer) in self.mapping(offset, length)
            )
        return False

    @property
    def minimum_address(self) -> int:
        return self._minaddr

    @property
    def maximum_address(self) -> int:
        return self._maxaddr
