policy_runtime.cpp

#include "bcd.hpp"
#include "bootdefs.hpp"
#include <string.h>
#include <stdlib.h>

// ============================================================================
// Advanced Boot Decision and Policy Runtime Definitions
// ============================================================================

enum BootDecisionState {
    BOOT_CONTINUE = 0,
    BOOT_RECOVERY = 1,
    BOOT_MENU     = 2,
    BOOT_FAIL     = 3,
    BOOT_FVE_UNLOCK_REQUIRED = 4,
    BOOT_POLICY_BLOCKED      = 5
};

struct BootOrchestratorState {
    PVOID               BootEntry;
    PVOID               ActivePolicyState;
    BootDecisionState   CurrentDecision;
    NTSTATUS            LastStatus;
    BOOLEAN             FveUnlocked;
    BOOLEAN             SecureBootActive;
    uint32_t            AppType;
    uint64_t            AppVersion;
};

// Reconstructed 16-byte telemetry correlation ID structure
struct BOOT_TELEMETRY_CORRELATION {
    GUID     ScenarioGuid;
    uint64_t Timestamp;
};

// ============================================================================
// BootPolicyOrchestrator Class
// ============================================================================

class BootPolicyOrchestrator {
private:
    BootOrchestratorState m_State;
    BOOT_TELEMETRY_CORRELATION m_Telemetry;

public:
    BootPolicyOrchestrator(PVOID bootEntry) {
        memset(&m_State, 0, sizeof(BootOrchestratorState));
        memset(&m_Telemetry, 0, sizeof(BOOT_TELEMETRY_CORRELATION));
        m_State.BootEntry = bootEntry;
        m_State.CurrentDecision = BOOT_CONTINUE;
        m_State.LastStatus = STATUS_SUCCESS;
    }

    ~BootPolicyOrchestrator() {
        if (m_State.ActivePolicyState) {
            SIPolicyReleaseMutableState(&m_State.ActivePolicyState);
            SIPolicyReleaseState(&m_State.ActivePolicyState);
        }
    }

    /**
     * @brief [GAP 1 & 2] Resolves Boot Decision states and parses BCD options.
     */
    NTSTATUS InitializeRuntime() {
        if (!m_State.BootEntry) {
            m_State.CurrentDecision = BOOT_FAIL;
            m_State.LastStatus = STATUS_INVALID_PARAMETER;
            return STATUS_INVALID_PARAMETER;
        }

        // Clone current Code Integrity policy active state
        m_State.LastStatus = SIPolicyCloneActiveState(&m_State.ActivePolicyState);
        if (m_State.LastStatus < 0) {
            m_State.CurrentDecision = BOOT_POLICY_BLOCKED;
            return m_State.LastStatus;
        }

        // Query firmware Secure Boot state
        BOOLEAN capable = FALSE;
        m_State.LastStatus = BlSIPolicyIsSecureBootCapable(&m_State.SecureBootActive, &capable);
        if (m_State.LastStatus < 0) {
            m_State.SecureBootActive = FALSE;
        }

        // Extract Telemetry Correlation Guid from BCD
        PBOOT_APPLICATION_ENTRY appEntry = (PBOOT_APPLICATION_ENTRY)m_State.BootEntry;
        memcpy(&m_Telemetry.ScenarioGuid, &appEntry->ApplicationId, sizeof(GUID));
        m_Telemetry.Timestamp = __rdtsc();

        return STATUS_SUCCESS;
    }

