# @file variable_format.py
# Module contains helper classes and functions to work with UEFI Variables.
#
# Copyright (c) Microsoft Corporation
#
# SPDX-License-Identifier: BSD-2-Clause-Patent
##
"""Module contains helper classes and functions to work with UEFI Variables."""

import struct
import sys
import uuid

import edk2toollib.uefi.uefi_multi_phase as ump

#
# UEFI GUIDs
#
EfiVariableGuid = uuid.UUID(fields=(0xDDCF3616, 0x3275, 0x4164, 0x98, 0xB6, 0xFE85707FFE7D))
EfiAuthenticatedVariableGuid = uuid.UUID(fields=(0xAAF32C78, 0x947B, 0x439A, 0xA1, 0x80, 0x2E144EC37792))

#
# UEFI #Defines
#
HEADER_ALIGNMENT = 4
VARIABLE_STORE_FORMATTED = 0x5A
VARIABLE_STORE_HEALTHY = 0xFE
VARIABLE_DATA = 0x55AA
VAR_IN_DELETED_TRANSITION = 0xFE  # Variable is in obsolete transition.
VAR_DELETED = 0xFD  # Variable is obsolete.
VAR_HEADER_VALID_ONLY = 0x7F  # Variable header has been valid.
VAR_ADDED = 0x3F  # Variable has been completely added.


class VariableStoreHeader(object):
    """Object representing the VARIABLE_STORE_HEADER struct.

    Can parse or produce an VARIABLE_STORE_HEADER structure/byte buffer.

    typedef struct {
        EFI_GUID  Signature;
        UINT32  Size;
        UINT8   Format;
        UINT8   State;
        UINT16  Reserved;
        UINT32  Reserved1;
    } VARIABLE_STORE_HEADER;
    """

    def __init__(self) -> "VariableStoreHeader":
        """Init the empty structure."""
        self.StructString = "=16sLBBHL"  # spell-checker: disable-line
        self.StructSize = struct.calcsize(self.StructString)
        self.Signature = None
        self.Size = None
        self.Format = None
        self.State = None
        self.Reserved0 = None
        self.Reserved1 = None
        self.Type = "Var"

    def load_from_file(self, file: str) -> "VariableStoreHeader":
        """Load the structure from a file."""
        # This function assumes that the file has been seeked
        # to the correct starting location.
        orig_seek = file.tell()
        struct_bytes = file.read(struct.calcsize(self.StructString))
        file.seek(orig_seek)

        # Load this object with the contents of the data.
        (signature_bin, self.Size, self.Format, self.State, self.Reserved0, self.Reserved1) = struct.unpack(
            self.StructString, struct_bytes
        )

        # Update the GUID to be a UUID object.
        if sys.byteorder == "big":
            self.Signature = uuid.UUID(bytes=signature_bin)
        else:
            self.Signature = uuid.UUID(bytes_le=signature_bin)

        # Check one last thing.
        if self.Signature != EfiVariableGuid and self.Signature != EfiAuthenticatedVariableGuid:
            raise Exception("VarStore is of unknown type! %s" % self.Signature)
        if self.Signature == EfiAuthenticatedVariableGuid:
            self.Type = "AuthVar"

        return self

    def serialize(self) -> bytes:
        """Serialize the structure."""
        signature_bin = self.Signature.bytes if sys.byteorder == "big" else self.Signature.bytes_le
        return struct.pack(
            self.StructString, signature_bin, self.Size, self.Format, self.State, self.Reserved0, self.Reserved1
        )


