#define WIN32_NO_STATUS
#include <windows.h>
#undef WIN32_NO_STATUS
#include <ntstatus.h>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <initializer_list>

#define NT_SUCCESS(s) (((NTSTATUS)(s)) >= 0)
#ifndef _IRQL_requires_max_
#define _IRQL_requires_max_(x)
#endif
#define PASSIVE_LEVEL 0
#define PAGED_CODE() ((void)0)
typedef void* PFILE_OBJECT;
typedef void* LPWIN_CERTIFICATE;

struct MINCRYPT_POLICY_INFO
{
    ULONG Size;
    NTSTATUS VerificationStatus;
    ULONG PolicyBits;
    void* ChainInfo;
};

#include "ci-validation-api.inc"

static unsigned checks;
static unsigned live;
static unsigned validationCalls;
static unsigned embeddedCalls;
static unsigned digestCalls;
static unsigned digestFrees;
static NTSTATUS nativeStatus = STATUS_SUCCESS;
static NTSTATUS embeddedStatus = STATUS_SUCCESS;
static NTSTATUS digestStatus = STATUS_SUCCESS;
static UINT32 nativeAlgorithm = CALG_SHA_256;
static UINT32 nativeSize = 32;
static ULONG fakeDigestSize = 32;
static bool omitCertificate;
static bool omitDigest;
static bool partialDigestOnFailure;
static const UCHAR image[] = {1, 2, 3};
static ULONG observedFlags;
static ULONG observedLevel;
static ULONG observedCapacity;
static PFILE_OBJECT observedObject;
static ALG_ID observedEmbeddedAlgorithm;
static ULONG observedSecuritySize;
static void* observedCertificate;

static void
Check(
    bool value,
    const char* message
)
{
    ++checks;

    if (!value)
    {
        std::fprintf(stderr, "FAIL: %s\n", message);
        std::exit(1);
    }
}

static void*
Allocate(
    SIZE_T size
)
{
    void* value = std::malloc(size);
    Check(value != nullptr, "host allocation");
    ++live;
    return value;
}

static void
Release(
    void* value
)
{
    Check(value && live, "single allocation owner");
    --live;
    std::free(value);
}

static void
CiFreePolicyInfo(
    MINCRYPT_POLICY_INFO* policy
)
{
    if (policy->ChainInfo)
    {
        Release(policy->ChainInfo);
    }

    *policy = {};
}

static void
Populate(
    MINCRYPT_POLICY_INFO* policy,
    MINCRYPT_POLICY_INFO* timestamp
)
{
    Check(!policy->ChainInfo && !timestamp->ChainInfo, "old result freed before reuse");
    policy->Size = timestamp->Size = sizeof(*policy);
    policy->ChainInfo = Allocate(1);
    timestamp->ChainInfo = Allocate(1);
    policy->PolicyBits = 0x123;
    timestamp->PolicyBits = 0x456;
    policy->VerificationStatus = STATUS_INVALID_IMAGE_HASH;
}

static NTSTATUS
CiValidateFileObject(
    PFILE_OBJECT object,
    UINT32 flags,
    UINT32 level,
    MINCRYPT_POLICY_INFO* policy,
    MINCRYPT_POLICY_INFO* timestamp,
    LARGE_INTEGER* time,
    UINT8* hash,
    UINT32* size,
    ALG_ID* algorithm
)
{
    ++validationCalls;
    observedObject = object;
    observedFlags = flags;
    observedLevel = level;
    observedCapacity = *size;
    Check(*size == 64, "fresh native hash capacity");
    Check(!time->QuadPart && !*algorithm, "native scalar outputs initialized");

    for (unsigned i = 0; i < 64; ++i)
    {
        Check(!hash[i], "native digest starts empty");
    }

    Populate(policy, timestamp);
    time->QuadPart = 99;
    std::memset(hash, 0x5a, 64);
    *size = nativeSize;
    *algorithm = nativeAlgorithm;
    return nativeStatus;
}

static NTSTATUS
DomitoCalculatePortableExecutableDigestEx(
    const UCHAR* bytes,
    ULONG size,
    UINT32 algorithm,
    PULONG digestSize,
    PVOID* digest,
    LPWIN_CERTIFICATE* certificate,
    PULONG securitySize
)
{
    ++digestCalls;
    Check(bytes == image && size == sizeof(image), "same PE snapshot forwarded");
    Check(
        algorithm == CALG_SHA1 || algorithm == CALG_SHA_256 || algorithm == CALG_SHA_384 || algorithm == CALG_SHA_512,
        "supported digest choice forwarded"
    );

    if (NT_SUCCESS(digestStatus) || partialDigestOnFailure)
    {
        *digest = omitDigest ? nullptr : Allocate(64);

        if (*digest)
        {
            std::memset(*digest, 0xa5, 64);
        }

        *digestSize = fakeDigestSize;
        *certificate = omitCertificate ? nullptr : const_cast<UCHAR*>(image);
        *securitySize = omitCertificate ? 0 : 123;
    }

    return digestStatus;
}

