#!/usr/bin/env python
# -*- coding: UTF-8 -*-
#
# Copyright 2019-2022 NXP
#
# SPDX-License-Identifier: BSD-3-Clause

"""The module provides support for On-The-Fly encoding for RTxxx devices."""

from struct import pack
from typing import List, Optional, Union

from Crypto.Cipher import AES

from spsdk import SPSDKError

from ..misc import align_block
from . import Counter, crypto_backend


class KeyBlob:
    """OTFAD KeyBlob: The class specifies AES key and counter initial value for specified address range.

    typedef struct KeyBlob
    {
        unsigned char key[kAesKeySizeBytes]; // 16 bytes, 128-bits, KEY[A15...A00]
        unsigned char ctr[kCtrSizeBytes];    // 8 bytes, 64-bits, CTR[C7...C0]
        unsigned int srtaddr;                // region start, STRADDR[31 - 10]
        unsigned int endaddr;                // region end, ENDADDR[31 - 10]; lowest three bits are used as flags
        // end of 32-byte area covered by CRC
        unsigned int zero_fill;      // zeros
        unsigned int key_blob_crc32; // crc32 over 1st 32-bytes
        // end of 40 byte (5*64-bit) key blob data
        unsigned char expanded_wrap_data[8]; // 8 bytes, used for wrap expanded data
        // end of 48 byte (6*64-bit) wrap data
        unsigned char unused_filler[16]; // unused fill to 64 bytes
    } keyblob_t;
    """

    _START_ADDR_MASK = 0x400 - 1
    # Region addresses are modulo 1024
    # The address ends with RO, ADE, VLD bits. From this perspective, only
    # bits [9:3] must be set to 1. The rest is configurable.
    _END_ADDR_MASK = 0x3F8

    # Key flags mask: RO, ADE, VLD
    _KEY_FLAG_MASK = 0x07
    # This field signals that the entire set of context registers (CTXn_KEY[0-3], CTXn_CTR[0-1],
    # CTXn_RGD_W[0-1] are read-only and cannot be modified. This field is sticky and remains
    # asserted until the next system reset. SR[RRAM] provides another level of register access
    # control and is independent of the RO indicator.
    KEY_FLAG_READ_ONLY = 0x4
    # AES Decryption Enable: For accesses hitting in a valid context, this bit indicates if the fetched data is to be
    # decrypted or simply bypassed.
    KEY_FLAG_ADE = 0x2
    # Valid: This field signals if the context is valid or not.
    KEY_FLAG_VLD = 0x1

    # key length in bytes
    KEY_SIZE = 16
    # counter length in bytes
    CTR_SIZE = 8
    # len of counter init value for export
    _EXPORT_CTR_IV_SIZE = 8
    # this constant seems to be fixed for SB2.1
    _EXPORT_NBLOCKS_5 = 5
    # binary export size
    _EXPORT_KEY_BLOB_SIZE = 64
    # QSPI image alignment length, 512 is supposed to be the safe alignment level for any QSPI device
    # this means that all QSPI images generated by this tool will be sizes of multiple 512
    _IMAGE_ALIGNMENT = 512

    def __init__(
        self,
        start_addr: int,
        end_addr: int,
        key: Optional[bytes] = None,
        counter_iv: Optional[bytes] = None,
        key_flags: int = KEY_FLAG_VLD | KEY_FLAG_ADE,
        # for testing
        zero_fill: Optional[bytes] = None,
        crc: Optional[bytes] = None,
    ):
        """Constructor.

        :param start_addr: start address of the region
        :param end_addr: end address of the region
        :param key_flags: see KEY_FLAG_xxx constants; default flags: RO = 0, ADE = 1, VLD = 1
        :param key: optional AES key; None to use random value
        :param counter_iv: optional counter init value for AES; None to use random value
        :param zero_fill: optional value for zero_fill (for testing only); None to use random value (recommended)
        :param crc: optional value for unused CRC fill (for testing only); None to use random value (recommended)
        :raises SPSDKError: Start or end address are not aligned
        :raises SPSDKError: When there is invalid key
        :raises SPSDKError: When there is invalid start/end address
        :raises SPSDKError: When key_flags exceeds mask
        """
        if key is None:
            key = crypto_backend().random_bytes(self.KEY_SIZE)
        if counter_iv is None:
            counter_iv = crypto_backend().random_bytes(self.CTR_SIZE)
        if (len(key) != self.KEY_SIZE) and (len(counter_iv) != self.CTR_SIZE):
            raise SPSDKError("Invalid key")
        if start_addr < 0 or start_addr >= end_addr or end_addr > 0xFFFFFFFF:
            raise SPSDKError("Invalid start/end address")
        if key_flags & ~self._KEY_FLAG_MASK != 0:
            raise SPSDKError(f"key_flags exceeds mask {hex(self._KEY_FLAG_MASK)}")
        if (start_addr & self._START_ADDR_MASK) != 0:
            raise SPSDKError(
                f"Start address must be aligned to {hex(self._START_ADDR_MASK + 1)} boundary"
            )
        if (end_addr & self._END_ADDR_MASK) != self._END_ADDR_MASK:
            raise SPSDKError(f"End address must be aligned to {hex(self._END_ADDR_MASK)} boundary")
        self.key = key
        self.ctr_init_vector = counter_iv
        self.start_addr = start_addr
        self.end_addr = end_addr
        self.key_flags = key_flags
        self.zero_fill = zero_fill
        self.crc_fill = crc

    def info(self) -> str:
        """Text info about the instance."""
        msg = ""
        msg += f"Key:        {self.key.hex()}\n"
        msg += f"Counter IV: {self.ctr_init_vector.hex()}\n"
        msg += f"Start Addr: {hex(self.start_addr)}\n"
        msg += f"End Addr:   {hex(self.end_addr)}\n"
        return msg

    def plain_data(self) -> bytes:
        """Plain data for selected key range.

        :return: key blob exported into binary form (serialization)
        :raises SPSDKError: Invalid value of zero fill parameter
        :raises SPSDKError: Invalid value crc
        :raises SPSDKError: Invalid length binary data
        """
        result = bytes()
        result += self.key
        result += self.ctr_init_vector
        result += pack("<I", self.start_addr)
        end_addr_with_flags = (self.end_addr & ~self._KEY_FLAG_MASK) | self.key_flags
        result += pack("<I", end_addr_with_flags)
        # zero fill
        if self.zero_fill:
            if len(self.zero_fill) != 4:
                raise SPSDKError("Invalid value")
            result += self.zero_fill
        else:
            result += crypto_backend().random_bytes(4)
        # CRC is not used, use random value
        if self.crc_fill:
            if len(self.crc_fill) != 4:
                raise SPSDKError("Invalid value crc")
            result += self.crc_fill
        else:
            result += crypto_backend().random_bytes(4)
        result += bytes([0] * 8)  # expanded_wrap_data
        result += bytes([0] * 16)  # unused filler
        if len(result) != 64:
            raise SPSDKError("Invalid length binary data")
        return result

    # pylint: disable=invalid-name
    def export(self, kek: Union[bytes, str], iv: bytes = bytes([0xA6] * 8)) -> bytes:
        """Creates key wrap for the key blob.

        :param kek: key to encode; 16 bytes long
        :param iv: counter initialization vector; 8 bytes; optional, OTFAD uses empty init value
        :return: Serialized key blob
        :raises SPSDKError: If any parameter is not valid
        :raises SPSDKError: If length of kek is not valid
        :raises SPSDKError: If length of data is not valid
        """
        if isinstance(kek, str):
            kek = bytes.fromhex(kek)
        if len(kek) != 16:
            raise SPSDKError("Invalid length of kek")
        if len(iv) != self._EXPORT_CTR_IV_SIZE:
            raise SPSDKError("Invalid length of initialization vector")
        n = self._EXPORT_NBLOCKS_5
        plaintext = self.plain_data()  # input data to be encrypted
        if len(plaintext) < n * 8:
            raise SPSDKError("Invalid length of data to be encrypted")

        # step 1: initialize the byte - sized data variables
        # set a = iv
        # for i = 1 to n
        # r[i] = plain_data[i]

        a = iv  # 64-bit integrity check register
        r = bytearray(8) + bytearray(plaintext[0 : 8 * n])  # 8-bit array of 64-bit registers

        # step 2: calculate intermediate values
        # for j = 0 to 5
        #   for i = 1 to n
        #       encr = AES(K, A | R[i])
        #       a = MSB(64, encr) ^ (n*j)+i
        #       r[i] = LSB(64, B)

        for j in range(0, 6):
            for i in range(1, n + 1):
                in_data = a + r[8 * i : 8 * i + 8]  # 128-bit temporary plaintext input vector
                aes = AES.new(kek, AES.MODE_ECB)
                encr = aes.encrypt(in_data)
                xor = encr[7] ^ ((n * j) + i)
                a = bytes(encr[:7]) + bytes([xor])
                r[8 * i : 8 * i + 8] = encr[8:]

        # step 3: output the results
        # set result[0] = A
        # for i = 1 to n
        #   result[i] = r[i]
        result = a + r[8 : 8 * n + 8]

        return align_block(
            result, self._EXPORT_KEY_BLOB_SIZE, padding=0
        )  # align to 64 bytes (0 padding)

    def _get_ctr_nonce(self) -> bytes:
        """Get the counter initial value for image encryption.

        :return: counter bytes
        :raises SPSDKError: If length of counter is not valid
        """
        #  CTRn_x[127-0] = {CTR_W0_x[C0...C3],    // 32 bits of pre-programmed CTR
        #  CTR_W1_x[C4...C7],                     // another 32 bits of CTR
        #  CTR_W0_x[C0...C3] ^ CTR_W1_x[C4...C7], // exclusive-OR of CTR values
        #  systemAddress[31-4], 0000b             // 0-modulo-16 system address */

        if len(self.ctr_init_vector) != 8:
            raise SPSDKError("Invalid length of counter init")

        result = bytearray(16)
        result[:4] = self.ctr_init_vector[:4]
        result[4:8] = self.ctr_init_vector[4:]
        for i in range(0, 4):
            result[8 + i] = self.ctr_init_vector[0 + i] ^ self.ctr_init_vector[4 + i]

        # result[15:12] = start_addr as a counter; nonce has these bytes zero and value passes as counter init value

        return bytes(result)

    def contains_addr(self, addr: int) -> bool:
        """Whether key blob contains specified address.

        :param addr: to be tested
        :return: True if yes, False otherwise
        """
        return self.start_addr <= addr <= self.end_addr

    def matches_range(self, image_start: int, image_end: int) -> bool:
        """Whether key blob matches address range of the image to be encrypted.

        :param image_start: start address of the image
        :param image_end: last address of the image
        :return: True if yes, False otherwise
        """
        return self.contains_addr(image_start) and self.contains_addr(image_end)

    def encrypt_image(self, base_address: int, data: bytes, byte_swap: bool) -> bytes:
        """Encrypt specified data.

        :param base_address: of the data in target memory; must be >= self.start_addr
        :param data: to be encrypted (e.g. plain image); base_address + len(data) must be <= self.end_addr
        :param byte_swap: this probably depends on the flash device, how bytes are organized there
                True should be used for FLASH on EVK RT6xx; False for FLASH on EVK RT5xx
        :return: encrypted data
        :raises SPSDKError: If start_addr or end_addr does not match with base_address (+ data length)
        :raises SPSDKError: If start address is not valid
        """
        if base_address % 16 != 0:
            raise SPSDKError("Invalid start address")  # Start address has to be 16 byte aligned
        data = align_block(data, self._IMAGE_ALIGNMENT)  # align data length
        data_len = len(data)

        # check start and end addresses
        if not self.matches_range(base_address, base_address + data_len - 1):
            raise SPSDKError(
                f"Image address range is not within key blob: {hex(self.start_addr)}-{hex(self.end_addr)}"
            )

        result = bytes()
        counter = Counter(
            self._get_ctr_nonce(), ctr_value=base_address, ctr_byteorder_encoding="big"
        )

        for index in range(0, data_len, 16):
            # prepare data in byte order
            if byte_swap:
                # swap 8 bytes + swap 8 bytes
                data_2_encr = (
                    data[-data_len + index + 7 : -data_len + index - 1 : -1]
                    + data[-data_len + index + 15 : -data_len + index + 7 : -1]
                )
            else:
                data_2_encr = data[index : index + 16]
            # encrypt
            encr_data = crypto_backend().aes_ctr_encrypt(self.key, data_2_encr, counter.value)
            # fix byte order in result
            if byte_swap:
                result += encr_data[-9:-17:-1] + encr_data[-1:-9:-1]  # swap 8 bytes + swap 8 bytes
            else:
                result += encr_data
            # update counter for encryption
            counter.increment(16)

        if len(result) != data_len:
            raise SPSDKError("Invalid length of encrypted data")
        return bytes(result)


