/*   ___         _       ___     _                _ _
 *  / __|___  __| |___  |_ _|_ _| |_ ___ __ _ _ _(_) |_ _  _
 * | (__/ _ \/ _` / -_)  | || ' \  _/ -_) _` | '_| |  _| || |
 *  \___\___/\__,_\___| |___|_||_\__\___\__, |_| |_|\__|\_, |
 *                                      |___/           |__/
 */

#include "Domito.Internal.h"
#include "Domito.MinCrypt.h"

/*  _   _ _   _ _ _ _          ___             _   _
 * | | | | |_(_) (_) |_ _  _  | __|  _ _ _  __| |_(_)___ _ _  ___
 * | |_| |  _| | | |  _| || | | _| || | ' \/ _|  _| / _ \ ' \(_-<
 *  \___/ \__|_|_|_|\__|\_, | |_| \_,_|_||_\__|\__|_\___/_||_/__/
 *                      |__/
 */

_IRQL_requires_max_(DISPATCH_LEVEL)
UINT32
DomitoGetPortableExecutableDigestKind(
    _In_ PUCHAR pPeBytes,
    _In_ ULONG PeSize,
    _In_ PIMAGE_DATA_DIRECTORY pImgDataDirectory
)
{
    if (!pPeBytes || !pImgDataDirectory)
    {
        return CALG_SHA1;
    }

    const IMAGE_DATA_DIRECTORY security = pImgDataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY];
    const ULONG securityRva = security.VirtualAddress;
    const ULONG certificateHeaderSize = FIELD_OFFSET(WIN_CERTIFICATE, bCertificate);

    if (!securityRva || securityRva > PeSize || security.Size > PeSize - securityRva ||
        security.Size < certificateHeaderSize)
    {
        return CALG_SHA1;
    }

    ULONG certificateSize;
    memcpy(&certificateSize, pPeBytes + securityRva, sizeof(certificateSize));

    if (certificateSize < certificateHeaderSize || certificateSize > security.Size)
    {
        return CALG_SHA1;
    }

    const PVOID pBase = pPeBytes + securityRva;
    const ULONG searchLen = certificateSize;
    PUCHAR pMatch = NULL;

    if (NT_SUCCESS(DomitoMemorySearchPattern(
            (const PUCHAR) "\x30\x21\x30\x09\x06\x05\x2b\x0e\x03\x02\x1a\x05\x00\x04\x14",
            0x00,
            15,
            pBase,
            searchLen,
            (PVOID*)&pMatch
        )))
    {
        return CALG_SHA1;
    }

    if (NT_SUCCESS(DomitoMemorySearchPattern(
            (const PUCHAR) "\x30\xcc\x30\x0d\x06\x09\x60\x86\x48\x01\x65\x03\x04\x02\xcc\x05\x00\x04\xcc",
            0xcc,
            19,
            pBase,
            searchLen,
            (PVOID*)&pMatch
        )))
    {
        if (pMatch == NULL)
        {
            return CALG_SHA1;
        }

        if (pMatch[1] == 0x31 && pMatch[14] == 0x01 && pMatch[18] == 0x20)
        {
            return CALG_SHA256;
        }
        else if (pMatch[1] == 0x41 && pMatch[14] == 0x02 && pMatch[18] == 0x30)
        {
            return CALG_SHA384;
        }
        else if (pMatch[1] == 0x51 && pMatch[14] == 0x03 && pMatch[18] == 0x40)
        {
            return CALG_SHA512;
        }
    }

    return CALG_SHA1;
}

// Bounded DER navigation for the outer Authenticode content only. This does not
// validate certificates, signerInfos, nested signatures, or cryptographic trust.
typedef struct _DOMITO_DER_VIEW
{
    const UCHAR* Bytes;
    ULONG Length;
} DOMITO_DER_VIEW;

static BOOLEAN
DomitoDerTake(
    DOMITO_DER_VIEW* Input,
    UCHAR Tag,
    DOMITO_DER_VIEW* Value
)
{
    if (Input->Length < 2 || Input->Bytes[0] != Tag)
    {
        return FALSE;
    }

    ULONG length = Input->Bytes[1];
    ULONG header = 2;

    if (length & 0x80)
    {
        const ULONG count = length & 0x7F;

        if (!count || count > sizeof(ULONG) || count > Input->Length - header || Input->Bytes[header] == 0)
        {
            return FALSE; // No indefinite or nonminimal DER lengths.
        }

        length = 0;

        for (ULONG i = 0; i < count; ++i)
        {
            length = (length << 8) | Input->Bytes[header++];
        }

        if (length < 0x80)
        {
            return FALSE;
        }
    }

    if (length > Input->Length - header)
    {
        return FALSE;
    }

    Value->Bytes = Input->Bytes + header;
    Value->Length = length;
    Input->Bytes += header + length;
    Input->Length -= header + length;
    return TRUE;
}