static void
DomitoFreePortableExecutableDigest(
    PVOID digest
)
{
    ++digestFrees;
    Release(digest);
}

static NTSTATUS
CiCheckSignedFile(
    PVOID hash,
    UINT32 size,
    ALG_ID algorithm,
    LPWIN_CERTIFICATE certificate,
    ULONG securitySize,
    MINCRYPT_POLICY_INFO* policy,
    LARGE_INTEGER* time,
    MINCRYPT_POLICY_INFO* timestamp
)
{
    ++embeddedCalls;
    observedEmbeddedAlgorithm = algorithm;
    observedSecuritySize = securitySize;
    observedCertificate = certificate;

    for (UINT32 i = 0; i < size; ++i)
    {
        Check(static_cast<UCHAR*>(hash)[i] == 0xa5, "computed digest forwarded unchanged");
    }

    Populate(policy, timestamp);
    time->QuadPart = 123;
    return embeddedStatus;
}

#ifdef DOMITO_TEST_OTHER_ARCH
#undef _M_AMD64
#undef _AMD64_
#endif
#include "ci-validation-under-test.inc"

static void
CheckEmptyDigest(
    const DOMITO_CI_VALIDATION_RESULT& result
)
{
    Check(!result.HashSize && !result.HashAlgId && !result.SigningTime.QuadPart, "failure digest/time cleared");

    for (auto value : result.Hash)
    {
        Check(!value, "failure hash bytes cleared");
    }
}

static void
Clear(
    DOMITO_CI_VALIDATION_RESULT* result
)
{
    DomitoFreeCiValidationResult(result);
    DOMITO_CI_VALIDATION_RESULT empty = {};
    Check(std::memcmp(result, &empty, sizeof(empty)) == 0, "full result cleared by destructor");
    Check(live == 0, "no leaked policy/digest allocation");
}

