#pragma once

#include <Domito.MinCrypt.h>

/*   ___
 *  / __|___ _ __  _ __  ___ _ _
 * | (__/ _ \ '  \| '  \/ _ \ ' \
 *  \___\___/_|_|_|_|_|_\___/_||_|
 *
 */

//
// Library initialization tasks. Call once in your DriverEntry
//
_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
NTSTATUS
DomitoInit();

_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
void
DomitoShutdown();

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

//
// Allocator function the library uses.
//
typedef _IRQL_requires_same_ _Function_class_(EVT_DOMITO_ALLOCATE_ROUTINE)
__drv_allocatesMem(
    mem
)
PVOID
NTAPI
EVT_DOMITO_ALLOCATE_ROUTINE(_In_ SIZE_T ByteSize);
typedef EVT_DOMITO_ALLOCATE_ROUTINE* PFN_DOMITO_ALLOCATE_ROUTINE;

//
// Freeing function the library uses.
//
typedef _IRQL_requires_same_ _Function_class_(EVT_DOMITO_FREE_ROUTINE) void NTAPI
EVT_DOMITO_FREE_ROUTINE(_In_ __drv_freesMem(mem) PVOID Memory);
typedef EVT_DOMITO_FREE_ROUTINE* PFN_DOMITO_FREE_ROUTINE;

//
// Get the original set of Domito memory functions.
//
EXTERN_C
void
DomitoGetOriginalMemoryFunctions(
    _Out_opt_ PFN_DOMITO_ALLOCATE_ROUTINE* Allocator,
    _Out_opt_ PFN_DOMITO_FREE_ROUTINE* Free
);

//
// Get the current set of Domito memory functions.
//
EXTERN_C
void
DomitoGetMemoryFunctions(
    _Out_opt_ PFN_DOMITO_ALLOCATE_ROUTINE* Allocator,
    _Out_opt_ PFN_DOMITO_FREE_ROUTINE* Free
);

//
// Replace Domito's memory allocation functions with a compatible pair during
// quiescent initialization. No in-flight calls or outstanding callback-backed
// allocations may span a pair change. Replacement is not synchronized.
EXTERN_C
void
DomitoSetMemoryFunctions(
    _In_opt_ PFN_DOMITO_ALLOCATE_ROUTINE Allocator,
    _In_opt_ PFN_DOMITO_FREE_ROUTINE Free
);

/*   ___               _                           _        
 *  / __|_ _ _  _ _ __| |_ ___  __ _ _ _ __ _ _ __| |_ _  _ 
 * | (__| '_| || | '_ \  _/ _ \/ _` | '_/ _` | '_ \ ' \ || |
 *  \___|_|  \_, | .__/\__\___/\__, |_| \__,_| .__/_||_\_, |
 *           |__/|_|           |___/         |_|       |__/ 
 */

//
// This structure encapsulates a signature used in verifying executable files.
//
#if !defined(WIN_CERTIFICATE)
typedef struct _WIN_CERTIFICATE
{
    DWORD dwLength;
    WORD wRevision;
    WORD wCertificateType;
    BYTE bCertificate[ANYSIZE_ARRAY];
} WIN_CERTIFICATE, *LPWIN_CERTIFICATE;
#endif

//
// UM definitions of WinCrypt.h
//

#if !defined(WIN_CERT_TYPE_X509)
#define WIN_CERT_TYPE_X509 (0x0001) // The bCertificate member contains an X.509 certificate.
#endif
#if !defined(WIN_CERT_TYPE_PKCS_SIGNED_DATA)
#define WIN_CERT_TYPE_PKCS_SIGNED_DATA (0x0002) // The bCertificate member contains a PKCS SignedData structure.
#endif
#if !defined(WIN_CERT_TYPE_PKCS1_SIGN)
#define WIN_CERT_TYPE_PKCS1_SIGN (0x0009) // The bCertificate member contains PKCS1_MODULE_SIGN fields.
#endif

#if !defined(CALG_SHA1)
#define CALG_SHA1 0x8004u
#endif
#if !defined(CALG_SHA256)
#define CALG_SHA256 0x800cu
#endif
#if !defined(CALG_SHA384)
#define CALG_SHA384 0x800du
#endif
#if !defined(CALG_SHA512)
#define CALG_SHA512 0x800eu
#endif