static BOOLEAN
DomitoDerOid(
    DOMITO_DER_VIEW* Input,
    const UCHAR* Oid,
    ULONG Size
)
{
    DOMITO_DER_VIEW value;
    return DomitoDerTake(Input, 0x06, &value) && value.Length == Size && memcmp(value.Bytes, Oid, Size) == 0;
}

_IRQL_requires_max_(DISPATCH_LEVEL)
NTSTATUS
DomitoGetPortableExecutableDigestKindEx(
    _In_reads_bytes_(PeSize) const UCHAR* PeBytes,
    _In_ ULONG PeSize,
    _In_reads_(IMAGE_DIRECTORY_ENTRY_SECURITY + 1) const IMAGE_DATA_DIRECTORY* Directories,
    _Out_ PUINT32 DigestCalg
)
{
    if (!DigestCalg)
    {
        return STATUS_INVALID_PARAMETER;
    }

    *DigestCalg = 0;

    if (!PeBytes || !Directories)
    {
        return STATUS_INVALID_PARAMETER;
    }

    const IMAGE_DATA_DIRECTORY security = Directories[IMAGE_DIRECTORY_ENTRY_SECURITY];

    if (!security.VirtualAddress && !security.Size)
    {
        return STATUS_NOT_FOUND;
    }

    if (!security.VirtualAddress || (security.VirtualAddress & 7) || security.VirtualAddress > PeSize ||
        security.Size > PeSize - security.VirtualAddress || security.Size < 8)
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    ULONG length;
    USHORT revision;
    USHORT type;
    const UCHAR* certificate = PeBytes + security.VirtualAddress;
    memcpy(&length, certificate, 4);
    memcpy(&revision, certificate + 4, 2);
    memcpy(&type, certificate + 6, 2);

    if (length <= 8 || length > security.Size)
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    if (revision != 0x0200 || type != WIN_CERT_TYPE_PKCS_SIGNED_DATA)
    {
        return STATUS_NOT_SUPPORTED;
    }

    const UCHAR signedDataOid[] = {0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x07, 0x02};
    const UCHAR indirectOid[] = {0x2B, 0x06, 0x01, 0x04, 0x01, 0x82, 0x37, 0x02, 0x01, 0x04};
    const UCHAR sha1Oid[] = {0x2B, 0x0E, 0x03, 0x02, 0x1A};
    const UCHAR sha2Prefix[] = {0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02};
    DOMITO_DER_VIEW payload = {certificate + 8, length - 8};
    DOMITO_DER_VIEW content;
    DOMITO_DER_VIEW wrapped;
    DOMITO_DER_VIEW signedData;
    DOMITO_DER_VIEW ignored;
    DOMITO_DER_VIEW contentInfo;
    DOMITO_DER_VIEW indirect;
    DOMITO_DER_VIEW digest;
    DOMITO_DER_VIEW algorithm;
    DOMITO_DER_VIEW oid;

    if (!DomitoDerTake(&payload, 0x30, &content) || !DomitoDerOid(&content, signedDataOid, sizeof(signedDataOid)) ||
        !DomitoDerTake(&content, 0xA0, &wrapped) || content.Length || !DomitoDerTake(&wrapped, 0x30, &signedData) ||
        wrapped.Length || !DomitoDerTake(&signedData, 0x02, &ignored) || !DomitoDerTake(&signedData, 0x31, &ignored) ||
        !DomitoDerTake(&signedData, 0x30, &contentInfo) ||
        !DomitoDerOid(&contentInfo, indirectOid, sizeof(indirectOid)) || !DomitoDerTake(&contentInfo, 0xA0, &wrapped) ||
        contentInfo.Length || !DomitoDerTake(&wrapped, 0x30, &indirect) || wrapped.Length ||
        !DomitoDerTake(&indirect, 0x30, &ignored) || !DomitoDerTake(&indirect, 0x30, &digest) || indirect.Length ||
        !DomitoDerTake(&digest, 0x30, &algorithm) || !DomitoDerTake(&algorithm, 0x06, &oid))
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    // Windows-signed entries can include up to seven zero alignment bytes in
    // dwLength itself. Accept only that bounded, aligned padding after the DER.
    if (payload.Length > 7 || (payload.Length && (length & 7)))
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    for (ULONG i = 0; i < payload.Length; ++i)
    {
        if (payload.Bytes[i])
        {
            return STATUS_INVALID_IMAGE_FORMAT;
        }
    }

    if (algorithm.Length && (!DomitoDerTake(&algorithm, 0x05, &ignored) || ignored.Length || algorithm.Length))
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    UINT32 calg = 0;
    ULONG expectedLength = 0;

    if (oid.Length == sizeof(sha1Oid) && !memcmp(oid.Bytes, sha1Oid, sizeof(sha1Oid)))
    {
        calg = CALG_SHA1;
        expectedLength = 20;
    }
    else if (oid.Length == sizeof(sha2Prefix) + 1 && !memcmp(oid.Bytes, sha2Prefix, sizeof(sha2Prefix)))
    {
        switch (oid.Bytes[sizeof(sha2Prefix)])
        {
        case 1:
            calg = CALG_SHA256;
            expectedLength = 32;
            break;
        case 2:
            calg = CALG_SHA384;
            expectedLength = 48;
            break;
        case 3:
            calg = CALG_SHA512;
            expectedLength = 64;
            break;
        }
    }

    if (!calg)
    {
        return STATUS_NOT_SUPPORTED;
    }

    if (!DomitoDerTake(&digest, 0x04, &ignored) || digest.Length || ignored.Length != expectedLength)
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    *DigestCalg = calg;
    return STATUS_SUCCESS;
}

