import binascii
import datetime as dt

from typing import List, Optional, Union

import logging

import attr

from construct import (
    Bytes,
    CString,
    EnumIntegerString,
    GreedyRange,
    Int32sl,
    Int32ul,
    Int64ul,
    StreamError,
    ConstError,
)
from io import BytesIO


from regipy.exceptions import (
    NoRegistrySubkeysException,
    RegistryKeyNotFoundException,
    RegistryValueNotFoundException,
    RegipyGeneralException,
    UnidentifiedHiveException,
    RegistryParsingException,
)
from regipy.hive_types import SUPPORTED_HIVE_TYPES
from regipy.security_utils import convert_sid, get_acls
from regipy.structs import (
    REGF_HEADER,
    HBIN_HEADER,
    CM_KEY_NODE,
    LF_LH_SK_ELEMENT,
    VALUE_KEY,
    INDEX_ROOT,
    REGF_HEADER_SIZE,
    INDEX_ROOT_SIGNATURE,
    LEAF_INDEX_SIGNATURE,
    FAST_LEAF_SIGNATURE,
    HASH_LEAF_SIGNATURE,
    BIG_DATA_BLOCK,
    INDEX_LEAF,
    DEFAULT_VALUE,
    SECURITY_KEY_v1_1,
    SECURITY_DESCRIPTOR,
    SID,
    VALUE_TYPE_ENUM,
)
from regipy.utils import (
    boomerang_stream,
    convert_wintime,
    identify_hive_type,
    MAX_LEN,
    trim_registry_data_for_error_msg,
    try_decode_binary,
)

logger = logging.getLogger(__name__)


@attr.s
class Cell:
    """
    Represents a Registry cell header
    """

    offset: int = attr.ib()
    cell_type: str = attr.ib()
    size: int = attr.ib()


@attr.s
class VKRecord:
    """
    The VK Record contains a value
    """

    value_type: EnumIntegerString = attr.ib()
    value_type_str: str = attr.ib()
    value: bytes = attr.ib()
    size: int = attr.ib(default=0)
    is_corrupted: bool = attr.ib(default=False)


@attr.s
class LIRecord:
    data: bytes = attr.ib()


@attr.s
class Value:
    name: str = attr.ib()
    value: Union[str, int, bytes] = attr.ib()
    value_type: str = attr.ib()
    is_corrupted: bool = attr.ib(default=False)


@attr.s
class Subkey:
    subkey_name: str = attr.ib()
    path: str = attr.ib()
    timestamp: dt.datetime = attr.ib()
    values_count: int = attr.ib()
    values: List[Value] = attr.ib(factory=list)

    # This field will be used if a partial hive was given, if not it would be None.
    actual_path: Optional[str] = attr.ib(default=None)


class RIRecord:
    data = None
    header = None

    def __init__(self, stream):
        self.header = INDEX_ROOT.parse_stream(stream)