//
// Converts a WinCrypt CALG_ID to a BCRYPT_ALGORITHM identifier.
//
PCWSTR

FORCEINLINE
DOMITO_CALG_TO_BCRYPT_ALGORITHM(
    _In_ UINT32 Calg
)
{
    switch (Calg)
    {
    case CALG_SHA1:
        return BCRYPT_SHA1_ALGORITHM;
    case CALG_SHA256:
        return BCRYPT_SHA256_ALGORITHM;
    case CALG_SHA384:
        return BCRYPT_SHA384_ALGORITHM;
    case CALG_SHA512:
        return BCRYPT_SHA512_ALGORITHM;
    default:
        return L"Unknown";
    }
}

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

//
// Heuristically detects the CALG_ID in the first WIN_CERTIFICATE using byte patterns.
// Returns CALG_SHA1 when absent, malformed, or unrecognized; this is not ASN.1
// validation or signature verification. The directory array must contain the
// IMAGE_DIRECTORY_ENTRY_SECURITY entry.
//
_IRQL_requires_max_(DISPATCH_LEVEL)
EXTERN_C
__declspec(deprecated(
    "Use DomitoGetPortableExecutableDigestKindEx; heuristic detection silently falls back to SHA-1. Check the returned NTSTATUS."
)) UINT32
DomitoGetPortableExecutableDigestKind(
    _In_ PUCHAR pPeBytes,
    _In_ ULONG PeSize,
    _In_ PIMAGE_DATA_DIRECTORY pImgDataDirectory
);

// Reads the first PKCS#7 entry's SpcIndirectDataContent DigestInfo algorithm
// using bounded DER navigation. Does not validate the signature or select among
// multiple/nested signatures. Absent table => STATUS_NOT_FOUND; unsupported OID
// or certificate type => STATUS_NOT_SUPPORTED; malformed DER => format error.
// Unlike the compatibility helper above, it never silently falls back to SHA-1.
_IRQL_requires_max_(DISPATCH_LEVEL)
EXTERN_C
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
);

// Algorithm-explicit Authenticode digest for the supported layout subset:
// bounded PE32/PE32+ headers, contiguous nonoverlapping raw sections in file
// offset order, optional overlay before a complete EOF certificate table.
// SHA-1/256/384/512 supported. Gaps/overlaps/nonterminal certificates are rejected.
// No signature trust decision is made. Outputs are cleared on failure; the
// certificate borrows stable input memory, and only the digest must be freed.
// See docs/authenticode-digests.md for policy and Windows SIP comparisons.
_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
NTSTATUS
DomitoCalculatePortableExecutableDigestEx(
    _In_reads_bytes_(PeSize) const UCHAR* PeBytes,
    _In_ ULONG PeSize,
    _In_ UINT32 HashAlgId,
    _Out_ PULONG DigestSize,
    _Outptr_result_bytebuffer_(*DigestSize) PVOID* Digest,
    _Outptr_result_maybenull_ LPWIN_CERTIFICATE* Certificate,
    _Out_ PULONG SecuritySize
);

//
// Calculates a PE digest using the existing linear hashing convention.
// Validates accessed header/certificate ranges and clears outputs on failure
// when all output pointers are valid. The input buffer must remain stable.
// The returned certificate pointer borrows the input buffer; free only the digest
// with DomitoFreePortableExecutableDigest. Algorithm detection is heuristic;
// general section-ordered Authenticode hashing is not implemented here.
//
_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
__declspec(deprecated(
    "Use DomitoCalculatePortableExecutableDigestEx with an explicit algorithm; review its supported PE layout and failure statuses."
)) NTSTATUS
DomitoCalculatePortableExecutableDigest(
    _In_ PUCHAR pPeBytes,
    _In_ ULONG PeSize,
    _Out_ PUINT32 pDigestCalgOut,
    _Out_ PULONG pDigestSizeOut,
    _Out_ PVOID* pDigestOut,
    _Outptr_result_maybenull_ LPWIN_CERTIFICATE* pCertOut,
    _Out_ PULONG pSizeOfSecurityDirectory
);