// Conservative section-layout boundary: contiguous raw sections (any header
// order), optional overlay before an EOF certificate table. Gaps, overlaps, and
// certificates in the middle are rejected rather than assigned an ambiguous hash.
static NTSTATUS
DomitoValidateAuthenticodeLayout(
    const UCHAR* Bytes,
    ULONG Size
)
{
    if (!Bytes || Size < sizeof(IMAGE_DOS_HEADER))
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    IMAGE_DOS_HEADER dos;
    memcpy(&dos, Bytes, sizeof(dos));

    if (dos.e_magic != IMAGE_DOS_SIGNATURE || dos.e_lfanew < 0)
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    const ULONG nt = (ULONG)dos.e_lfanew;
    const ULONG prefix = sizeof(ULONG) + sizeof(IMAGE_FILE_HEADER);

    if (nt > Size || prefix > Size - nt)
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    ULONG signature;
    IMAGE_FILE_HEADER file;
    memcpy(&signature, Bytes + nt, sizeof(signature));
    memcpy(&file, Bytes + nt + sizeof(signature), sizeof(file));

    if (signature != IMAGE_NT_SIGNATURE || !file.NumberOfSections || file.NumberOfSections > 96)
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    const ULONG optional = nt + prefix;

    if (file.SizeOfOptionalHeader < 2 || file.SizeOfOptionalHeader > Size - optional)
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    USHORT magic;
    memcpy(&magic, Bytes + optional, sizeof(magic));
    ULONG directory;

    if (magic == IMAGE_NT_OPTIONAL_HDR32_MAGIC)
    {
        directory = FIELD_OFFSET(IMAGE_OPTIONAL_HEADER32, DataDirectory);
    }
    else if (magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC)
    {
        directory = FIELD_OFFSET(IMAGE_OPTIONAL_HEADER64, DataDirectory);
    }
    else
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    if (file.SizeOfOptionalHeader < directory)
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    ULONG count;
    ULONG headers;
    memcpy(&count, Bytes + optional + directory - 4, sizeof(count));
    memcpy(&headers, Bytes + optional + FIELD_OFFSET(IMAGE_OPTIONAL_HEADER32, SizeOfHeaders), sizeof(headers));

    if (count > (file.SizeOfOptionalHeader - directory) / sizeof(IMAGE_DATA_DIRECTORY))
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    if (count <= IMAGE_DIRECTORY_ENTRY_SECURITY)
    {
        return STATUS_NOT_SUPPORTED; // SIP exclusion behavior differs for absent slots.
    }

    const ULONG table = optional + file.SizeOfOptionalHeader;
    const ULONG tableSize = file.NumberOfSections * sizeof(IMAGE_SECTION_HEADER);

    if (tableSize > Size - table || headers < table + tableSize || headers > Size)
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    struct
    {
        ULONG Offset;
        ULONG Length;
    } sections[96];

    ULONG used = 0;

    for (ULONG i = 0; i < file.NumberOfSections; ++i)
    {
        IMAGE_SECTION_HEADER section;
        memcpy(&section, Bytes + table + i * sizeof(section), sizeof(section));

        if (!section.SizeOfRawData)
        {
            continue;
        }

        if (section.PointerToRawData < headers || section.PointerToRawData > Size ||
            section.SizeOfRawData > Size - section.PointerToRawData)
        {
            return STATUS_INVALID_IMAGE_FORMAT;
        }

        ULONG at = used;
        while (at && sections[at - 1].Offset > section.PointerToRawData)
        {
            sections[at] = sections[at - 1];
            --at;
        }

        sections[at].Offset = section.PointerToRawData;
        sections[at].Length = section.SizeOfRawData;
        ++used;
    }

    ULONG end = headers;

    for (ULONG i = 0; i < used; ++i)
    {
        if (sections[i].Offset != end)
        {
            return STATUS_NOT_SUPPORTED;
        }

        end += sections[i].Length;
    }

    if (count > IMAGE_DIRECTORY_ENTRY_SECURITY)
    {
        IMAGE_DATA_DIRECTORY security;
        memcpy(
            &security,
            Bytes + optional + directory + IMAGE_DIRECTORY_ENTRY_SECURITY * sizeof(security),
            sizeof(security)
        );

        if (security.VirtualAddress || security.Size)
        {
            if (!security.VirtualAddress || !security.Size || (security.VirtualAddress & 7) ||
                security.VirtualAddress < end || security.VirtualAddress > Size ||
                security.Size != Size - security.VirtualAddress)
            {
                return STATUS_NOT_SUPPORTED;
            }

            // Validate every aligned certificate entry, not just the first.
            ULONG offset = security.VirtualAddress;
            while (offset < Size)
            {
                if (Size - offset < 8)
                {
                    return STATUS_INVALID_IMAGE_FORMAT;
                }

                ULONG length;
                USHORT revision;
                USHORT type;
                memcpy(&length, Bytes + offset, 4);
                memcpy(&revision, Bytes + offset + 4, 2);
                memcpy(&type, Bytes + offset + 6, 2);

                if (length <= 8 || length > Size - offset)
                {
                    return STATUS_INVALID_IMAGE_FORMAT;
                }

                if (revision != 0x0200 || type != WIN_CERT_TYPE_PKCS_SIGNED_DATA)
                {
                    return STATUS_NOT_SUPPORTED;
                }

                const ULONG padding = (8 - (length & 7)) & 7;

                if (padding > Size - offset - length)
                {
                    return STATUS_INVALID_IMAGE_FORMAT;
                }

                offset += length + padding;
            }
        }
    }

    return STATUS_SUCCESS;
}