    /**
     * @brief [GAP 3] Evaluates BitLocker/FVE volume unlocking dependencies.
     */
    NTSTATUS EvaluateVolumeDecryption() {
        if (m_State.CurrentDecision != BOOT_CONTINUE) {
            return m_State.LastStatus;
        }

        PBOOT_APPLICATION_ENTRY appEntry = (PBOOT_APPLICATION_ENTRY)m_State.BootEntry;
        BOOLEAN isFveActive = ((appEntry->Flags & 4) != 0);

        if (isFveActive) {
            PVOID fveDevice = nullptr;
            m_State.LastStatus = BlGetBootOptionDevice(appEntry->BcdOptions, Bcd::OSLoaderDevice_OSDevice, &fveDevice, 0);
            
            if (m_State.LastStatus >= 0 && fveDevice) {
                // Call BitLocker partition unlock routine
                m_State.LastStatus = FvebpUnlockSinglePartition(1, fveDevice, nullptr, 2, 0, 0);
                if (m_State.LastStatus < 0) {
                    m_State.CurrentDecision = BOOT_FVE_UNLOCK_REQUIRED;
                    FvebpLogError(12, m_State.LastStatus, nullptr, nullptr, 0);
                } else {
                    m_State.FveUnlocked = TRUE;
                }
            } else {
                m_State.CurrentDecision = BOOT_FAIL;
                m_State.LastStatus = STATUS_UNSUCCESSFUL;
            }

            if (fveDevice) {
                BlMmFreeHeap(fveDevice);
            }
        } else {
            m_State.FveUnlocked = TRUE;
        }

        return m_State.LastStatus;
    }

    /**
     * @brief [GAP 4] Merges active OS security policies with firmware UEFI PK/KEK keys.
     */
    NTSTATUS MergeFirmwareSecurityPolicies() {
        if (m_State.CurrentDecision != BOOT_CONTINUE) {
            return m_State.LastStatus;
        }

        if (m_State.SecureBootActive) {
            // Retrieve Secure Boot active database configuration
            SbUseFirmwareRevocationLists();

            if (SbpActivePolicy) {
                PVOID mergedPayload = nullptr;
                uint32_t mergedSize = 0;
                
                // Overlay OS policy rules on top of baseline Firmware db/dbx keys
                m_State.LastStatus = SbMergeActivePolicyWithPolicy(m_State.AppVersion, 0, &mergedPayload, &mergedSize);
                if (m_State.LastStatus >= 0) {
                    m_State.LastStatus = BlSecureBootPackagePolicyForKernel(mergedPayload, mergedSize);
                }

                if (mergedPayload) {
                    BlMmFreeHeap(mergedPayload);
                }
            }
        }

        return m_State.LastStatus;
    }

    /**
     * @brief [GAP 5] Classifies target boot application types.
     */
    NTSTATUS ClassifyBootApplication(uint32_t appTypeIndex) {
        if (m_State.CurrentDecision != BOOT_CONTINUE) {
            return m_State.LastStatus;
        }

        m_State.LastStatus = BlImgGetApplicationType(appTypeIndex, &m_State.AppType, &m_State.AppVersion);
        if (m_State.LastStatus < 0) {
            // Default to OS Loader (winload.efi) fallback classification
            m_State.AppType = 1; 
            m_State.AppVersion = 0x100000000LL;
            m_State.LastStatus = STATUS_SUCCESS;
        }

        return m_State.LastStatus;
    }

    /**
     * @brief [GAP 6] Validates strict cryptographic binding between policy and executable image.
     */
    NTSTATUS ValidateImagePolicyBinding(PVOID imageBase, uint32_t imageSize, const wchar_t* imagePath) {
        if (m_State.CurrentDecision != BOOT_CONTINUE) {
            return m_State.LastStatus;
        }

        // Verify PE Authenticode signature
        m_State.LastStatus = BlImgVerifySignedPeImageFileContents(imagePath, imageBase, imageSize);
        if (m_State.LastStatus < 0) {
            m_State.CurrentDecision = BOOT_POLICY_BLOCKED;
            return m_State.LastStatus;
        }

        // Strict binding: extract version resources from image
        int size = 0;
        PVOID versionResource = BlResourceFindVersionFromImage((uint64_t)imageBase, imageSize, 0, &size);
        if (versionResource && size > 0) {
            uint16_t fileName[128] = {0};
            uint64_t version = 0;

            m_State.LastStatus = SIPolicyGetOriginalFilenameAndVersionFromVersionResource(
                versionResource,
                size,
                (int)(uintptr_t)fileName,
                (int)(uintptr_t)&version,
                0, 0, 0
            );

            if (m_State.LastStatus >= 0) {
                // Ensure parsed policy version matches target image version bounds
                if (m_State.AppVersion != 0 && version < m_State.AppVersion) {
                    m_State.CurrentDecision = BOOT_POLICY_BLOCKED;
                    m_State.LastStatus = -1058471932; // STATUS_SIPOLICY_POLICY_BLOCKED
                }
            }
        }

        return m_State.LastStatus;
    }