#
# TODO: VariableHeader and AuthenticatedVariableHeader are not truly
#       header structures. They're entire variables. This code should be
#       cleaned up.
#
class VariableHeader(object):
    """Object representing the VARIABLE_HEADER struct.

    Can parse or produce an VARIABLE_HEADER structure/byte buffer.

    typedef struct {
        UINT16      StartId;
        UINT8       State;
        UINT8       Reserved;
        UINT32      Attributes;
        UINT32      NameSize;
        UINT32      DataSize;
        EFI_GUID    VendorGuid;
    } VARIABLE_HEADER;
    """

    def __init__(self) -> "VariableHeader":
        """Init the structure."""
        self.StructString = "=HBBLLL16s"  # spell-checker: disable-line
        self.StructSize = struct.calcsize(self.StructString)
        self.StartId = VARIABLE_DATA
        self.State = VAR_ADDED
        self.Attributes = ump.EFI_VARIABLE_NON_VOLATILE | ump.EFI_VARIABLE_BOOTSERVICE_ACCESS
        self.NameSize = 0
        self.DataSize = 0
        self.VendorGuid = uuid.uuid4()
        self.Name = None
        self.Data = None

    def populate_structure_fields(self, in_bytes: bytes) -> None:
        """Populate the structure field from bytes."""
        (self.StartId, self.State, reserved, self.Attributes, self.NameSize, self.DataSize, self.VendorGuid) = (
            struct.unpack(self.StructString, in_bytes)
        )

    def load_from_bytes(self, in_bytes: bytes) -> "VariableHeader":
        """Load the structure from a bytes."""
        # Load this object with the contents of the data.
        self.populate_structure_fields(in_bytes[0 : self.StructSize])

        # Update the GUID to be a UUID object.
        if sys.byteorder == "big":
            self.VendorGuid = uuid.UUID(bytes=self.VendorGuid)
        else:
            self.VendorGuid = uuid.UUID(bytes_le=self.VendorGuid)

        # Before loading data, make sure that this is a valid variable.
        if self.StartId != VARIABLE_DATA:
            raise EOFError("No variable data!")

        # Finally, load the data.
        data_offset = self.StructSize
        self.Name = in_bytes[data_offset : (data_offset + self.NameSize)].decode("utf-16")
        self.Name = self.Name[:-1]  # Strip the terminating char.
        data_offset += self.NameSize
        self.Data = in_bytes[data_offset : (data_offset + self.DataSize)]

        return self

    def load_from_file(self, file: str) -> "VariableHeader":
        """Load the struct from a file."""
        # This function assumes that the file has been seeked
        # to the correct starting location.
        orig_seek = file.tell()
        struct_bytes = file.read(struct.calcsize(self.StructString))

        # Load this object with the contents of the data.
        self.populate_structure_fields(struct_bytes)

        # Update the GUID to be a UUID object.
        if sys.byteorder == "big":
            self.VendorGuid = uuid.UUID(bytes=self.VendorGuid)
        else:
            self.VendorGuid = uuid.UUID(bytes_le=self.VendorGuid)

        # Before loading data, make sure that this is a valid variable.
        if self.StartId != VARIABLE_DATA:
            file.seek(orig_seek)
            raise EOFError("No variable data!")

        # Finally, load the data.
        self.Name = file.read(self.NameSize).decode("utf-16")[:-1]  # Strip the terminating char.
        self.Data = file.read(self.DataSize)

        file.seek(orig_seek)
        return self

    def get_buffer_data_size(self) -> int:
        """Get the buffer data size."""
        return self.StructSize + self.NameSize + self.DataSize

    def get_buffer_padding_size(self) -> int:
        """Get the buffer padding size."""
        buffer_data_size = self.get_buffer_data_size()
        padding_size = 0
        if buffer_data_size % HEADER_ALIGNMENT != 0:
            padding_size += HEADER_ALIGNMENT - (buffer_data_size % HEADER_ALIGNMENT)
        return padding_size

    def get_buffer_size(self) -> int:
        """Get the buffer size."""
        return self.get_buffer_data_size() + self.get_buffer_padding_size()

    def get_packed_name(self) -> bytes:
        """Get the name attribute."""
        # Make sure to replace the terminating char.
        # name_bytes = b"\x00".join([char for char in (self.Name + b'\x00')])
        name_bytes = self.Name.encode("utf-16")

        # Python encode will leave an "0xFFFE" on the front
        # to declare the encoding type. UEFI does not use this.
        name_bytes = name_bytes[2:]

        # Python encode skips the terminating character, so let's add that.
        name_bytes += b"\x00\x00"

        return name_bytes

    def set_name(self, new_name: str) -> None:
        """Set the name attribute."""
        self.Name = new_name
        self.NameSize = len(self.get_packed_name())

    def set_data(self, new_data: bytes) -> None:
        """Set the data attribute."""
        self.Data = new_data
        self.DataSize = len(new_data)

    def pack_struct(self) -> bytes:
        """Pack the object."""
        vendor_guid = self.VendorGuid.bytes if sys.byteorder == "big" else self.VendorGuid.bytes_le
        return struct.pack(
            self.StructString, self.StartId, self.State, 0, self.Attributes, self.NameSize, self.DataSize, vendor_guid
        )

    def serialize(self, with_padding: bool = False) -> bytes:
        """Serialize the object."""
        bytes = self.pack_struct()

        # Now add the name and data.
        bytes += self.get_packed_name()
        bytes += self.Data

        # Add padding if necessary.
        if with_padding:
            bytes += b"\xff" * self.get_buffer_padding_size()

        return bytes


class AuthenticatedVariableHeader(VariableHeader):
    """Object representing the AUTHENTICATED_VARIABLE_HEADER struct.

    Can parse or produce an AUTHENTICATED_VARIABLE_HEADER structure/byte buffer.

    typedef struct {
        UINT16      StartId;
        UINT8       State;
        UINT8       Reserved;
        UINT32      Attributes;
        UINT64      MonotonicCount;
        EFI_TIME    TimeStamp;
        UINT32      PubKeyIndex;
        UINT32      NameSize;
        UINT32      DataSize;
        EFI_GUID    VendorGuid;
    } AUTHENTICATED_VARIABLE_HEADER;
    """

    def __init__(self) -> "AuthenticatedVariableHeader":
        """Initializes the struct."""
        super(AuthenticatedVariableHeader, self).__init__()
        self.StructString = "=HBBLQ16sLLL16s"  # spell-checker: disable-line
        self.StructSize = struct.calcsize(self.StructString)
        self.MonotonicCount = 0
        self.TimeStamp = b""
        self.PubKeyIndex = 0

    def populate_structure_fields(self, in_bytes: bytes) -> None:
        """Populates the struct."""
        (
            self.StartId,
            self.State,
            reserved,
            self.Attributes,
            self.MonotonicCount,
            self.TimeStamp,
            self.PubKeyIndex,
            self.NameSize,
            self.DataSize,
            self.VendorGuid,
        ) = struct.unpack(self.StructString, in_bytes)

    def pack_struct(self, with_padding: bool = False) -> bytes:
        """Packs the struct."""
        vendor_guid = self.VendorGuid.bytes if sys.byteorder == "big" else self.VendorGuid.bytes_le
        return struct.pack(
            self.StructString,
            self.StartId,
            self.State,
            0,
            self.Attributes,
            self.MonotonicCount,
            self.TimeStamp,
            self.PubKeyIndex,
            self.NameSize,
            self.DataSize,
            vendor_guid,
        )


if __name__ == "__main__":
    pass