_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
static NTSTATUS
DomitoCalculatePortableExecutableDigestInternal(
    _In_ UINT32 ExplicitCalg,
    _In_ PUCHAR pPeBytes,
    _In_ ULONG PeSize,
    _Out_ PUINT32 pDigestCalgOut,
    _Out_ PULONG pDigestSizeOut,
    _Out_ PVOID* pDigestOut,
    _Outptr_result_maybenull_ LPWIN_CERTIFICATE* pCertOut,
    _Out_ PULONG pSizeOfSecurityDirectory
)
{
    if (!pDigestCalgOut || !pDigestSizeOut || !pDigestOut || !pCertOut || !pSizeOfSecurityDirectory)
    {
        return STATUS_INVALID_PARAMETER;
    }

    *pDigestCalgOut = 0;
    *pDigestSizeOut = 0;
    *pDigestOut = nullptr;
    *pCertOut = nullptr;
    *pSizeOfSecurityDirectory = 0;

    if (!pPeBytes)
    {
        return STATUS_INVALID_PARAMETER;
    }

    if (PeSize < sizeof(IMAGE_DOS_HEADER))
    {
        return STATUS_INVALID_IMAGE_NOT_MZ;
    }

    // Copy fields rather than assuming the caller's buffer is aligned.
    IMAGE_DOS_HEADER dos = {};
    memcpy(&dos, pPeBytes, sizeof(dos));

    if (dos.e_magic != IMAGE_DOS_SIGNATURE)
    {
        return STATUS_INVALID_IMAGE_NOT_MZ;
    }

    if (dos.e_lfanew < 0)
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    const ULONG ntOffset = (ULONG)dos.e_lfanew;
    const ULONG ntPrefixSize = sizeof(ULONG) + sizeof(IMAGE_FILE_HEADER);

    if (ntOffset > PeSize || ntPrefixSize > PeSize - ntOffset)
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    ULONG signature;
    IMAGE_FILE_HEADER fileHeader;
    memcpy(&signature, pPeBytes + ntOffset, sizeof(signature));
    memcpy(&fileHeader, pPeBytes + ntOffset + sizeof(signature), sizeof(fileHeader));

    if (signature != IMAGE_NT_SIGNATURE)
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    const ULONG optionalOffset = ntOffset + ntPrefixSize;
    const ULONG optionalSize = fileHeader.SizeOfOptionalHeader;

    if (optionalSize < sizeof(USHORT) || optionalSize > PeSize - optionalOffset)
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    USHORT magic;
    memcpy(&magic, pPeBytes + optionalOffset, sizeof(magic));
    ULONG directoryOffset;
    ULONG checksumOffset;
    ULONG directoryCountOffset;

    if (magic == IMAGE_NT_OPTIONAL_HDR32_MAGIC)
    {
        directoryOffset = FIELD_OFFSET(IMAGE_OPTIONAL_HEADER32, DataDirectory);
        checksumOffset = FIELD_OFFSET(IMAGE_OPTIONAL_HEADER32, CheckSum);
        directoryCountOffset = FIELD_OFFSET(IMAGE_OPTIONAL_HEADER32, NumberOfRvaAndSizes);
    }
    else if (magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC)
    {
        directoryOffset = FIELD_OFFSET(IMAGE_OPTIONAL_HEADER64, DataDirectory);
        checksumOffset = FIELD_OFFSET(IMAGE_OPTIONAL_HEADER64, CheckSum);
        directoryCountOffset = FIELD_OFFSET(IMAGE_OPTIONAL_HEADER64, NumberOfRvaAndSizes);
    }
    else
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    if (optionalSize < directoryOffset)
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    ULONG directoryCount;
    memcpy(&directoryCount, pPeBytes + optionalOffset + directoryCountOffset, sizeof(directoryCount));

    if (directoryCount > (optionalSize - directoryOffset) / sizeof(IMAGE_DATA_DIRECTORY))
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    checksumOffset += optionalOffset;
    const BOOLEAN hasSecurityEntry = directoryCount > IMAGE_DIRECTORY_ENTRY_SECURITY;
    IMAGE_DATA_DIRECTORY directories[IMAGE_DIRECTORY_ENTRY_SECURITY + 1] = {};
    ULONG securityEntryOffset = PeSize;
    ULONG dataOffset = PeSize;

    if (hasSecurityEntry)
    {
        securityEntryOffset =
            optionalOffset + directoryOffset + IMAGE_DIRECTORY_ENTRY_SECURITY * sizeof(IMAGE_DATA_DIRECTORY);
        memcpy(
            &directories[IMAGE_DIRECTORY_ENTRY_SECURITY], pPeBytes + securityEntryOffset, sizeof(IMAGE_DATA_DIRECTORY)
        );
        dataOffset = securityEntryOffset + sizeof(IMAGE_DATA_DIRECTORY);
    }

    const IMAGE_DATA_DIRECTORY security = directories[IMAGE_DIRECTORY_ENTRY_SECURITY];
    const BOOLEAN hasEmbeddedSig = security.VirtualAddress != 0;

    if (hasEmbeddedSig != (security.Size != 0))
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    if (hasEmbeddedSig)
    {
        // The security directory contains a FILE OFFSET, not an image RVA.
        const ULONG certificateHeaderSize = FIELD_OFFSET(WIN_CERTIFICATE, bCertificate);

        if (security.VirtualAddress < dataOffset || security.VirtualAddress > PeSize ||
            security.Size > PeSize - security.VirtualAddress || security.Size < certificateHeaderSize)
        {
            return STATUS_INVALID_IMAGE_FORMAT;
        }

        ULONG certificateSize;
        memcpy(&certificateSize, pPeBytes + security.VirtualAddress, sizeof(certificateSize));

        if (certificateSize < certificateHeaderSize || certificateSize > security.Size)
        {
            return STATUS_INVALID_IMAGE_FORMAT;
        }
    }

    PUCHAR pHash = nullptr;
    BCRYPT_ALG_HANDLE hbAlg = nullptr;
    BCRYPT_HASH_HANDLE hbHash = nullptr;
    ULONG hashLength = 0;
    ULONG resultLength = 0;
    // Intentional compatibility branch; the explicit API never calls the heuristic.
#pragma warning(push)
#pragma warning(disable : 4996)
    const UINT32 peDigestKind =
        ExplicitCalg ? ExplicitCalg : DomitoGetPortableExecutableDigestKind(pPeBytes, PeSize, directories);
#pragma warning(pop)
    NTSTATUS status =
        BCryptOpenAlgorithmProvider(&hbAlg, DOMITO_CALG_TO_BCRYPT_ALGORITHM(peDigestKind), MS_PRIMITIVE_PROVIDER, 0);

    if (!NT_SUCCESS(status))
    {
        goto cleanup;
    }

    status = BCryptGetProperty(hbAlg, BCRYPT_HASH_LENGTH, (PUCHAR)&hashLength, sizeof(hashLength), &resultLength, 0);

    if (!NT_SUCCESS(status))
    {
        goto cleanup;
    }

#pragma warning(disable : 4996)
    pHash = (PUCHAR)ExAllocatePoolWithTag(NonPagedPoolNx, hashLength, DOMITO_POOL_TAG);
#pragma warning(default : 4996)
    if (!pHash)
    {
        status = STATUS_INSUFFICIENT_RESOURCES;
        goto cleanup;
    }

    status = BCryptCreateHash(hbAlg, &hbHash, nullptr, 0, nullptr, 0, 0);

    if (!NT_SUCCESS(status))
    {
        goto cleanup;
    }

    // Keep the existing linear hashing convention, but hash validated ranges
    // directly: no fixed-size scratch buffer and no unchecked memcpy lengths.
    // General section-ordered Authenticode hashing is a separate change.
    status = BCryptHashData(hbHash, pPeBytes, checksumOffset, 0);

    if (!NT_SUCCESS(status))
    {
        goto cleanup;
    }

    status = BCryptHashData(
        hbHash, pPeBytes + checksumOffset + sizeof(ULONG), securityEntryOffset - checksumOffset - sizeof(ULONG), 0
    );

    if (!NT_SUCCESS(status))
    {
        goto cleanup;
    }

    if (hasSecurityEntry)
    {
        status = BCryptHashData(
            hbHash, pPeBytes + dataOffset, (hasEmbeddedSig ? security.VirtualAddress : PeSize) - dataOffset, 0
        );

        if (!NT_SUCCESS(status))
        {
            goto cleanup;
        }
    }

    status = BCryptFinishHash(hbHash, pHash, hashLength, 0);

    if (!NT_SUCCESS(status))
    {
        goto cleanup;
    }

    *pDigestCalgOut = peDigestKind;
    *pDigestSizeOut = hashLength;
    *pDigestOut = pHash;
    pHash = nullptr; // Transfer ownership only after all work has succeeded.

    if (hasEmbeddedSig)
    {
        *pCertOut = (LPWIN_CERTIFICATE)(pPeBytes + security.VirtualAddress);
        *pSizeOfSecurityDirectory = security.Size;
    }

cleanup:

    if (hbHash)
    {
        BCryptDestroyHash(hbHash);
    }

    if (pHash)
    {
        ExFreePoolWithTag(pHash, DOMITO_POOL_TAG);
    }

    if (hbAlg)
    {
        BCryptCloseAlgorithmProvider(hbAlg, 0);
    }

    return status;
}

