#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, live, validationCalls, embeddedCalls, digestCalls, digestFrees;
static NTSTATUS nativeStatus = STATUS_SUCCESS, embeddedStatus = STATUS_SUCCESS, digestStatus = STATUS_SUCCESS;
static UINT32 nativeAlgorithm = CALG_SHA_256, nativeSize = 32;
static ULONG fakeDigestSize = 32;
static bool omitCertificate, omitDigest, partialDigestOnFailure;
static const UCHAR image[] = {1,2,3};
static ULONG observedFlags, observedLevel, 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
    );
}