class Otfad:
    """OTFAD: On-the-Fly AES Decryption Module."""

    def __init__(self) -> None:
        """Constructor."""
        self._key_blobs: List[KeyBlob] = []

    def __getitem__(self, index: int) -> KeyBlob:
        return self._key_blobs[index]

    def add_key_blob(self, key_blob: KeyBlob) -> None:
        """Add key for specified address range.

        :param key_blob: to be added
        """
        self._key_blobs.append(key_blob)

    def encrypt_image(self, image: bytes, base_addr: int, byte_swap: bool) -> bytes:
        """Encrypt image.

        :param image: plain image to be encrypted
        :param base_addr: where the image will be located in target processor
        :param byte_swap: this probably depends on the flash device, how bytes are organized there
                True should be used for FLASH on EVK RT6xx; False for FLASH on EVK RT5xx
        :return: encrypted image
        :raises SPSDKError: If address range does not match to any key blob
        """
        image_end = base_addr + len(image) - 1
        for key_blob in self._key_blobs:
            if key_blob.matches_range(base_addr, image_end):
                return key_blob.encrypt_image(base_addr, image, byte_swap)

        raise SPSDKError("The image address range does not match to key blob")

    def encrypt_key_blobs(self, kek: Union[bytes, str]) -> bytes:
        """Encrypt key blobs with specified key.

        :param kek: key to encode key blobs
        :return: encrypted binary key blobs joined together
        """
        result = bytes()
        for key_blob in self._key_blobs:
            result += key_blob.export(kek)
        return align_block(
            result, 256
        )  # this is for compatibility with elftosb, probably need FLASH sector size

    def info(self) -> str:
        """Text info about the instance."""
        msg = "Key-Blob\n"
        for index, key_blob in enumerate(self._key_blobs):
            msg += f"Key-Blob {str(index)}:\n"
            msg += key_blob.info()
        return msg