_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
NTSTATUS
DomitoCalculatePortableExecutableDigest(
    PUCHAR PeBytes,
    ULONG PeSize,
    PUINT32 DigestCalg,
    PULONG DigestSize,
    PVOID* Digest,
    LPWIN_CERTIFICATE* Certificate,
    PULONG SecuritySize
)
{
    return DomitoCalculatePortableExecutableDigestInternal(
        0, PeBytes, PeSize, DigestCalg, DigestSize, Digest, Certificate, SecuritySize
    );
}

_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
NTSTATUS
DomitoCalculatePortableExecutableDigestEx(
    const UCHAR* PeBytes,
    ULONG PeSize,
    UINT32 HashAlgId,
    PULONG DigestSize,
    PVOID* Digest,
    LPWIN_CERTIFICATE* Certificate,
    PULONG SecuritySize
)
{
    if (!DigestSize || !Digest || !Certificate || !SecuritySize)
    {
        return STATUS_INVALID_PARAMETER;
    }

    *DigestSize = 0;
    *Digest = nullptr;
    *Certificate = nullptr;
    *SecuritySize = 0;

    if (!PeBytes)
    {
        return STATUS_INVALID_PARAMETER;
    }

    if (HashAlgId != CALG_SHA1 && HashAlgId != CALG_SHA256 && HashAlgId != CALG_SHA384 && HashAlgId != CALG_SHA512)
    {
        return STATUS_NOT_SUPPORTED;
    }

    NTSTATUS status = DomitoValidateAuthenticodeLayout(PeBytes, PeSize);

    if (!NT_SUCCESS(status))
    {
        return status;
    }

    UINT32 actualCalg = 0;
    // Contiguous sorted sections make the validated section sequence identical
    // to linear byte order; overlay precedes an EOF certificate table.
    return DomitoCalculatePortableExecutableDigestInternal(
        HashAlgId, (PUCHAR)PeBytes, PeSize, &actualCalg, DigestSize, Digest, Certificate, SecuritySize
    );
}