int
main()
{
    DOMITO_CI_VALIDATION_RESULT result = {};
    PFILE_OBJECT object = reinterpret_cast<PFILE_OBJECT>(static_cast<ULONG_PTR>(0x1234));
    DomitoFreeCiValidationResult(nullptr);
    Clear(&result);
    Clear(&result);
    Check(
        DomitoValidateFileObjectEx(object, 0, 0, DomitoCiSignatureSelectionNative, nullptr) == STATUS_INVALID_PARAMETER,
        "null native result"
    );
    Check(
        DomitoValidatePortableExecutableEmbeddedSignature(image, sizeof(image), CALG_SHA_256, nullptr) ==
            STATUS_INVALID_PARAMETER,
        "null embedded result"
    );
    Check(!validationCalls && !embeddedCalls && !digestCalls, "null result made no validation calls");

    for (UINT32 flags : {0UL, 0x20UL, 0xa0UL, 0x80000024UL})
    {
        for (UINT8 level : {UINT8(0), UINT8(7), UINT8(255)})
        {
            result.HashSize = 1;
            result.HashAlgId = 0xdead;
            Check(
                DomitoValidateFileObjectEx(object, flags, level, DomitoCiSignatureSelectionNative, &result) ==
                    STATUS_SUCCESS,
                "native status"
            );
            Check(
                observedObject == object && observedFlags == flags && observedLevel == level && observedCapacity == 64,
                "native inputs unchanged"
            );
            Check(
                result.FileObjectValidationFlags == flags && result.RequestedSigningLevel == level,
                "export input controls recorded"
            );
            Check(
                result.HashSize == 32 && result.HashAlgId == CALG_SHA_256 && result.Hash[0] == 0x5a &&
                    result.SigningTime.QuadPart == 99,
                "native result"
            );
            // API return success does not silently reinterpret a policy status.
            Check(result.PolicyInfo.VerificationStatus == STATUS_INVALID_IMAGE_HASH, "policy observation retained");
            Clear(&result);
        }
    }

    const unsigned beforeExperimental = validationCalls;
    NTSTATUS experimental =
        DomitoValidateFileObjectEx(object, 0x20, 9, DomitoCiSignatureSelectionFirstSuccessfulExperimental, &result);
#ifdef DOMITO_TEST_OTHER_ARCH
    Check(
        experimental == STATUS_NOT_SUPPORTED && validationCalls == beforeExperimental,
        "experimental mode excluded on non-x64"
    );
    CheckEmptyDigest(result);
#else
    Check(experimental == STATUS_SUCCESS && validationCalls == beforeExperimental + 1, "experimental native call");
    Check(observedFlags == 0xa0 && observedLevel == 9, "only experimental bit added");
#endif
    Clear(&result);
    const unsigned beforeInvalid = validationCalls;
    Check(
        DomitoValidateFileObjectEx(nullptr, 0x20, 0, DomitoCiSignatureSelectionNative, &result) ==
            STATUS_INVALID_PARAMETER,
        "null file object"
    );
    Check(
        DomitoValidateFileObjectEx(object, 0x20, 0, static_cast<DOMITO_CI_SIGNATURE_SELECTION>(-1), &result) ==
            STATUS_INVALID_PARAMETER,
        "invalid selection"
    );
    Check(validationCalls == beforeInvalid, "invalid requests never call validation");
    Clear(&result);

    for (NTSTATUS failure : {STATUS_INVALID_IMAGE_HASH, STATUS_BUFFER_TOO_SMALL, STATUS_INSUFFICIENT_RESOURCES})
    {
        nativeStatus = failure;
        Check(
            DomitoValidateFileObjectEx(object, 0x20, 0, DomitoCiSignatureSelectionNative, &result) == failure,
            "native failure preserved"
        );
        CheckEmptyDigest(result);
        Check(live == 2 && result.PolicyInfo.PolicyBits == 0x123, "failure policy diagnostics remain owned");
        Clear(&result);
    }

    nativeStatus = STATUS_SUCCESS;
    nativeSize = 65;
    Check(
        DomitoValidateFileObjectEx(object, 0, 0, DomitoCiSignatureSelectionNative, &result) ==
            STATUS_INVALID_BUFFER_SIZE,
        "oversized returned hash rejected"
    );
    CheckEmptyDigest(result);
    Clear(&result);
    nativeSize = 31;
    Check(
        DomitoValidateFileObjectEx(object, 0, 0, DomitoCiSignatureSelectionNative, &result) ==
            STATUS_INVALID_BUFFER_SIZE,
        "algorithm length mismatch rejected"
    );
    Clear(&result);
    nativeSize = 32;
    nativeAlgorithm = 0xdead;
    Check(
        DomitoValidateFileObjectEx(object, 0, 0, DomitoCiSignatureSelectionNative, &result) == STATUS_NOT_SUPPORTED,
        "unknown returned algorithm rejected"
    );
    Clear(&result);

    for (ALG_ID algorithm : {CALG_SHA1, CALG_SHA_256, CALG_SHA_384, CALG_SHA_512})
    {
        nativeAlgorithm = algorithm;
        nativeSize = DomitoCiDigestSize(algorithm);
        Check(
            DomitoValidateFileObjectEx(object, 0, 0, DomitoCiSignatureSelectionNative, &result) == STATUS_SUCCESS,
            "known native digest sizes"
        );

        for (UINT32 i = result.HashSize; i < sizeof(result.Hash); ++i)
        {
            Check(!result.Hash[i], "unused native digest bytes cleared");
        }

        Clear(&result);
        fakeDigestSize = DomitoCiDigestSize(algorithm);
        Check(
            DomitoValidatePortableExecutableEmbeddedSignature(image, sizeof(image), algorithm, &result) ==
                STATUS_SUCCESS,
            "embedded success"
        );
        Check(
            observedEmbeddedAlgorithm == algorithm && observedSecuritySize == 123 && observedCertificate == image,
            "embedded certificate/algorithm forwarded"
        );
        Check(
            result.HashSize == fakeDigestSize && result.HashAlgId == algorithm && result.Hash[0] == 0xa5 &&
                result.SigningTime.QuadPart == 123,
            "embedded result owns copied digest"
        );
        Check(
            !result.FileObjectValidationFlags && !result.RequestedSigningLevel && live == 2,
            "embedded metadata and digest cleanup"
        );
        Clear(&result);
    }

    nativeAlgorithm = CALG_SHA_256;
    nativeSize = 32;
    fakeDigestSize = 32;
    nativeStatus = STATUS_MORE_ENTRIES;
    Check(
        DomitoValidateFileObjectEx(object, 0, 0, DomitoCiSignatureSelectionNative, &result) == STATUS_MORE_ENTRIES,
        "nonnegative native status preserved"
    );
    Clear(&result);
    nativeStatus = STATUS_SUCCESS;
    Check(
        DomitoValidatePortableExecutableEmbeddedSignature(nullptr, 0, CALG_SHA_256, &result) ==
            STATUS_INVALID_PARAMETER,
        "null PE"
    );
    unsigned beforeDigest = digestCalls;
    Check(
        DomitoValidatePortableExecutableEmbeddedSignature(image, sizeof(image), 0xdead, &result) ==
            STATUS_NOT_SUPPORTED,
        "unsupported algorithm"
    );
    Check(digestCalls == beforeDigest, "unsupported algorithm does not hash");

    for (NTSTATUS failure : {STATUS_NOT_SUPPORTED, STATUS_INVALID_IMAGE_FORMAT, STATUS_INSUFFICIENT_RESOURCES})
    {
        digestStatus = failure;
        unsigned beforeEmbedded = embeddedCalls;
        Check(
            DomitoValidatePortableExecutableEmbeddedSignature(image, sizeof(image), CALG_SHA_256, &result) == failure,
            "digest failure propagation"
        );
        Check(embeddedCalls == beforeEmbedded, "digest failure never calls CI");
        CheckEmptyDigest(result);
        Clear(&result);
    }

    partialDigestOnFailure = true;
    unsigned beforeFree = digestFrees;
    Check(
        DomitoValidatePortableExecutableEmbeddedSignature(image, sizeof(image), CALG_SHA_256, &result) ==
            STATUS_INSUFFICIENT_RESOURCES,
        "partial digest failure"
    );
    Check(digestFrees == beforeFree + 1, "partial digest allocation freed");
    Clear(&result);
    partialDigestOnFailure = false;
    digestStatus = STATUS_SUCCESS;
    omitCertificate = true;
    unsigned beforeEmbedded = embeddedCalls;
    Check(
        DomitoValidatePortableExecutableEmbeddedSignature(image, sizeof(image), CALG_SHA_256, &result) ==
            STATUS_INVALID_IMAGE_HASH,
        "unsigned PE rejected"
    );
    Check(embeddedCalls == beforeEmbedded, "unsigned PE does not call CI");
    Clear(&result);
    omitCertificate = false;
    omitDigest = true;
    Check(
        DomitoValidatePortableExecutableEmbeddedSignature(image, sizeof(image), CALG_SHA_256, &result) ==
            STATUS_INVALID_BUFFER_SIZE,
        "missing digest rejected"
    );
    Clear(&result);
    omitDigest = false;
    fakeDigestSize = 31;
    Check(
        DomitoValidatePortableExecutableEmbeddedSignature(image, sizeof(image), CALG_SHA_256, &result) ==
            STATUS_INVALID_BUFFER_SIZE,
        "digest helper length mismatch"
    );
    Clear(&result);
    fakeDigestSize = 32;
    embeddedStatus = STATUS_INVALID_IMAGE_HASH;
    Check(
        DomitoValidatePortableExecutableEmbeddedSignature(image, sizeof(image), CALG_SHA_256, &result) ==
            STATUS_INVALID_IMAGE_HASH,
        "CI embedded failure propagated without fallback"
    );
    CheckEmptyDigest(result);
    Check(live == 2 && result.TimeStampPolicyInfo.PolicyBits == 0x456, "embedded failure diagnostics owned");
    Clear(&result);
    embeddedStatus = STATUS_SUCCESS;

    for (unsigned i = 0; i < 16; ++i)
    {
        Check(
            DomitoValidateFileObjectEx(object, 0x20, 0, DomitoCiSignatureSelectionNative, &result) == STATUS_SUCCESS,
            "repeated native"
        );
        Check(
            DomitoValidatePortableExecutableEmbeddedSignature(image, sizeof(image), CALG_SHA_256, &result) ==
                STATUS_SUCCESS,
            "reuse native result as embedded result"
        );
        Check(
            !result.FileObjectValidationFlags && !result.RequestedSigningLevel && live == 2,
            "reuse resets metadata without leaking"
        );
    }

    // Invalid calls must release a previously valid result, too.
    Check(
        DomitoValidateFileObjectEx(nullptr, 0, 0, DomitoCiSignatureSelectionNative, &result) ==
            STATUS_INVALID_PARAMETER,
        "invalid reuse cleans old result"
    );
    Clear(&result);
    std::printf(
        "PASS: %u CI validation checks (%s architecture branch), no live allocations\n",
        checks,
#ifdef DOMITO_TEST_OTHER_ARCH
        "non-x64"
#else
        "x64"
#endif
    );
}