#include "pch.h"

#define MAC_ADDRESS_SIZE 6
#define SMBIOS_UUID_SIZE 16

//
// Intentionally add 2 bytes of padding so that it is easy to check if the MAC
// address has been generated by casting to a ULONG64 and comparing to zero.
// The 2 bytes of padding will be zero initialized by the loader and will never
// be written, so should always remain zero.
//

__declspec(align(8))
UCHAR KdTargetMacAddress[MAC_ADDRESS_SIZE + 2];

NTSTATUS
GetSmBiosUuid(
    __in PLOADER_PARAMETER_BLOCK LoaderBlock,
    __out_bcount(SMBIOS_UUID_SIZE) PUCHAR SmBiosUuid
    )
/*++

Routine Description:

    This routine will get the UUID in SMBios Table.

Arguments:

    LoaderBlock -

    SmBiosUuid - Returned SMBIOS UUID

Return Value:

    NT Status code.

--*/
{
    NTSTATUS Status;
    PLOADER_PARAMETER_EXTENSION LoaderExtension;
    PSMBIOS3_EPS_HEADER SMBiosEPSHeader;
    PHYSICAL_ADDRESS SMBiosTablePhysicalAddress;
    ULONG SMBiosTableLength;
    PUCHAR SMBiosTable;
    PSMBIOS_SYSTEM_INFORMATION_STRUCT SMBiosInfo;
    ULONG SMBiosTablePages;
    ULONG Offset;

    Status = STATUS_SUCCESS;    

    if (LoaderBlock == NULL) {
        Status = STATUS_INVALID_PARAMETER;
        goto GetSmBiosUuidExit;
    }

    if (SmBiosUuid == NULL) {
        Status = STATUS_INVALID_PARAMETER;
        goto GetSmBiosUuidExit;
    }

    RtlZeroMemory(SmBiosUuid, SMBIOS_UUID_SIZE);

    LoaderExtension = LoaderBlock->Extension;
    if ((LoaderExtension == NULL) ||
        (LoaderExtension->Size < sizeof(LOADER_PARAMETER_EXTENSION))) {

        Status = STATUS_INVALID_PARAMETER;
        goto GetSmBiosUuidExit;
    }

    SMBiosEPSHeader = LoaderExtension->SMBiosEPSHeader;
    if (SMBiosEPSHeader == NULL) {
        Status = STATUS_INVALID_PARAMETER;
        goto GetSmBiosUuidExit;
    }

    //
    // Map SMBios Table to memory
    //

    SMBiosTablePhysicalAddress.QuadPart = SMBiosEPSHeader->StructureTableAddress;
    SMBiosTableLength = SMBiosEPSHeader->StructureTableMaximumSize;
    SMBiosTablePages = ADDRESS_AND_SIZE_TO_SPAN_PAGES(
        SMBiosTablePhysicalAddress.LowPart, SMBiosTableLength
        );
    SMBiosTable = KdMapPhysicalMemory64(SMBiosTablePhysicalAddress,
                                        SMBiosTablePages, 
                                        FALSE);

    if (SMBiosTable == NULL) {
        Status = STATUS_INSUFFICIENT_RESOURCES;
        goto GetSmBiosUuidExit;
    }

    //
    // Look for SMBios System Information
    //
    SMBiosInfo = NULL;
    Offset = 0;
    while (Offset < SMBiosTableLength) {
        PSMBIOS_STRUCT_HEADER Header;

        Header = (PSMBIOS_STRUCT_HEADER)&SMBiosTable[Offset];
        if (Header->Type == SMBIOS_SYSTEM_INFORMATION) {
            SMBiosInfo = (PSMBIOS_SYSTEM_INFORMATION_STRUCT)Header;
            break;
        }

        //
        // Move to next header, but include two null bytes
        //
        Offset += Header->Length + 2;

        // 
        // Just in case there is more data than header declared
        //
        while ((SMBiosTable[Offset - 2] != 0) ||
            (SMBiosTable[Offset - 1] != 0)) {
            Offset++;
            if (Offset >= SMBiosTableLength) break;
        }
    }

    //
    // Exit if SMBios System Information wasn't found
    //
    if ((SMBiosInfo == NULL) ||
        (SMBiosInfo->Length <= SMBIOS_SYSTEM_INFORMATION_LENGTH_20)) {

        Status = STATUS_NOT_FOUND;
        KdUnmapVirtualAddress(SMBiosTable, SMBiosTablePages, FALSE);
        goto GetSmBiosUuidExit;
    }

    RtlCopyMemory(SmBiosUuid, SMBiosInfo->Uuid, SMBIOS_UUID_SIZE);
    KdUnmapVirtualAddress(SMBiosTable, SMBiosTablePages, FALSE);

GetSmBiosUuidExit:
    return Status;
}

VOID
MacAddressFromSmBiosUuid(
    __in PLOADER_PARAMETER_BLOCK LoaderBlock,
    __out_bcount(MAC_ADDRESS_SIZE) PUCHAR MacAddress
    )
/*++

Routine Description:

    This routine will generate a pseudo-unique MAC address based on the UUID
    in SMBios Table. Goal is return each time same MAC address on particular
    address which will be different per device. To avoid DHCP servers overload
    we don't want to generate different values on each boot.

Arguments:

    LoaderBlock - 

    MacAddress - Returned MAC address

Return Value:

    None.

--*/
{
    NTSTATUS Status;
    UCHAR SmBiosUuid[SMBIOS_UUID_SIZE];
    SYMCRYPT_HMAC_SHA256_EXPANDED_KEY HmacKey;
    BYTE HmacResult[SYMCRYPT_HMAC_SHA256_RESULT_SIZE];

    //
    // Return any previously calculated MAC address.
    //

    if (*(PULONG64)KdTargetMacAddress != 0) {
        goto MacAddressFromSmBiosUuidExit;
    }

    //
    // Bail if no LoaderBlock was passed, since that is required to calculate
    // a reasonable address.
    //

    if (LoaderBlock == NULL) {
        return;
    }

    //
    // In case of failure we will try to use a PRIVATE MAC address
    //

    KdTargetMacAddress[0] = 0xAC;
    KdTargetMacAddress[1] = 0xDE;
    KdTargetMacAddress[2] = 0x48;
    KdTargetMacAddress[3] = 0x00;
    KdTargetMacAddress[4] = 0x00;
    KdTargetMacAddress[5] = 0x00;

    Status = GetSmBiosUuid(LoaderBlock, SmBiosUuid);
    if (!NT_SUCCESS(Status)) {
        goto MacAddressFromSmBiosUuidExit;
    }

    //
    // At this moment we have table with UUID, to avoid problems on systems
    // where UUID doesn't have unique node value (see SMBIOS reference or
    // RFC4122) we will use SHA256 to generate unique value.
    //
    // In addition we mark generated address as locally administred and clear
    // multicast bit.
    //

    SymCryptHmacSha256ExpandKey(&HmacKey, "MAC_KEY", sizeof("MAC_KEY"));
    SymCryptHmacSha256(&HmacKey, SmBiosUuid, sizeof(SmBiosUuid), HmacResult);
    RtlCopyMemory(KdTargetMacAddress, HmacResult, MAC_ADDRESS_SIZE);
    SymCryptWipeKnownSize(&HmacKey, sizeof(SYMCRYPT_HMAC_SHA256_EXPANDED_KEY));
    KdTargetMacAddress[0] |= 0x02;
    KdTargetMacAddress[0] &= ~0x01;

MacAddressFromSmBiosUuidExit:
    RtlCopyMemory(MacAddress, KdTargetMacAddress, MAC_ADDRESS_SIZE);
    return;
}
