/*++

Copyright (c) 2024, Quinn Stephens.
Provided under the BSD 3-Clause license.

Module Name:

   efifw.c

Abstract:

   Provides EFI firmware utilities.

--*/

#include <ntrtl.h>
#include "bootlib.h"
#include "efi.h"
#include "efilib.h"

BOOT_FIRMWARE_DATA EfiFirmwareData;
PBOOT_FIRMWARE_DATA EfiFirmwareParameters;
EFI_HANDLE EfiImageHandle;
EFI_SYSTEM_TABLE *EfiST;
EFI_BOOT_SERVICES *EfiBS;
EFI_RUNTIME_SERVICES *EfiRT;
SIMPLE_TEXT_OUTPUT_INTERFACE *EfiConOut;
SIMPLE_INPUT_INTERFACE *EfiConIn;
EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *EfiConInEx;

NTSTATUS
BlpFwInitialize (
    IN ULONG Stage,
    IN PBOOT_FIRMWARE_DATA FirmwareData
    )

/*++

Routine Description:

    Internal routine to initialize the boot library.

Arguments:

    Stage - Which stage of initialization to perform.

        Stage 0: Initialize global firmware-related data and pointers.
            Once this stage is complete, ConsolePrint() and ConsolePrintf() can be used.

    FirmwareData - Pointer to BOOT_FIRMWARE_DATA.

Return Value:

    STATUS_SUCCESS if successful,
    STATUS_INVALID_PARAMETER if FirmwareData is invalid.

--*/

{
    NTSTATUS Status;
    EFI_KEY_TOGGLE_STATE KeyToggleState;

    if (FirmwareData == NULL || FirmwareData->Version == 0) {
        return STATUS_INVALID_PARAMETER;
    }

    if (Stage == 0) {
        RtlCopyMemory(&EfiFirmwareData, FirmwareData, sizeof(BOOT_FIRMWARE_DATA));
        EfiFirmwareParameters = &EfiFirmwareData;

        EfiImageHandle = FirmwareData->ImageHandle;
        EfiST = FirmwareData->SystemTable;
        EfiBS = EfiST->BootServices;
        EfiRT = EfiST->RuntimeServices;
        EfiConOut = EfiST->ConOut;
        EfiConIn = EfiST->ConIn;
        EfiConInEx = NULL;
    } else if (Stage == 1) {
        //
        // Open the extended console input protocol.
        // If successful, tell it to capture partial key events.
        //
        Status = EfiOpenProtocol(
            EfiST->ConsoleInHandle,
            &EfiSimpleTextInputExProtocol,
            (VOID**)&EfiConInEx
        );
        if (NT_SUCCESS(Status)) {
            KeyToggleState = EFI_KEY_STATE_EXPOSED | EFI_TOGGLE_STATE_VALID;
            EfiConInExSetState(EfiConInEx, &KeyToggleState);
        }

        //
        // Disable the watchdog timer.
        //
        EfiSetWatchdogTimer(0, 0, 0, NULL);
    }

    return STATUS_SUCCESS;
}