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

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

import logging
import os
from copy import deepcopy
from struct import pack
from typing import Any, Dict, List, Optional, Union

from crcmod.predefined import mkPredefinedCrcFun
from cryptography.hazmat.primitives.keywrap import aes_key_wrap

from spsdk import SPSDKError
from spsdk import version as spsdk_version
from spsdk.apps.utils.utils import filepath_from_config, get_key
from spsdk.exceptions import SPSDKValueError
from spsdk.utils.database import Database
from spsdk.utils.exceptions import SPSDKRegsErrorBitfieldNotFound
from spsdk.utils.images import BinaryImage
from spsdk.utils.misc import (
    align_block,
    load_binary,
    reverse_bits_in_bytes,
    split_data,
    value_to_bytes,
    value_to_int,
)
from spsdk.utils.registers import Registers
from spsdk.utils.schema_validator import ConfigTemplate, ValidationSchemas

from . import OTFAD_DATA_FOLDER, OTFAD_DATABASE_FILE, OTFAD_SCH_FILE, Counter, crypto_backend

logger = logging.getLogger(__name__)


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, SRTADDR[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
    # Encryption block size
    _ENCRYPTION_BLOCK_SIZE = 16

    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 binaries: optional data chunks of this key blob
        :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)
        if self.end_addr or self.key_flags:
            end_addr_with_flags = (
                ((self.end_addr - 1) & ~self._KEY_FLAG_MASK) | self.key_flags | self._END_ADDR_MASK
            )
        else:
            end_addr_with_flags = 0
        result += pack("<I", end_addr_with_flags)
        header_crc: bytes = mkPredefinedCrcFun("crc-32-mpeg")(result).to_bytes(4, "little")
        # 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 += header_crc
        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),
        byte_swap_cnt: int = 0,
    ) -> 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
        :param byte_swap_cnt: Encrypted keyblob reverse byte count, 0 means NO reversing is enabled
        :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")

        blobs = bytes()
        wrap = aes_key_wrap(kek, plaintext[:40])
        if byte_swap_cnt > 0:
            for i in range(0, len(wrap), byte_swap_cnt):
                blobs += wrap[i : i + byte_swap_cnt][::-1]
        else:
            blobs += wrap

        return align_block(
            blobs, 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,
        counter_value: Optional[int] = None,
    ) -> 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
        :param counter_value: Optional counter value, if not specified start address of keyblob will be used
        :return: encrypted data
        :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._ENCRYPTION_BLOCK_SIZE)  # align data length
        data_len = len(data)

        # check start and end addresses
        # Support dual image boot, do not raise exception
        if not self.matches_range(base_address, base_address + data_len - 1):
            logger.warning(
                f"Image address range is not within key blob: "
                f"{hex(self.start_addr)}-{hex(self.end_addr)}."
                " Ignore this if flash remap feature is used"
            )
        result = bytes()
        # SPSDK-936 change base address to start address of keyblob to allow flash remap on RT5xx/RT6xx
        if not counter_value:
            counter_value = self.start_addr

        counter = Counter(
            self._get_ctr_nonce(), ctr_value=counter_value, 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)

    @property
    def is_encrypted(self) -> bool:
        """Get the required encryption or not.

        :return: True if blob is encrypted, False otherwise.
        """
        return (bool)(
            (self.key_flags & (self.KEY_FLAG_ADE | self.KEY_FLAG_VLD))
            == (self.KEY_FLAG_ADE | self.KEY_FLAG_VLD)
        )


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

    OTFAD_DATA_UNIT = 0x400

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

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

    def __setitem__(self, index: int, value: KeyBlob) -> None:
        self._key_blobs.remove(self._key_blobs[index])
        self._key_blobs.insert(index, value)

    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 with all available keyblobs.

        :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
        """
        encrypted_data = bytearray(image)
        addr = base_addr
        for block in split_data(image, self.OTFAD_DATA_UNIT):
            for key_blob in self._key_blobs:
                if key_blob.matches_range(addr, addr + len(block)):
                    logger.debug(
                        f"Encrypting {hex(addr)}:{hex(len(block) + addr)}"
                        f" with keyblob: \n {key_blob.info()}"
                    )
                    encrypted_data[
                        addr - base_addr : len(block) + addr - base_addr
                    ] = key_blob.encrypt_image(addr, block, byte_swap, counter_value=addr)
            addr += len(block)

        return bytes(encrypted_data)

    def get_key_blobs(self) -> bytes:
        """Get key blobs.

        :return: Binary key blobs joined together
        """
        result = bytes()
        for key_blob in self._key_blobs:
            result += key_blob.plain_data()
        return align_block(
            result, 256
        )  # this is for compatibility with elftosb, probably need FLASH sector size

    def encrypt_key_blobs(
        self,
        kek: Union[bytes, str],
        key_scramble_mask: Optional[int] = None,
        key_scramble_align: Optional[int] = None,
        byte_swap_cnt: int = 0,
    ) -> bytes:
        """Encrypt key blobs with specified key.

        :param kek: key to encode key blobs
        :param key_scramble_mask: 32-bit scramble key, if KEK scrambling is desired.
        :param key_scramble_align: 8-bit scramble align, if KEK scrambling is desired.
        :param byte_swap_cnt: Encrypted keyblob reverse byte count, 0 means NO reversing is enabled
        :raises SPSDKValueError: Invalid input value.
        :return: encrypted binary key blobs joined together
        """
        if isinstance(kek, str):
            kek = bytes.fromhex(kek)
        scramble_enabled = key_scramble_mask is not None and key_scramble_align is not None
        if scramble_enabled:
            assert key_scramble_mask and key_scramble_align
            if key_scramble_mask >= 1 << 32:
                raise SPSDKValueError("OTFAD Key scramble mask has invalid length")
            if key_scramble_align >= 1 << 8:
                raise SPSDKValueError("OTFAD Key scramble align has invalid length")

            logger.debug("The scrambling of keys is enabled.")
            key_scramble_mask_inv = reverse_bits_in_bytes(
                key_scramble_mask.to_bytes(4, byteorder="big")
            )
            logger.debug(f"The inverted scramble key is: {key_scramble_mask_inv.hex()}")
        result = bytes()
        scrambled = bytes()
        for i, key_blob in enumerate(self._key_blobs):
            if scramble_enabled:
                assert key_scramble_mask and key_scramble_align
                scrambled = bytearray(kek)
                long_ix = (key_scramble_align >> (i * 2)) & 0x03
                for j in range(4):
                    scrambled[(long_ix * 4) + j] ^= key_scramble_mask_inv[j]

            logger.debug(
                f"Used KEK for keyblob{i} encryption is: {scrambled.hex() if scramble_enabled else kek.hex()}"
            )

            result += key_blob.export(
                scrambled if scramble_enabled else kek, byte_swap_cnt=byte_swap_cnt
            )
        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


class OtfadNxp(Otfad):
    """OTFAD: On-the-Fly AES Decryption Module with reflecting of NXP parts."""

    def __init__(
        self,
        family: str,
        kek: Union[bytes, str],
        table_address: int = 0,
        key_blobs: Optional[List[KeyBlob]] = None,
        key_scramble_mask: Optional[int] = None,
        key_scramble_align: Optional[int] = None,
        binaries: Optional[BinaryImage] = None,
    ) -> None:
        """Constructor.

        :param family: Device family
        :param kek: KEK to encrypt OTFAD table
        :param table_address: Absolute address of OTFAD table.
        :param key_blobs: Optional Key blobs to add to OTFAD, defaults to None
        :param key_scramble_mask: If defined, the key scrambling algorithm will be applied.
            ('key_scramble_align' must be defined also)
        :param key_scramble_align: If defined, the key scrambling algorithm will be applied.
            ('key_scramble_mask' must be defined also)
        :raises SPSDKValueError: Unsupported family
        """
        super().__init__()

        if family not in self.get_supported_families():
            raise SPSDKValueError(f"Unsupported family{family} by OTFAD")

        if (key_scramble_align is None and key_scramble_mask) or (
            key_scramble_align and key_scramble_mask is None
        ):
            raise SPSDKValueError("Key Scrambling is not fully defined")

        self.family = family
        self.kek = bytes.fromhex(kek) if isinstance(kek, str) else kek
        self.key_scramble_mask = key_scramble_mask
        self.key_scramble_align = key_scramble_align
        self.table_address = table_address
        self.database = Database(OTFAD_DATABASE_FILE)
        self.blobs_min_cnt = self.database.get_device_value("key_blob_min_cnt", device=family)
        self.blobs_max_cnt = self.database.get_device_value("key_blob_max_cnt", device=family)
        self.byte_swap = self.database.get_device_value("byte_swap", device=family)
        self.key_blob_rec_size = self.database.get_device_value("key_blob_rec_size", device=family)
        self.keyblob_byte_swap_cnt = self.database.get_device_value(
            "keyblob_byte_swap_cnt", device=family
        )
        assert self.keyblob_byte_swap_cnt in [0, 2, 4, 8, 16]
        self.binaries = binaries

        if key_blobs:
            for key_blob in key_blobs:
                self.add_key_blob(key_blob)

        # Just fill up the minimum count of key blobs
        while len(self._key_blobs) < self.blobs_min_cnt:
            self.add_key_blob(
                KeyBlob(
                    start_addr=0,
                    end_addr=0,
                    key=bytes([0] * KeyBlob.KEY_SIZE),
                    counter_iv=bytes([0] * KeyBlob.CTR_SIZE),
                    key_flags=0,
                    zero_fill=bytes([0] * 4),
                )
            )

    @staticmethod
    def get_blhost_script_otp_keys(
        family: str, otp_master_key: bytes, otfad_key_seed: bytes
    ) -> str:
        """Create BLHOST script to load fuses needed to run OTFAD with OTP fuses.

        :param family: Device family.
        :param otp_master_key: OTP Master Key.
        :param otfad_key_seed: OTFAD Key Seed.
        :return: BLHOST script that loads the keys into fuses.
        """
        database = Database(OTFAD_DATABASE_FILE)
        xml_fuses = database.get_device_value("reg_fuses", device=family, default=None)
        if not xml_fuses:
            logger.debug(f"The {family} has no OTFAD fuses definition")
            return ""

        xml_fuses = os.path.join(OTFAD_DATA_FOLDER, xml_fuses)

        fuses = Registers(family, base_endianness="little")
        grouped_regs = database.get_device_value("grouped_registers", device=family, default=None)
        fuses.load_registers_from_xml(xml_fuses, grouped_regs=grouped_regs)
        reg_omk = fuses.find_reg("OTP_MASTER_KEY")
        reg_oks = fuses.find_reg("OTFAD_KEK_SEED")
        reg_omk.set_value(otp_master_key)
        reg_oks.set_value(otfad_key_seed)
        ret = (
            "# BLHOST OTFAD keys fuse programming script\n"
            f"# Generated by SPSDK {spsdk_version}\n"
            f"# Chip: {family}\n\n"
        )

        ret += f"# OTP MASTER KEY(Big Endian): {reg_omk.get_hex_value()}\n\n"
        for reg in reg_omk.sub_regs:
            ret += f"# {reg.name} fuse.\n"
            ret += f"efuse-program-once {reg.offset} 0x{reg.get_bytes_value().hex()} --no-verify\n"

        ret += f"\n# OTFAD KEK SEED (Big Endian): {reg_oks.get_hex_value()}\n\n"
        for reg in reg_oks.sub_regs:
            ret += f"# {reg.name} fuse.\n"
            ret += f"efuse-program-once {reg.offset} 0x{reg.get_bytes_value().hex()} --no-verify\n"

        return ret

    def get_blhost_script_otp_kek(self, index: int = 1) -> str:
        """Create BLHOST script to load fuses needed to run OTFAD with OTP fuses just for OTFAD key.

        :param index: Index of OTFAD peripheral [1, 2, ..., n].
        :return: BLHOST script that loads the keys into fuses.
        """
        if not self.database.get_device_value("has_kek_fuses", self.family, default=False):
            logger.debug(f"The {self.family} has no OTFAD KEK fuses")
            return ""

        peripheral_list = self.database.get_device_value("peripheral_list", device=self.family)
        if str(index) not in peripheral_list:
            logger.debug(f"The {self.family} has no OTFAD{index} peripheral")
            return ""

        filter_out_list = [f"OTFAD{i}" for i in peripheral_list if str(index) != i]
        xml_fuses = self.database.get_device_value("reg_fuses", device=self.family, default=None)
        if not xml_fuses:
            logger.debug(f"The {self.family} has no OTFAD fuses definition")
            return ""

        xml_fuses = os.path.join(OTFAD_DATA_FOLDER, xml_fuses)

        fuses = Registers(self.family, base_endianness="little")
        self.database.index = index

        grouped_regs = self.database.get_device_value(
            "grouped_registers", device=self.family, default=None
        )

        fuses.load_registers_from_xml(xml_fuses, filter_out_list, grouped_regs)

        scramble_enabled = (
            self.key_scramble_mask is not None and self.key_scramble_align is not None
        )

        otfad_key_fuse = self.database.get_device_value("otfad_key_fuse", self.family)
        otfad_cfg_fuse = self.database.get_device_value("otfad_cfg_fuse", self.family)

        fuses.find_reg(otfad_key_fuse).set_value(self.kek)
        otfad_cfg = fuses.find_reg(otfad_cfg_fuse)

        try:
            otfad_cfg.find_bitfield(
                self.database.get_device_value("otfad_enable_bitfield", self.family)
            ).set_value(1)
        except SPSDKRegsErrorBitfieldNotFound:
            logger.debug(f"Bitfield for OTFAD ENABLE not found for {self.family}")

        if scramble_enabled:
            otfad_cfg.find_bitfield(
                self.database.get_device_value("otfad_scramble_enable_bitfield", self.family)
            ).set_value(1)
            otfad_cfg.find_bitfield(
                self.database.get_device_value("otfad_scramble_align_bitfield", self.family)
            ).set_value(self.key_scramble_align)
            fuses.find_reg(
                self.database.get_device_value("otfad_scramble_key", self.family)
            ).set_value(self.key_scramble_mask)

        ret = (
            f"# BLHOST OTFAD{index} KEK fuses programming script\n"
            f"# Generated by SPSDK {spsdk_version}\n"
            f"# Chip: {self.family}, peripheral: OTFAD{index} !\n\n"
        )

        ret += f"# OTP KEK (Big Endian): {self.kek.hex()}\n\n"
        for reg in fuses.find_reg(otfad_key_fuse).sub_regs:
            ret += f"# {reg.name} fuse.\n"
            ret += f"efuse-program-once {reg.offset} 0x{reg.get_bytes_value().hex()} --no-verify\n"

        ret += f"\n\n# {otfad_cfg.name} fuse.\n"
        for bitfield in otfad_cfg.get_bitfields():
            ret += f"#   {bitfield.name}: {bitfield.get_enum_value()}\n"
        ret += f"efuse-program-once {otfad_cfg.offset} 0x{otfad_cfg.get_bytes_value().hex()} --no-verify\n"

        if scramble_enabled:
            scramble = fuses.find_reg(
                self.database.get_device_value("otfad_scramble_key", self.family)
            )
            ret += f"\n# {scramble.name} fuse.\n"
            ret += f"efuse-program-once {scramble.offset} 0x{scramble.get_bytes_value().hex()} --no-verify\n"

        return ret

    def export_image(
        self,
        plain_data: bool = False,
        swap_bytes: bool = False,
        join_sub_images: bool = True,
        table_address: int = 0,
    ) -> Optional[BinaryImage]:
        """Get the OTFAD Key Blob Binary Image representation.

        :param plain_data: Binary representation in plain data format, defaults to False
        :param swap_bytes: For some platforms the swap bytes is needed in encrypted format, defaults to False.
        :param join_sub_images: If it's True, all the binary sub-images are joined into one, defaults to True.
        :param table_address: Absolute address of OTFAD table.
        :return: OTFAD key blob data in BinaryImage.
        """
        if self.binaries is None:
            return None
        binaries: BinaryImage = deepcopy(self.binaries)
        for binary in binaries.sub_images:
            if binary.binary:
                binary.binary = align_block(binary.binary, KeyBlob._ENCRYPTION_BLOCK_SIZE)
            for segment in binary.sub_images:
                if segment.binary:
                    segment.binary = align_block(segment.binary, KeyBlob._ENCRYPTION_BLOCK_SIZE)

        binaries.validate()

        if not plain_data:
            for binary in binaries.sub_images:
                if binary.binary:
                    binary.binary = self.encrypt_image(
                        binary.binary,
                        table_address + binary.absolute_address,
                        swap_bytes,
                    )
                for segment in binary.sub_images:
                    if segment.binary:
                        segment.binary = self.encrypt_image(
                            segment.binary,
                            segment.absolute_address + table_address,
                            swap_bytes,
                        )

        if join_sub_images:
            binaries.join_images()
            binaries.validate()

        return binaries

    def binary_image(
        self,
        plain_data: bool = False,
        data_alignment: int = 16,
        otfad_table_name: str = "OTFAD_Table",
    ) -> BinaryImage:
        """Get the OTFAD Binary Image representation.

        :param plain_data: Binary representation in plain format, defaults to False
        :param data_alignment: Alignment of data part key blobs.
        :param otfad_table_name: name of the output file that contains OTFAD table
        :return: OTFAD in BinaryImage.
        """
        otfad = BinaryImage("OTFAD", offset=self.table_address)
        # Add mandatory OTFAD table
        otfad_table = (
            self.get_key_blobs()
            if plain_data
            else self.encrypt_key_blobs(
                self.kek,
                self.key_scramble_mask,
                self.key_scramble_align,
                self.keyblob_byte_swap_cnt,
            )
        )
        otfad.add_image(
            BinaryImage(
                otfad_table_name,
                size=self.key_blob_rec_size * self.blobs_max_cnt,
                offset=0,
                description=f"OTFAD description table for {self.family}",
                binary=otfad_table,
                alignment=256,
            )
        )
        binaries = self.export_image(table_address=self.table_address)

        if binaries:
            binaries.alignment = data_alignment
            binaries.validate()
            otfad.add_image(binaries)
        return otfad

    @staticmethod
    def get_supported_families() -> List[str]:
        """Get all supported families for AHAB container.

        :return: List of supported families.
        """
        database = Database(OTFAD_DATABASE_FILE)
        return database.devices.device_names

    @staticmethod
    def get_validation_schemas(family: str) -> List[Dict[str, Any]]:
        """Get list of validation schemas.

        :param family: Family for which the template should be generated.
        :return: Validation list of schemas.
        """
        if family not in OtfadNxp.get_supported_families():
            return []

        database = Database(OTFAD_DATABASE_FILE)
        schemas = ValidationSchemas.get_schema_file(OTFAD_SCH_FILE)
        family_sch = deepcopy(schemas["otfad_family"])
        family_sch["properties"]["family"]["enum"] = OtfadNxp.get_supported_families()
        family_sch["properties"]["family"]["template_value"] = family
        ret = [family_sch, schemas["otfad_output"], schemas["otfad"]]
        additional_schemes = database.get_device_value("additional_template", family, default=[])
        ret.extend([schemas[x] for x in additional_schemes])
        return ret

    @staticmethod
    def get_validation_schemas_family() -> List[Dict[str, Any]]:
        """Get list of validation schemas for family key.

        :return: Validation list of schemas.
        """
        schemas = ValidationSchemas.get_schema_file(OTFAD_SCH_FILE)
        family_sch = deepcopy(schemas["otfad_family"])
        family_sch["properties"]["family"]["enum"] = OtfadNxp.get_supported_families()
        return [family_sch]

    @staticmethod
    def generate_config_template(family: str) -> Dict[str, Any]:
        """Generate OTFAD configuration template.

        :param family: Family for which the template should be generated.
        :return: Dictionary of individual templates (key is name of template, value is template itself).
        """
        val_schemas = OtfadNxp.get_validation_schemas(family)
        database = Database(OTFAD_DATABASE_FILE)

        if val_schemas:
            template_note = database.get_device_value(
                "additional_template_text", family, default=""
            )
            title = f"On-The-Fly AES decryption Configuration template for {family}."

            yaml_data = ConfigTemplate(title, val_schemas, note=template_note).export_to_yaml()

            return {f"{family}_otfad": yaml_data}

        return {}

    @staticmethod
    def load_from_config(
        config: Dict[str, Any], config_dir: str, search_paths: Optional[List[str]] = None
    ) -> "OtfadNxp":
        """Converts the configuration option into an OTFAD image object.

        "config" content array of containers configurations.

        :param config: array of OTFAD configuration dictionaries.
        :param config_dir: directory where the config is located
        :param search_paths: List of paths where to search for the file, defaults to None
        :return: initialized OTFAD object.
        """
        otfad_config: List[Dict[str, Any]] = config["key_blobs"]
        family = config["family"]
        database = Database(OTFAD_DATABASE_FILE)
        kek = get_key(config["kek"], expected_size=16, search_paths=search_paths)
        logger.debug(f"Loaded KEK: {kek.hex()}")
        table_address = value_to_int(config["otfad_table_address"])
        start_address = min([value_to_int(addr["start_address"]) for addr in otfad_config])

        key_scramble_mask = None
        key_scramble_align = None
        if database.get_device_value("supports_key_scrambling", device=family, default=False):
            if "key_scramble" in config.keys():
                key_scramble = config["key_scramble"]
                key_scramble_mask = value_to_int(key_scramble["key_scramble_mask"])
                key_scramble_align = value_to_int(key_scramble["key_scramble_align"])

        data_blobs: Optional[List[Dict]] = config.get("data_blobs")
        binaries = None
        if data_blobs:
            # pylint: disable-next=nested-min-max
            start_address = min(
                min([value_to_int(addr["address"]) for addr in data_blobs]),
                start_address,
            )
            binaries = BinaryImage(
                filepath_from_config(
                    config, "encrypted_name", "encrypted_blobs", config_dir, config["output_folder"]
                ),
                offset=start_address - table_address,
            )
            for data_blob in data_blobs:
                data = load_binary(data_blob["data"], search_paths=search_paths)
                address = value_to_int(data_blob["address"])

                binary = BinaryImage(
                    os.path.basename(data_blob["data"]),
                    offset=address - table_address - binaries.offset,
                    binary=data,
                )
                binaries.add_image(binary)

        otfad = OtfadNxp(
            family=family,
            kek=kek,
            table_address=table_address,
            key_scramble_align=key_scramble_align,
            key_scramble_mask=key_scramble_mask,
            binaries=binaries,
        )

        for i, key_blob_cfg in enumerate(otfad_config):
            aes_key = value_to_bytes(key_blob_cfg["aes_key"], byte_cnt=KeyBlob.KEY_SIZE)
            aes_ctr = value_to_bytes(key_blob_cfg["aes_ctr"], byte_cnt=KeyBlob.CTR_SIZE)
            start_addr = value_to_int(key_blob_cfg["start_address"])
            end_addr = value_to_int(key_blob_cfg["end_address"])
            aes_decryption_enable = key_blob_cfg.get("aes_decryption_enable", True)
            valid = key_blob_cfg.get("valid", True)
            read_only = key_blob_cfg.get("read_only", True)
            flags = 0
            if aes_decryption_enable:
                flags |= KeyBlob.KEY_FLAG_ADE
            if valid:
                flags |= KeyBlob.KEY_FLAG_VLD
            if read_only:
                flags |= KeyBlob.KEY_FLAG_READ_ONLY

            otfad[i] = KeyBlob(
                start_addr=start_addr,
                end_addr=end_addr,
                key=aes_key,
                counter_iv=aes_ctr,
                key_flags=flags,
                zero_fill=bytes([0] * 4),
            )

        return otfad
