import logging

from regipy.exceptions import NoRegistrySubkeysException, RegistryKeyNotFoundException
from regipy.hive_types import SYSTEM_HIVE_TYPE
from regipy.plugins.plugin import Plugin
from regipy.utils import convert_wintime

logger = logging.getLogger(__name__)

USBSTOR_KEY_PATH = r"Enum\USBSTOR"
DISK_GUID_PATH = r"Device Parameters\Partmgr"
PROPERTIES_NAME_GUID = r"{540b947e-8b40-45bc-a8a2-6a0b894cbda2}"
PROPERTIES_DATES_GUID = r"{83da6326-97a6-4088-9453-a1923f573b29}"
DEVICE_NAME_KEY = "0004"
FIRST_INSTALLED_TIME_KEY = "0065"
LAST_CONNECTED_TIME_KEY = "0066"
LAST_REMOVED_TIME_KEY = "0067"
LAST_INSTALLED_TIME_KEY = "0064"


class USBSTORPlugin(Plugin):
    NAME = "usbstor_plugin"
    DESCRIPTION = "Parse the connected USB devices history"
    COMPATIBLE_HIVE = SYSTEM_HIVE_TYPE

    def run(self):
        try:
            for subkey_path in self.registry_hive.get_control_sets(USBSTOR_KEY_PATH):
                usbstor_key = self.registry_hive.get_key(subkey_path)
                for usbstor_drive in usbstor_key.iter_subkeys():
                    try:
                        disk, manufacturer, title, version = usbstor_drive.name.split(
                            "&"
                        )
                    except ValueError:
                        manufacturer, title, version = None, None, None

                    for serial_subkey in usbstor_drive.iter_subkeys():
                        timestamp = convert_wintime(
                            serial_subkey.header.last_modified, as_json=self.as_json
                        )
                        serial_number = serial_subkey.name
                        key_path = (
                            rf"{subkey_path}\{usbstor_drive.name}\{serial_number}"
                        )

                        try:
                            device_guid_key = self.registry_hive.get_key(
                                rf"{subkey_path}\{usbstor_drive.name}\{serial_number}\{DISK_GUID_PATH}"
                            )
                            disk_guid = device_guid_key.get_value("DiskId")
                        except RegistryKeyNotFoundException:
                            disk_guid = None

                        properties_subkey = serial_subkey.get_subkey("Properties")
                        if not properties_subkey:
                            (
                                device_name,
                                first_installed_time,
                                last_connected_time,
                                last_removed_time,
                                last_installed_time,
                            ) = (None, None, None, None, None)
                        else:
                            try:
                                device_name_guid_key = properties_subkey.get_subkey(
                                    PROPERTIES_NAME_GUID
                                )
                                device_name_key = device_name_guid_key.get_subkey(
                                    DEVICE_NAME_KEY
                                )
                                device_name = device_name_key.get_value()
                            except (
                                RegistryKeyNotFoundException,
                                NoRegistrySubkeysException,
                            ):
                                device_name = None

                            (
                                first_installed_time,
                                last_connected_time,
                                last_removed_time,
                                last_installed_time,
                            ) = (None, None, None, None)

                            dates_subkey = properties_subkey.get_subkey(
                                PROPERTIES_DATES_GUID, raise_on_missing=False
                            )
                            if dates_subkey:
                                first_installed_key = dates_subkey.get_subkey(
                                    FIRST_INSTALLED_TIME_KEY, raise_on_missing=False
                                )
                                if first_installed_key:
                                    first_installed_time = (
                                        first_installed_key.get_value(
                                            as_json=self.as_json
                                        )
                                    )

                                last_connected_key = dates_subkey.get_subkey(
                                    LAST_CONNECTED_TIME_KEY, raise_on_missing=False
                                )
                                if last_connected_key:
                                    last_connected_time = last_connected_key.get_value(
                                        as_json=self.as_json
                                    )

                                last_removed_key = dates_subkey.get_subkey(
                                    LAST_REMOVED_TIME_KEY, raise_on_missing=False
                                )
                                if last_removed_key:
                                    last_removed_time = last_removed_key.get_value(
                                        as_json=self.as_json
                                    )

                                last_installed_key = dates_subkey.get_subkey(
                                    LAST_INSTALLED_TIME_KEY, raise_on_missing=False
                                )
                                if last_installed_key:
                                    last_installed_time = last_installed_key.get_value(
                                        as_json=self.as_json
                                    )

                        self.entries.append(
                            {
                                "last_write": timestamp,
                                "key_path": key_path,
                                "last_connected": last_connected_time,
                                "last_removed": last_removed_time,
                                "first_installed": first_installed_time,
                                "last_installed": last_installed_time,
                                "serial_number": serial_number,
                                "device_name": device_name,
                                "disk_guid": disk_guid,
                                "manufacturer": manufacturer,
                                "version": version,
                                "title": title,
                            }
                        )

        except RegistryKeyNotFoundException as ex:
            logger.error(
                f"Could not find {self.NAME} plugin data at: {USBSTOR_KEY_PATH}: {ex}"
            )