//
// Frees the memory allocated by DomitoCalculatePortableExecutableDigest.
//
_IRQL_requires_max_(DISPATCH_LEVEL)
EXTERN_C
void
DomitoFreePortableExecutableDigest(
    _In_ PVOID pDigestOut
);

#if (NTDDI_VERSION >= NTDDI_WIN10)
// Native forwards caller flags unchanged. The experimental x64 mode additionally
// sets 0x80; it stops after the first accepted candidate in examined CI builds,
// and also changes other policy/fallback paths. Neither mode selects a publisher
// or enumerates signatures. No root-acceptance flag is added automatically.
typedef enum _DOMITO_CI_SIGNATURE_SELECTION
{
    DomitoCiSignatureSelectionNative = 0,
    DomitoCiSignatureSelectionFirstSuccessfulExperimental = 1
} DOMITO_CI_SIGNATURE_SELECTION;

// Zero-initialize before first use. Both policies own their CI chain allocations,
// including possible diagnostic allocations on failure. Do not copy as an owner.
// Hash is a file digest, not a certificate thumbprint. Hash/algorithm/signing time
// are cleared on failure; policies may retain diagnostics. Check return NTSTATUS.
typedef struct _DOMITO_CI_VALIDATION_RESULT
{
    MINCRYPT_POLICY_INFO PolicyInfo;
    MINCRYPT_POLICY_INFO TimeStampPolicyInfo;
    LARGE_INTEGER SigningTime;
    UINT8 Hash[MINCRYPT_MAX_HASH_LENGTH];
    UINT32 HashSize;
    ALG_ID HashAlgId;
    UINT32 FileObjectValidationFlags; // Flags passed to CI export; zero for embedded helper.
    UINT8 RequestedSigningLevel;      // File-object helper only.
} DOMITO_CI_VALIDATION_RESULT, *PDOMITO_CI_VALIDATION_RESULT;

// Safe for NULL or an empty/initialized result; resets all fields after cleanup.
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
void
DomitoFreeCiValidationResult(
    _Inout_opt_ DOMITO_CI_VALIDATION_RESULT* Result
);

// Referenced FILE_OBJECT at PASSIVE_LEVEL; keep file bytes stable. Previous result
// policies are released before reuse, including invalid-argument failures.
// Supplies a fresh 64-byte hash buffer on every call. FirstSuccessfulExperimental
// is opt-in and x64-only; other architectures return STATUS_NOT_SUPPORTED for it.
// Only one selected signer plus its timestamp policy is returned. The caller
// must enforce its expected certificate and intended trust policy separately.
_Success_(return >= 0)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
NTSTATUS
DomitoValidateFileObjectEx(
    _In_ PFILE_OBJECT FileObject,
    _In_ UINT32 ValidationFlags,
    _In_ UINT8 RequestedSigningLevel,
    _In_ DOMITO_CI_SIGNATURE_SELECTION Selection,
    _Inout_ DOMITO_CI_VALIDATION_RESULT* Result
);

// Hash and verify the first WIN_CERTIFICATE in a stable PE buffer. Uses the
// explicit digest helper's supported contiguous-section/EOF-certificate subset.
// SHA-1/256/384/512 are explicit; no algorithm retry or catalog fallback occurs.
// CI may select a nested signature under its algorithm rules: this is not a
// universal force-primary API. Matching SHA-256 selects the company primary in
// the tested dual-signed BthPS3 fixture. Caller still checks the returned signer.
_Success_(return >= 0)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
NTSTATUS
DomitoValidatePortableExecutableEmbeddedSignature(
    _In_reads_bytes_(PeSize) const UCHAR* PeBytes,
    _In_ ULONG PeSize,
    _In_ ALG_ID HashAlgId,
    _Inout_ DOMITO_CI_VALIDATION_RESULT* Result
);
#endif // NTDDI_VERSION >= NTDDI_WIN10

#define DOMITO_CI_SOURCE_EMBEDDED 1UL
#define DOMITO_CI_SOURCE_CATALOG 2UL