_IRQL_requires_max_(DISPATCH_LEVEL)
void
DomitoFreePortableExecutableDigest(
    _In_ PVOID pDigestOut
)
{
    if (!pDigestOut)
    {
        return;
    }

    ExFreePoolWithTag(pDigestOut, DOMITO_POOL_TAG);
}

_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
#pragma code_seg("PAGED")
static NTSTATUS
DomitoValidateFileLegacyModeInternal(
    _In_ HANDLE FileHandle,
    _In_ PVOID Hash,
    _In_ UINT32 HashSize,
    _In_ ALG_ID HashAlgId,
    _In_ const IMAGE_DATA_DIRECTORY* SecurityDirectory,
    _Inout_ MINCRYPT_POLICY_INFO* PolicyInfo,
    _Out_ LARGE_INTEGER* SigningTime,
    _Inout_ MINCRYPT_POLICY_INFO* TimeStampPolicyInfo,
    _In_ UINT32 AcceptedCatalogPolicyBits,
    _In_ BOOLEAN RequireTrustedCatalog,
    _Out_ PULONG ValidationSource
)
{
    PAGED_CODE();

    if (!ValidationSource)
    {
        return STATUS_INVALID_PARAMETER;
    }

    *ValidationSource = 0;

    if (!Hash || !SecurityDirectory || !PolicyInfo || !SigningTime || !TimeStampPolicyInfo ||
        PolicyInfo == TimeStampPolicyInfo)
    {
        return STATUS_INVALID_PARAMETER;
    }

    NTSTATUS status = STATUS_SUCCESS;
    PVOID certDirectory = nullptr;
    KAPC_STATE systemContext = {};

    do
    {
        SigningTime->QuadPart = 0;

        CiFreePolicyInfo(PolicyInfo);
        CiFreePolicyInfo(TimeStampPolicyInfo);

        if (HashAlgId != CALG_SHA1 && HashAlgId != CALG_SHA256)
        {
            status = STATUS_NOT_SUPPORTED;
            break;
        }

        if (HashSize != (UINT32)(HashAlgId == CALG_SHA1 ? MINCRYPT_SHA1_LENGTH : MINCRYPT_SHA256_LENGTH))
        {
            status = STATUS_INVALID_PARAMETER;
            break;
        }

        if ((SecurityDirectory->Size != 0) != (SecurityDirectory->VirtualAddress != 0))
        {
            status = STATUS_INVALID_IMAGE_FORMAT;
            break;
        }

        if (SecurityDirectory->Size != 0u && SecurityDirectory->VirtualAddress != 0u)
        {
            if ((SecurityDirectory->VirtualAddress & 7) || SecurityDirectory->Size < 12)
            {
                status = STATUS_INVALID_IMAGE_FORMAT;
                break;
            }

            // Bound the allocation/read against the file before trusting its table.
            PFILE_OBJECT fileObject = nullptr;
            status = ObReferenceObjectByHandle(
                FileHandle, FILE_READ_DATA, *IoFileObjectType, KernelMode, (PVOID*)&fileObject, nullptr
            );

            if (!NT_SUCCESS(status))
            {
                break;
            }

            LARGE_INTEGER fileSize = {};
            status = FsRtlGetFileSize(fileObject, &fileSize);
            ObDereferenceObject(fileObject);

            if (!NT_SUCCESS(status))
            {
                break;
            }

            if (fileSize.QuadPart < 0 || (ULONGLONG)SecurityDirectory->VirtualAddress > (ULONGLONG)fileSize.QuadPart ||
                SecurityDirectory->Size > (ULONGLONG)fileSize.QuadPart - SecurityDirectory->VirtualAddress)
            {
                status = STATUS_INVALID_IMAGE_FORMAT;
                break;
            }
#pragma warning(disable : 4996)
            certDirectory = ExAllocatePoolWithTag(PagedPool, SecurityDirectory->Size, DOMITO_POOL_TAG);
#pragma warning(default : 4996)
            if (certDirectory == NULL)
            {
                status = STATUS_INSUFFICIENT_RESOURCES;
                break;
            }

            LARGE_INTEGER offset = {};
            offset.QuadPart = SecurityDirectory->VirtualAddress;
            status = DomitoReadFileAt(FileHandle, certDirectory, SecurityDirectory->Size, &offset);

            if (!NT_SUCCESS(status))
            {
                break;
            }

            KeStackAttachProcess(PsInitialSystemProcess, &systemContext);
            {
                status = CiCheckSignedFile(
                    Hash,
                    HashSize,
                    HashAlgId,
                    certDirectory,
                    SecurityDirectory->Size,
                    PolicyInfo,
                    SigningTime,
                    TimeStampPolicyInfo
                );
            }

            KeUnstackDetachProcess(&systemContext);

            if (NT_SUCCESS(status))
            {
                *ValidationSource = DOMITO_CI_SOURCE_EMBEDDED;
                break;
            }

            if (status != STATUS_INVALID_IMAGE_HASH)
            {
                break;
            }
        }

        CiFreePolicyInfo(PolicyInfo);
        CiFreePolicyInfo(TimeStampPolicyInfo);
        SigningTime->QuadPart = 0;
        KeStackAttachProcess(PsInitialSystemProcess, &systemContext);
        {
            status = CiVerifyHashInCatalog(
                Hash,
                HashSize,
                HashAlgId,
                FALSE,
                RequireTrustedCatalog,
                AcceptedCatalogPolicyBits,
                PolicyInfo,
                NULL,
                SigningTime,
                TimeStampPolicyInfo
            );

            if (status == STATUS_INVALID_IMAGE_HASH)
            {
                CiFreePolicyInfo(PolicyInfo);
                CiFreePolicyInfo(TimeStampPolicyInfo);
                SigningTime->QuadPart = 0;
                status = CiVerifyHashInCatalog(
                    Hash,
                    HashSize,
                    HashAlgId,
                    TRUE,
                    RequireTrustedCatalog,
                    AcceptedCatalogPolicyBits,
                    PolicyInfo,
                    NULL,
                    SigningTime,
                    TimeStampPolicyInfo
                );
            }
        }

        KeUnstackDetachProcess(&systemContext);

        if (NT_SUCCESS(status))
        {
            *ValidationSource = DOMITO_CI_SOURCE_CATALOG;
        }

    } while (FALSE);

    if (certDirectory)
    {
        ExFreePoolWithTag(certDirectory, DOMITO_POOL_TAG);
    }

    return status;
}

