# Copyright 2021-2022 Nitrokey Developers
#
# Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
# http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
# http://opensource.org/licenses/MIT>, at your option. This file may not be
# copied, modified, or distributed except according to those terms.

import logging
import platform
import re
import sys
from typing import Optional, TypeVar

from nitrokey.trussed import Uuid, Version

from . import FirmwareMetadata, ProgressCallback, TrussedBootloader, Variant
from .lpc55_upload.mboot.interfaces.usb import MbootUSBInterface
from .lpc55_upload.mboot.mcuboot import McuBoot
from .lpc55_upload.mboot.properties import PropertyTag
from .lpc55_upload.sbfile.sb2.images import BootImageV21
from .lpc55_upload.utils.interfaces.device.usb_device import UsbDevice

RKTH = bytes.fromhex("050aad3e77791a81e59c5b2ba5a158937e9460ee325d8ccba09734b8fdebb171")
KEK = bytes([0xAA] * 32)
UUID_LEN = 4
FILENAME_PATTERN = re.compile("(firmware|alpha)-nk3..-lpc55-(?P<version>.*)\\.sb2$")

T = TypeVar("T", bound="TrussedBootloaderLpc55")

logger = logging.getLogger(__name__)


class TrussedBootloaderLpc55(TrussedBootloader):
    """A Nitrokey 3 device running the LPC55 bootloader."""

    def __init__(self, device: UsbDevice):
        self._validate_vid_pid(device.vid, device.pid)
        self._path = device.path
        self.device = McuBoot(MbootUSBInterface(device))

    def __enter__(self: T) -> T:
        self.device.open()
        return self

    @property
    def variant(self) -> Variant:
        return Variant.LPC55

    @property
    def path(self) -> str:
        if isinstance(self._path, bytes):
            return self._path.decode("UTF-8")
        return self._path

    @property
    def status(self) -> str:
        return self.device.status_string

    def close(self) -> None:
        self.device.close()

    def reboot(self) -> bool:
        if not self.device.reset(reopen=False):
            # On Windows, this function returns false even if the reset was successful
            if platform.system() == "Windows":
                logger.warning("Failed to reboot Nitrokey 3 bootloader")
            else:
                raise Exception("Failed to reboot Nitrokey 3 bootloader")
        return True

    def uuid(self) -> Optional[Uuid]:
        uuid = self.device.get_property(PropertyTag.UNIQUE_DEVICE_IDENT)
        if not uuid:
            raise ValueError("Missing response for UUID property query")
        if len(uuid) != UUID_LEN:
            raise ValueError(f"UUID response has invalid length {len(uuid)}")

        # See GetProperties::device_uuid in the lpc55 crate:
        # https://github.com/lpc55/lpc55-host/blob/main/src/bootloader/property.rs#L222
        wrong_endian = (uuid[3] << 96) + (uuid[2] << 64) + (uuid[1] << 32) + uuid[0]
        right_endian = wrong_endian.to_bytes(16, byteorder="little")
        return Uuid(int.from_bytes(right_endian, byteorder="big"))

    def update(
        self,
        image: bytes,
        callback: Optional[ProgressCallback] = None,
        check_errors: bool = False,
    ) -> None:
        success = self.device.receive_sb_file(
            image,
            progress_callback=callback,
            check_errors=check_errors,
        )
        logger.debug(f"Firmware update finished with status {self.status}")
        if success:
            self.reboot()
        else:
            raise Exception(f"Firmware update failed with status {self.status}")

    @classmethod
    def _list_vid_pid(cls: type[T], vid: int, pid: int) -> list[T]:
        devices = []
        for device in UsbDevice.enumerate(vid=vid, pid=pid):
            try:
                devices.append(cls(device))
            except ValueError:
                logger.warn(
                    f"Invalid Nitrokey 3 LPC55 bootloader returned by enumeration: {device}"
                )
        return devices

    @classmethod
    def _open(cls: type[T], path: str) -> Optional[T]:
        devices = UsbDevice.enumerate(path=path)
        if len(devices) == 0:
            logger.warn(f"No HID device at {path}")
            return None
        if len(devices) > 1:
            logger.warn(f"Multiple HID devices at {path}: {devices}")
            return None

        try:
            return cls(devices[0])
        except ValueError:
            logger.warn(
                f"No Nitrokey 3 bootloader at path {path}", exc_info=sys.exc_info()
            )
            return None


def parse_firmware_image(data: bytes) -> FirmwareMetadata:
    image = BootImageV21.parse(data, kek=KEK)
    bcd_version = image.header.product_version
    version = Version(
        major=bcd_version.major, minor=bcd_version.minor, patch=bcd_version.service
    )
    metadata = FirmwareMetadata(version=version)
    if image.cert_block:
        if image.cert_block.rkth == RKTH:
            metadata.signed_by = "Nitrokey"
            metadata.signed_by_nitrokey = True
        else:
            metadata.signed_by = f"unknown issuer (RKTH: {image.cert_block.rkth.hex()})"
    return metadata
