# @file
# Module contains helper classes and functions to work with UEFI FVs.
#
#
# Copyright (c) Microsoft Corporation
#
# SPDX-License-Identifier: BSD-2-Clause-Patent
##
"""Module containing helper classes and functions for working with UEFI Fvs."""

import struct
import sys
import uuid
from typing import IO

#
# UEFI GUIDs
#
EfiSystemNvDataFvGuid = uuid.UUID(fields=(0xFFF12B8D, 0x7696, 0x4C8B, 0xA9, 0x85, 0x2747075B4F50))


#
# UEFI #Defines
#
EFI_FVH_SIGNATURE = b"_FVH"


class EfiFirmwareVolumeHeader(object):
    """An object representing an EFI_FIRMWARE_VOLUME_HEADER.

    Can parse or produce an EFI_FIRMWARE_VOLUME_HEADER structure/byte buffer.

    ```
    typedef struct {
        UINT8                     ZeroVector[16];
        EFI_GUID                  FileSystemGuid;
        UINT64                    FvLength;
        UINT32                    Signature;
        EFI_FVB_ATTRIBUTES_2      Attributes;
        UINT16                    HeaderLength;
        UINT16                    Checksum;
        UINT16                    ExtHeaderOffset;
        UINT8                     Reserved[1];
        UINT8                     Revision;
        EFI_FV_BLOCK_MAP_ENTRY    BlockMap[1];
    } EFI_FIRMWARE_VOLUME_HEADER;
    ```
    """

    def __init__(self) -> "EfiFirmwareVolumeHeader":
        """Inits an empty object."""
        self.StructString = "=16s16sQ4sLHHHBBQQ"  # spell-checker: disable-line
        self.ZeroVector = None
        self.FileSystemGuid = None
        self.FvLength = None
        self.Attributes = None
        self.HeaderLength = None
        self.Checksum = None
        self.ExtHeaderOffset = None
        self.Reserved = None
        self.Revision = None
        self.Blockmap0 = None
        self.Blockmap1 = None

    def load_from_file(self, file: IO) -> "EfiFirmwareVolumeHeader":
        """Loads data into the object from a filestream.

        Args:
            file (obj): An open file that has been seeked to the correct location.

        Returns:
            (EfiFirmwareVolumeHeader): self

        Raises:
            (Exception): Invalid signature in fs
        """
        # 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.
        (
            self.ZeroVector,
            file_system_guid_bin,
            self.FvLength,
            self.Signature,
            self.Attributes,
            self.HeaderLength,
            self.Checksum,
            self.ExtHeaderOffset,
            self.Reserved,
            self.Revision,
            self.Blockmap0,
            self.Blockmap1,
        ) = struct.unpack(self.StructString, struct_bytes)

        # Make sure that this structure is what we think it is.
        if self.Signature != EFI_FVH_SIGNATURE:
            raise Exception("File does not appear to point to a valid EfiFirmwareVolumeHeader!")

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

        return self

    def serialize(self) -> bytes:
        r"""Serializes the object.

        Returns:
            (str): string representing packed data as bytes (i.e. b'\x01\x00\x03')
        """
        file_system_guid_bin = self.FileSystemGuid.bytes if sys.byteorder == "big" else self.FileSystemGuid.bytes_le
        return struct.pack(
            self.StructString,
            self.ZeroVector,
            file_system_guid_bin,
            self.FvLength,
            self.Signature,
            self.Attributes,
            self.HeaderLength,
            self.Checksum,
            self.ExtHeaderOffset,
            self.Reserved,
            self.Revision,
            self.Blockmap0,
            self.Blockmap1,
        )


class EfiFirmwareVolumeExtHeader(object):
    """An object representing an EFI_FIRMWARE_VOLUME_EXT_HEADER.

    Can parse or produce an EFI_FIRMWARE_VOLUME_EXT_HEADER structure/byte buffer.

    ```
    typedef struct {
        EFI_GUID                  FileSystemGuid;
        UINT32                    ExtHeaderSize;
    } EFI_FIRMWARE_VOLUME_EXT_HEADER;
    ```
    """

    def __init__(self) -> "EfiFirmwareVolumeExtHeader":
        """Inits an empty object."""
        self.StructString = "=16sL"
        self.FileSystemGuid = None
        self.ExtHeaderSize = None

    def load_from_file(self, file: IO) -> "EfiFirmwareVolumeExtHeader":
        """Loads data into the object from a filestream.

        Args:
            file: An open file that has been seeked to the correct location.

        Returns:
            (EfiFirmwareVolumeExtHeader): self
        """
        # 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.
        (self.FileSystemGuid, self.ExtHeaderSize) = struct.unpack(self.StructString, struct_bytes)

        return self