// Policy-explicit legacy helper. SHA-1/256 only; hash size must match algorithm.
// Policy objects must be zero-initialized or own valid CI results; both are freed
// before each attempt. Caller owns both final policy outputs, including failure.
// Embedded validation is tried first; only INVALID_IMAGE_HASH permits catalog
// fallback/reload. AcceptedCatalogPolicyBits uses ANY-overlap catalog filtering;
// RequireTrustedCatalog gates the observed marker, not publisher identity.
// Source is 0 on failure, EMBEDDED/CATALOG on success. Stable file bytes and a
// handle valid in the caller's current process context are required during I/O.
_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
NTSTATUS
DomitoValidateFileLegacyModeEx(
    _In_ HANDLE FileHandle,
    _In_reads_bytes_(HashSize) 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
);

//
// Compatibility wrapper around DomitoValidateFileLegacyModeEx, with catalog
// mask DOMITO_CI_LEGACY_CATALOG_POLICY_BITS and RequireTrustedCatalog = FALSE.
// The same initialization, ownership, hash, handle, and stable-file rules apply.
//
_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
__declspec(deprecated(
    "Use DomitoValidateFileLegacyModeEx with explicit catalog policy and validation-source output; do not change trust policy implicitly."
)) NTSTATUS
DomitoValidateFileLegacyMode(
    _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
);

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

//
// Finds the base address of a driver module. Clears the optional output first.
//
_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
NTSTATUS
DomitoFindModuleBaseAddress(
    _In_ PANSI_STRING ModuleName,
    _Inout_opt_ PVOID* ModuleBase
);

//
// Finds an export by name (case-insensitive) in a valid loader-mapped module.
// ModuleBase must remain mapped; arbitrary file buffers are not supported.
// Export table spans are checked against SizeOfImage. Forwarded exports return
// STATUS_NOT_SUPPORTED; the optional output is cleared before lookup.
//
_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
NTSTATUS
DomitoFindExportedFunctionAddress(
    _In_ PVOID ModuleBase,
    _In_ PANSI_STRING FunctionName,
    _Inout_opt_ PVOID* FunctionAddress
);

//
// Scans a provided buffer for a memory pattern.
//
_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(DISPATCH_LEVEL)
EXTERN_C
NTSTATUS
DomitoMemorySearchPattern(
    _In_ PCUCHAR pcPattern,
    _In_ UCHAR uWildcard,
    _In_ SIZE_T puLen,
    _In_ PVOID pcBase,
    _In_ SIZE_T puSize,
    _Outptr_result_maybenull_ PVOID* ppMatch
);

//
// Reads exactly BufferSize bytes at a nonnegative absolute file offset.
// Same completion, buffer, handle, APC, and short-read contract as DomitoReadFile.
//
_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
NTSTATUS
DomitoReadFileAt(
    _In_ HANDLE FileHandle,
    _Out_writes_bytes_(BufferSize) PVOID Buffer,
    _In_ ULONG BufferSize,
    _In_ const LARGE_INTEGER* ByteOffset
);

//
// Reads exactly BufferSize bytes from file offset zero. Short reads return
// STATUS_END_OF_FILE; Buffer may then contain a partial read. Zero length succeeds.
// Waits for pending I/O on a private event before returning. Requires PASSIVE_LEVEL
// with special kernel APCs enabled, a readable handle valid in the current process
// (or a kernel handle), and a buffer that remains valid for the entire call.
// For synchronous handles, the explicit read also updates the file position.
//
_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
NTSTATUS
DomitoReadFile(
    _In_ HANDLE FileHandle,
    _Out_ PVOID Buffer,
    _In_ ULONG BufferSize
);

//
// Gets the main image name after resolving the ULONG PID once to a process handle.
// Requires DomitoInit. Output is NULL on failure; successful result is one Domito
// allocation containing UNICODE_STRING and its buffer. Use DomitoFreeProcessImageName.
// Do not change memory callbacks with outstanding allocations. PID reuse can race
// the initial open; a later path reopen is not proof of original image identity.
//
_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
__declspec(deprecated(
    "Use DomitoGetProcessImageNameEx with a caller-referenced PEPROCESS to avoid a fresh PID lookup; free the result with DomitoFreeProcessImageName."
)) NTSTATUS
DomitoGetProcessImageName(
    _In_ ULONG ProcessId,
    _Outptr_result_maybenull_ PUNICODE_STRING* ProcessImageName
);

// Additive consumer helpers, included here for existing Domito.h clients.
#include "Domito.Consumer.h"