class RegistryHive:
    CONTROL_SETS = [r"\ControlSet001", r"\ControlSet002"]

    def __init__(self, hive_path, hive_type=None, partial_hive_path=None):
        """
        Represents a registry hive
        :param hive_path: Path to the registry hive
        :param hive_type: The hive type can be specified if this is a partial hive,
                          or for some other reason regipy cannot identify the hive type
        :param partial_hive_path: The path from which the partial hive actually starts, for example:
                                  hive_type=ntuser partial_hive_path="/Software" would mean
                                  this is actually a HKCU hive, starting from HKCU/Software
        """

        self.partial_hive_path = None
        self.hive_type = None

        with open(hive_path, "rb") as f:
            self._stream = BytesIO(f.read())

        with boomerang_stream(self._stream) as s:
            self.header = REGF_HEADER.parse_stream(s)

            # Get the first cell in root HBin, which is the root NKRecord:
            root_hbin = self.get_hbin_at_offset()
            root_hbin_cell = next(root_hbin.iter_cells(s))
            self.root = NKRecord(root_hbin_cell, s)
        self.name = self.header.file_name

        if hive_type:
            if hive_type.lower() in SUPPORTED_HIVE_TYPES:
                self.hive_type = hive_type
            else:
                raise UnidentifiedHiveException(
                    f"{hive_type} is not a supported hive type: "
                    f"only the following are supported: {SUPPORTED_HIVE_TYPES}"
                )
        else:
            try:
                self.hive_type = identify_hive_type(self.name)
            except UnidentifiedHiveException:
                logger.info(
                    f"Hive type for {hive_path} was not identified: {self.name}"
                )

        if partial_hive_path:
            self.partial_hive_path = partial_hive_path

    def recurse_subkeys(
        self,
        nk_record=None,
        path_root=None,
        as_json=False,
        is_init=True,
        fetch_values=True,
    ):
        """
        Recurse over a subkey, and yield all of its subkeys and values
        :param nk_record: an instance of NKRecord from which to start iterating, if None, will start from Root
        :param path_root: If we are iterating an incomplete hive, for example a hive tree starting
                          from ControlSet001 and not SYSTEM, there is no way to know that.
                          This string will be added as prefix to all paths.
        :param as_json: Whether to normalize the data as JSON or not
        :param fetch_values: If False, subkey values will not be returned, but the iteration will be faster
        """
        # If None, will start iterating from Root NK entry
        if not nk_record:
            nk_record = self.root

        # Iterate over subkeys
        if nk_record.header.subkey_count:
            for subkey in nk_record.iter_subkeys():
                if path_root:
                    subkey_path = (
                        r"{}\{}".format(path_root, subkey.name)
                        if path_root
                        else r"\{}".format(subkey.name)
                    )
                else:
                    subkey_path = f"\\{subkey.name}"

                # Leaf Index records do not contain subkeys
                if isinstance(subkey, LIRecord):
                    continue

                if subkey.subkey_count:
                    yield from self.recurse_subkeys(
                        nk_record=subkey,
                        path_root=subkey_path,
                        as_json=as_json,
                        is_init=False,
                        fetch_values=fetch_values,
                    )

                values = []
                if fetch_values:
                    if subkey.values_count:
                        try:
                            if as_json:
                                values = [
                                    attr.asdict(x)
                                    for x in subkey.iter_values(as_json=as_json)
                                ]
                            else:
                                values = list(subkey.iter_values(as_json=as_json))
                        except RegistryParsingException:
                            logger.exception(
                                f"Failed to parse hive value at path: {trim_registry_data_for_error_msg(path_root)}"
                            )

                ts = convert_wintime(subkey.header.last_modified)
                yield Subkey(
                    subkey_name=subkey.name,
                    path=subkey_path,
                    timestamp=ts.isoformat() if as_json else ts,
                    values=values,
                    values_count=subkey.values_count,
                    actual_path=(
                        f"{self.partial_hive_path}{subkey_path}"
                        if self.partial_hive_path
                        else None
                    ),
                )

        if is_init:
            # Get the values of the subkey
            values = []
            if nk_record.values_count:
                try:
                    if as_json:
                        values = [
                            attr.asdict(x)
                            for x in nk_record.iter_values(as_json=as_json)
                        ]
                    else:
                        values = list(nk_record.iter_values(as_json=as_json))
                except RegistryParsingException as ex:
                    logger.exception(
                        f"Failed to parse hive value at path: {trim_registry_data_for_error_msg(path_root)}: {ex}"
                    )
                    values = []

            ts = convert_wintime(nk_record.header.last_modified)
            subkey_path = path_root or "\\"
            yield Subkey(
                subkey_name=nk_record.name,
                path=subkey_path,
                timestamp=ts.isoformat() if as_json else ts,
                values=values,
                values_count=len(values),
                actual_path=(
                    f"{self.partial_hive_path}\\{subkey_path}"
                    if self.partial_hive_path
                    else None
                ),
            )

    def get_hbin_at_offset(self, offset=0):
        """
        This offset is from start of data (meaning that it starts from 4096)
        If not offset is given, will return the root Hbin
        :return:
        """
        self._stream.seek(REGF_HEADER_SIZE + offset)
        return HBin(self._stream)

    def get_key(self, key_path):
        if self.partial_hive_path:
            if key_path.startswith(self.partial_hive_path):
                key_path = key_path.partition(self.partial_hive_path)[-1]
            else:
                raise RegistryKeyNotFoundException(
                    f"Did not find subkey at {key_path}, because this is a partial hive"
                )

        logger.debug("Getting key: {}".format(key_path))

        # If the key path is \ we are just refering to root
        if key_path == "\\":
            return self.root

        # If the path contain slashes, this is a full path. Split it
        if "\\" in key_path:
            key_path_parts = key_path.split("\\")[1:]
        else:
            key_path_parts = [key_path]

        previous_key_name = []

        subkey = self.root.get_subkey(key_path_parts.pop(0), raise_on_missing=False)

        if not subkey:
            raise RegistryKeyNotFoundException(
                "Did not find subkey at {}".format(key_path)
            )

        if not key_path_parts:
            return subkey

        for path_part in key_path_parts:
            new_path = "\\".join(previous_key_name)
            previous_key_name.append(subkey.name)
            subkey = subkey.get_subkey(path_part, raise_on_missing=False)

            if not subkey:
                raise RegistryKeyNotFoundException(
                    "Did not find {} at {}".format(path_part, new_path)
                )
        return subkey

    def get_control_sets(self, registry_path):
        """
        Get the optional control sets for a registry hive
        :param registry_path:
        :return: A list of paths, including the control sets
        """

        found_control_sets = []
        for cs in self.CONTROL_SETS:
            try:
                found_control_sets.append(self.get_key(cs))
            except RegistryKeyNotFoundException:
                continue
        result = [
            r"\{}\{}".format(subkey.name, registry_path)
            for subkey in found_control_sets
        ]
        logger.debug("Found control sets: {}".format(result))
        return result