_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
NTSTATUS
DomitoValidateFileLegacyMode(
    HANDLE FileHandle,
    PVOID Hash,
    UINT32 HashSize,
    ALG_ID HashAlgId,
    const IMAGE_DATA_DIRECTORY* SecurityDirectory,
    MINCRYPT_POLICY_INFO* PolicyInfo,
    LARGE_INTEGER* SigningTime,
    MINCRYPT_POLICY_INFO* TimeStampPolicyInfo
)
{
    ULONG source;
    return DomitoValidateFileLegacyModeInternal(
        FileHandle,
        Hash,
        HashSize,
        HashAlgId,
        SecurityDirectory,
        PolicyInfo,
        SigningTime,
        TimeStampPolicyInfo,
        DOMITO_CI_LEGACY_CATALOG_POLICY_BITS,
        FALSE,
        &source
    );
}

_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
NTSTATUS
DomitoValidateFileLegacyModeEx(
    HANDLE FileHandle,
    PVOID Hash,
    UINT32 HashSize,
    ALG_ID HashAlgId,
    const IMAGE_DATA_DIRECTORY* SecurityDirectory,
    MINCRYPT_POLICY_INFO* PolicyInfo,
    LARGE_INTEGER* SigningTime,
    MINCRYPT_POLICY_INFO* TimeStampPolicyInfo,
    UINT32 AcceptedCatalogPolicyBits,
    BOOLEAN RequireTrustedCatalog,
    PULONG ValidationSource
)
{
    return DomitoValidateFileLegacyModeInternal(
        FileHandle,
        Hash,
        HashSize,
        HashAlgId,
        SecurityDirectory,
        PolicyInfo,
        SigningTime,
        TimeStampPolicyInfo,
        AcceptedCatalogPolicyBits,
        RequireTrustedCatalog,
        ValidationSource
    );
}

#pragma code_seg()