    /**
     * @brief [GAP 9] Traverses policy inheritance hierarchies and BCD dependency graphs.
     */
    NTSTATUS ResolvePolicyHierarchy() {
        if (m_State.CurrentDecision != BOOT_CONTINUE) {
            return m_State.LastStatus;
        }

        PBOOT_APPLICATION_ENTRY appEntry = (PBOOT_APPLICATION_ENTRY)m_State.BootEntry;
        BOOLEAN inheritParent = FALSE;

        // Traverse inherited boot menu and safety policies
        if (BlGetBootOptionBoolean(appEntry->BcdOptions, Bcd::OSLoaderBoolean_AssociatedResume, &inheritParent) >= 0 && inheritParent) {
            PVOID parentDevice = nullptr;
            m_State.LastStatus = BlGetBootOptionDevice(appEntry->BcdOptions, Bcd::LibraryDevice_ApplicationDevice, &parentDevice, 0);
            if (m_State.LastStatus >= 0 && parentDevice) {
                // Incorporate parent volume settings
                BlSIPolicyCheckSystemDevices(parentDevice);
                BlMmFreeHeap(parentDevice);
            }
        }

        return STATUS_SUCCESS;
    }

    /**
     * @brief [GAP 7 & 8] Orchestrates rollback routing and logs ETW scenarios.
     */
    void HandleRollbackAndTelemetry(uint32_t transitionReason) {
        // Log telemetry event and trace timeline
        BmpLogEtwApplicationScenarioTime(BOOT_APPLICATION_EXECUTION_TIME, &m_Telemetry.ScenarioGuid, m_Telemetry.Timestamp);
        
        struct BSD_LOG_BLOCK {
            GUID     Scenario;
            uint32_t Decision;
            uint32_t Reserved;
        } bsdLog;

        bsdLog.Scenario = m_Telemetry.ScenarioGuid;
        bsdLog.Decision = (uint32_t)m_State.CurrentDecision;
        bsdLog.Reserved = 0;

        BlBsdLogEntry(1, 18, &bsdLog);

        if (m_State.CurrentDecision != BOOT_CONTINUE) {
            // Trigger automatic fallback recovery sequences
            PBOOT_APPLICATION_ENTRY appEntry = (PBOOT_APPLICATION_ENTRY)m_State.BootEntry;
            if (BlBootOptionExists(appEntry->BcdOptions, Bcd::LibraryInteger_AutoRecoveryEnabled)) {
                m_State.CurrentDecision = BOOT_RECOVERY;
                BmLaunchRecoverySequence(m_State.BootEntry, transitionReason);
            } else {
                m_State.CurrentDecision = BOOT_MENU;
                BmDisplayDumpError(m_State.BootEntry, transitionReason);
            }
        }
    }

    /**
     * @brief [GAP 10] Finalizes and handsoff activated policies to the Operating System kernel.
     */
    NTSTATUS FinalizeKernelTrustActivation() {
        if (m_State.CurrentDecision != BOOT_CONTINUE) {
            return m_State.LastStatus;
        }

        // Commit finalized security state
        SIPolicyCommitState(m_State.ActivePolicyState);

        // Actual kernel trust activation layer
        ImgpProcessFlightRootTrustPolicy();

        m_State.CurrentDecision = BOOT_CONTINUE;
        return STATUS_SUCCESS;
    }

    BootDecisionState GetDecision() const { return m_State.CurrentDecision; }
    NTSTATUS GetLastStatus() const { return m_State.LastStatus; }
};