class HBin:
    def __init__(self, stream):
        """
        :param stream: a stream at the start of the hbin block
        """
        self.header = HBIN_HEADER.parse_stream(stream)
        self.hbin_data_offset = stream.tell()

    def iter_cells(self, stream):
        stream.seek(self.hbin_data_offset)
        offset = stream.tell()
        while offset < self.hbin_data_offset + self.header.size - HBIN_HEADER.sizeof():
            hbin_cell_size = Int32sl.parse_stream(stream)

            # If the cell size is positive, it means it is unallocated. We are not interested in those on a regular run
            if hbin_cell_size >= 0:
                continue

            bytes_to_read = (hbin_cell_size * -1) - 4

            cell_type = Bytes(2).parse_stream(stream)

            # Yield the cell
            yield Cell(
                cell_type=cell_type.decode(), offset=stream.tell(), size=bytes_to_read
            )

            # Go to the next cell
            offset += stream.tell() + bytes_to_read


class NKRecord:
    """
    The NKRecord represents a Name Key entry
    """

    def __init__(self, cell, stream):
        stream.seek(cell.offset)
        self.header = CM_KEY_NODE.parse_stream(stream)
        self._stream = stream

        # Sometimes the key names are ASCII and sometimes UTF-16 little endian
        if self.header.flags.KEY_COMP_NAME:
            # Compressed (ASCII) key name
            self.name = self.header.key_name_string.decode("ascii", errors="replace")
        else:
            # Unicode (UTF-16) key name
            self.name = self.header.key_name_string.decode(
                "utf-16-le", errors="replace"
            )
            logger.debug(f'Unicode key name identified: "{self.name}"')

        self.subkey_count = self.header.subkey_count
        self.values_count = self.header.values_count
        self.volatile_subkeys_count = self.header.volatile_subkey_count

    def get_subkey(self, key_name, raise_on_missing=True):
        if not self.subkey_count and raise_on_missing:
            raise NoRegistrySubkeysException(
                "No subkeys for {}".format(self.header.key_name_string)
            )

        for subkey in self.iter_subkeys():
            # This should not happen
            if not isinstance(subkey, NKRecord):
                raise RegipyGeneralException(f"Unknown record type: {subkey}")

            if subkey.name.upper() == key_name.upper():
                return subkey

        if raise_on_missing:
            raise NoRegistrySubkeysException(
                "No subkey {} for {}".format(key_name, self.header.key_name_string)
            )

    def iter_subkeys(self):

        if not self.header.subkey_count:
            return None

        # Go to the offset where the subkey list starts (+4 is because of the cell header)
        target_offset = REGF_HEADER_SIZE + 4 + self.header.subkeys_list_offset
        self._stream.seek(target_offset)

        # Read the signature
        try:
            signature = Bytes(2).parse_stream(self._stream)
        except StreamError as ex:
            raise RegistryParsingException(
                f"Bad subkey at offset {target_offset}: {ex}"
            )

        # LF,  LH and RI contain subkeys
        if signature in [
            HASH_LEAF_SIGNATURE,
            FAST_LEAF_SIGNATURE,
            LEAF_INDEX_SIGNATURE,
        ]:
            yield from self._parse_subkeys(self._stream, signature=signature)
        # RI contains pointers to arrays of subkeys
        elif signature == INDEX_ROOT_SIGNATURE:
            ri_record = RIRecord(self._stream)
            if ri_record.header.element_count > 0:
                for element in ri_record.header.elements:
                    # We skip 6 because of the signature as well as the cell header
                    element_target_offset = (
                        REGF_HEADER_SIZE + 4 + element.subkey_list_offset
                    )
                    self._stream.seek(element_target_offset)
                    yield from self._parse_subkeys(self._stream)

    @staticmethod
    def _parse_subkeys(stream, signature=None):
        """
        Parse an LI , LF or LH Record
        :param stream: A stream at the header of the LH or LF entry, skipping the signature
        :return:
        """
        if not signature:
            signature = stream.read(2)

        if signature in [HASH_LEAF_SIGNATURE, FAST_LEAF_SIGNATURE]:
            subkeys = LF_LH_SK_ELEMENT.parse_stream(stream)
        elif signature == LEAF_INDEX_SIGNATURE:
            subkeys = INDEX_LEAF.parse_stream(stream)
        else:
            raise RegistryParsingException(
                f"Expected a known signature, got: {signature} at offset {stream.tell()}"
            )

        for subkey in subkeys.elements:
            stream.seek(REGF_HEADER_SIZE + subkey.key_node_offset)

            # This cell should always be allocated, therefor we expect a negative size
            cell_size = Int32sl.parse_stream(stream) * -1

            # We read to this offset and skip 2 bytes, because that is the cell size we just read
            nk_cell = Cell(cell_type="nk", offset=stream.tell() + 2, size=cell_size)
            nk_record = NKRecord(cell=nk_cell, stream=stream)
            yield nk_record

    @staticmethod
    def read_value(vk, stream):
        """
        Read a registry value
        :param vk: A parse VK record
        :param stream: The registry stream
        :return: A VKRecord
        """
        stream.seek(REGF_HEADER_SIZE + 4 + vk.data_offset)
        data_type = vk.data_type
        data = stream.read(vk.data_size)
        return VKRecord(
            value_type=data_type,
            value_type_str=str(data_type),
            value=data,
            size=vk.data_size,
        )

    @staticmethod
    def _parse_indirect_block(stream, value):
        # This is an indirect datablock (Bigger than 16344, therefor we handle it differently)
        # The value inside the vk entry actually contains a pointer to the buffers containing the data
        big_data_block_header = BIG_DATA_BLOCK.parse(value.value)

        # Go to the start of the segment offset list
        stream.seek(REGF_HEADER_SIZE + big_data_block_header.offset_to_list_of_segments)
        buffer = BytesIO()

        # Read them sequentially until we got all the size of the VK
        value_size = value.size
        while value_size > 0:
            data_segment_offset = Int32ul.parse_stream(stream)
            with boomerang_stream(stream) as tmpstream:
                tmpstream.seek(REGF_HEADER_SIZE + 4 + data_segment_offset)
                tmpbuffer = tmpstream.read(min(0x3FD8, value_size))
                value_size -= len(tmpbuffer)
                buffer.write(tmpbuffer)
        buffer.seek(0)
        return buffer.read()

    def iter_values(self, as_json=False, max_len=MAX_LEN, trim_values=True):
        """
        Get the values of a subkey. Will raise if no values exist
        :param as_json: Whether to normalize the data as JSON or not
        :param max_len: Max length of value to return
        :param trim_values: whether to trim values to MAX_LEN
        :return: List of values for the subkey
        """
        if not self.values_count:
            return

        # Get the offset of the values key. We skip 4 because of Cell Header
        target_offset = REGF_HEADER_SIZE + 4 + self.header.values_list_offset
        self._stream.seek(target_offset)

        for _ in range(self.values_count):
            is_corrupted = False
            try:
                vk_offset = Int32ul.parse_stream(self._stream)
            except StreamError:
                logger.info(f"Skipping bad registry VK at {self._stream.tell()}")
                raise RegistryParsingException(
                    f"Bad registry VK at {self._stream.tell()}"
                )

            with boomerang_stream(self._stream) as substream:
                actual_vk_offset = REGF_HEADER_SIZE + 4 + vk_offset
                substream.seek(actual_vk_offset)
                try:
                    vk = VALUE_KEY.parse_stream(substream)
                except (ConstError, StreamError):
                    logger.error(
                        f"Could not parse VK at {substream.tell()}, registry hive is probably corrupted."
                    )
                    return

                value = self.read_value(vk, substream)

                if vk.name_size == 0:
                    value_name = "(default)"
                elif vk.flags.VALUE_COMP_NAME:
                    # Compressed (ASCII) value name
                    value_name = vk.name.decode("ascii", errors="replace")
                else:
                    # Unicode (UTF-16) value name
                    value_name = vk.name.decode("utf-16-le", errors="replace")
                    logger.debug(f'Unicode value name identified: "{value_name}"')

                # If the value is bigger than this value, it means this is a DEVPROP structure
                # https://doxygen.reactos.org/d0/dba/devpropdef_8h_source.html
                # https://sourceforge.net/p/mingw-w64/mingw-w64/ci/668a1d3e85042c409e0c292e621b3dc0aa26177c/tree/
                # mingw-w64-headers/include/devpropdef.h?diff=dd86a3b7594dadeef9d6a37c4b6be3ca42ef7e94
                # We currently do not support these, We are going to make the best effort to dump as string.
                # This int casting will always work because the data_type is construct's EnumIntegerString
                if int(vk.data_type) > 0xFFFF0000:
                    data_type = VALUE_TYPE_ENUM.parse(
                        Int32ul.build(int(vk.data_type) & 0xFFFF)
                    )
                    logger.info(
                        f"Value at {hex(actual_vk_offset)} contains DEVPROP structure of type {data_type}"
                    )

                # Skip this unknown data type, research pending :)
                # TODO: Add actual parsing
                elif int(vk.data_type) == 0x200000:
                    logger.info(
                        f"Skipped unknown data type value at {actual_vk_offset}"
                    )
                    continue
                else:
                    data_type = str(vk.data_type)

                if data_type in ["REG_SZ", "REG_EXPAND", "REG_EXPAND_SZ"]:
                    if vk.data_size >= 0x80000000:
                        # data is contained in the data_offset field
                        value.size -= 0x80000000
                        actual_value = vk.data_offset
                    elif vk.data_size > 0x3FD8 and value.value[:2] == b"db":
                        data = self._parse_indirect_block(substream, value)
                        actual_value = try_decode_binary(
                            data, as_json=as_json, trim_values=trim_values
                        )
                    else:
                        actual_value = try_decode_binary(
                            value.value, as_json=as_json, trim_values=trim_values
                        )
                elif data_type in ["REG_BINARY", "REG_NONE"]:
                    if vk.data_size >= 0x80000000:
                        # data is contained in the data_offset field
                        actual_value = vk.data_offset
                    elif vk.data_size > 0x3FD8 and value.value[:2] == b"db":
                        try:
                            actual_value = self._parse_indirect_block(substream, value)

                            actual_value = (
                                try_decode_binary(
                                    actual_value, as_json=True, trim_values=trim_values
                                )
                                if as_json
                                else actual_value
                            )
                        except ConstError:
                            logger.error(f"Bad value at {actual_vk_offset}")
                            continue
                    else:
                        # Return the actual data
                        actual_value = (
                            binascii.b2a_hex(value.value).decode()[:max_len]
                            if trim_values
                            else value.value
                        )
                elif data_type == "REG_SZ":
                    actual_value = try_decode_binary(
                        value.value, as_json=as_json, trim_values=trim_values
                    )
                elif data_type == "REG_DWORD":
                    # If the data size is bigger than 0x80000000, data is actually stored in the VK data offset.
                    actual_value = (
                        vk.data_offset
                        if vk.data_size >= 0x80000000
                        else Int32ul.parse(value.value)
                    )
                elif data_type == "REG_QWORD":
                    actual_value = (
                        vk.data_offset
                        if vk.data_size >= 0x80000000
                        else Int64ul.parse(value.value)
                    )
                elif data_type == "REG_MULTI_SZ":
                    parsed_value = GreedyRange(CString("utf-16-le")).parse(value.value)
                    # Because the ListContainer object returned by Construct cannot be turned into a list,
                    # we do this trick
                    actual_value = [x for x in parsed_value if x]
                # We currently dumps this as hex string or raw
                # TODO: Add actual parsing
                elif data_type in [
                    "REG_RESOURCE_REQUIREMENTS_LIST",
                    "REG_RESOURCE_LIST",
                ]:
                    actual_value = (
                        binascii.b2a_hex(value.value).decode()[:max_len]
                        if trim_values
                        else value.value
                    )
                elif data_type == "REG_FILETIME":
                    actual_value = convert_wintime(
                        Int64ul.parse(value.value), as_json=as_json
                    )
                else:
                    actual_value = try_decode_binary(
                        value.value, as_json=as_json, trim_values=trim_values
                    )
                yield Value(
                    name=value_name,
                    value_type=data_type,
                    value=actual_value,
                    is_corrupted=is_corrupted,
                )

    def get_value(
        self,
        value_name=DEFAULT_VALUE,
        as_json=False,
        raise_on_missing=False,
        case_sensitive=True,
    ):
        """
        Get a value by name. Will raise if raise_on_missing is set,
        if no value name is given, will return the content of the default value
        :param value_name: The value name to look for
        :param as_json: Whether to normalize the data as JSON or not
        :param raise_on_missing: Will raise exception if value is missing, else will return None
        :return:
        """
        value_name = value_name if case_sensitive else value_name.lower()
        for value in self.iter_values(as_json=as_json, trim_values=False):
            v = value.name if case_sensitive else value.name.lower()
            if v == value_name:
                return value.value

        if raise_on_missing:
            raise RegistryValueNotFoundException(
                "Did not find the value {} on subkey {}".format(value_name, self.name)
            )
        return None

    def get_values(self, as_json=False, trim_values=False):
        return [x for x in self.iter_values(as_json=as_json, trim_values=trim_values)]

    def get_security_key_info(self):
        self._stream.seek(REGF_HEADER_SIZE + self.header.security_key_offset)
        # TODO: If parsing fails, parse with SECURITY_KEY_v1_2
        security_key = SECURITY_KEY_v1_1.parse_stream(self._stream)
        security_descriptor = SECURITY_DESCRIPTOR.parse(
            security_key.security_descriptor
        )

        with boomerang_stream(self._stream) as s:
            security_base_offset = (
                REGF_HEADER_SIZE + self.header.security_key_offset + 24
            )

            s.seek(security_base_offset + security_descriptor.owner)
            owner_sid = convert_sid(SID.parse_stream(s))

            s.seek(security_base_offset + security_descriptor.group)
            group_sid = convert_sid(SID.parse_stream(s))

            sacl_aces = None
            if security_descriptor.offset_sacl > 0:
                s.seek(security_base_offset + security_descriptor.offset_sacl)
                sacl_aces = get_acls(s)

            dacl_aces = None
            if security_descriptor.offset_dacl > 0:
                s.seek(security_base_offset + security_descriptor.offset_dacl)
                dacl_aces = get_acls(s)
            return {
                "owner": owner_sid,
                "group": group_sid,
                "dacl": dacl_aces,
                "sacl": sacl_aces,
            }

    def get_class_name(self) -> str:
        """
        Gets the key class name as would be returned via
        the `lpClass` argument of the `RegQueryInfoKey()` function.
        """

        # Get the offset of the class name string. We skip 4 because of Cell Header
        read_offset = REGF_HEADER_SIZE + 4 + self.header.class_name_offset

        self._stream.seek(read_offset)
        class_name = self._stream.read(self.header.class_name_size)

        return class_name.decode("utf-16-le", errors="replace")

    def __dict__(self):
        return {
            "name": self.name,
            "subkey_count": self.subkey_count,
            "value_count": self.values_count,
            "values": (
                {x["name"]: x["value"] for x in self.iter_values()}
                if self.values_count
                else None
            ),
            "subkeys": (
                {x["name"] for x in self.iter_subkeys()} if self.subkey_count else None
            ),
            "timestamp": convert_wintime(self.header.last_modified, as_json=True),
            "volatile_subkeys": self.volatile_subkeys_count,
        }
