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

#pragma once

//
// SDK/WDK
// 
#include <ntifs.h>
#include <ntintsafe.h>
#include <ntimage.h>
#include <bcrypt.h>

//
// Public
// 
#include "Domito.h"


/*  _  _ _   ___  _ _       _        
 * | \| | |_|   \| | |  ___| |_ __   
 * | .` |  _| |) | | | / -_)  _/ _|_ 
 * |_|\_|\__|___/|_|_| \___|\__\__(_)
 *                                   
 */

 // Structure representing a loaded module
typedef struct _SYSTEM_MODULE_INFORMATION_ENTRY
{
    PVOID Unknown1;
    PVOID Unknown2;
    PVOID Base;
    ULONG Size;
    ULONG Flags;
    USHORT Index;
    USHORT NameLength;
    USHORT LoadCount;
    USHORT PathLength;
    CHAR ImageName[256];
} SYSTEM_MODULE_INFORMATION_ENTRY, * PSYSTEM_MODULE_INFORMATION_ENTRY;

// Structure representing the loaded module information
typedef struct _SYSTEM_MODULE_INFORMATION
{
    ULONG Count;
    SYSTEM_MODULE_INFORMATION_ENTRY Module[ANYSIZE_ARRAY];
} SYSTEM_MODULE_INFORMATION, * PSYSTEM_MODULE_INFORMATION;

// Function prototype for ZwQuerySystemInformation
EXTERN_C NTSYSAPI NTSTATUS NTAPI ZwQuerySystemInformation(
    _In_ ULONG SystemInformationClass,
    _Inout_ PVOID SystemInformation,
    _In_ ULONG SystemInformationLength,
    _Out_opt_ PULONG ReturnLength
);

typedef struct _LDR_DATA_TABLE_ENTRY
{
    LIST_ENTRY64 InLoadOrderLinks;
    PVOID ExceptionTable;
    ULONG ExceptionTableSize;
    PVOID GpValue;
    PVOID NonPagedDebugInfo;
    PVOID ImageBase;
    PVOID EntryPoint;
    ULONG SizeOfImage;
    UNICODE_STRING FullImageName;
    UNICODE_STRING BaseImageName;
    ULONG Flags;
    USHORT LoadCount;
    USHORT TlsIndex;
    LIST_ENTRY64 HashLinks;
    PVOID SectionPointer;
    ULONG CheckSum;
    ULONG TimeDateStamp;
    PVOID LoadedImports;
    PVOID EntryPointActivationContext;
    PVOID PatchInformation;
} LDR_DATA_TABLE_ENTRY, * PLDR_DATA_TABLE_ENTRY;

typedef PVOID(NTAPI* t_RtlImageDirectoryEntryToData)(
    IN PVOID Base,
    IN BOOLEAN MappedAsImage,
    IN USHORT DirectoryEntry,
    OUT PULONG Size
    );

typedef NTSTATUS(NTAPI* t_ZwQueryInformationProcess) (
    __in HANDLE ProcessHandle,
    __in PROCESSINFOCLASS ProcessInformationClass,
    __out_bcount(ProcessInformationLength) PVOID ProcessInformation,
    __in ULONG ProcessInformationLength,
    __out_opt PULONG ReturnLength
    );

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

typedef struct
{
    t_RtlImageDirectoryEntryToData RtlImageDirectoryEntryToData;

    t_ZwQueryInformationProcess ZwQueryInformationProcess;
    
} DOMITO_COMMON;

extern DOMITO_COMMON G_Common;


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

//
// Default pool tag
// 
#define DOMITO_POOL_TAG     'imoD'

//
// Function pointers for malloc/free variants
// 
typedef struct
{
    PFN_DOMITO_ALLOCATE_ROUTINE Allocate;

    PFN_DOMITO_FREE_ROUTINE Free;
} DOMITO_MEMORY;

//
// Global instance, individual field can be adjusted by the caller
// 
extern DOMITO_MEMORY G_Memory;


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

typedef decltype(&CiFreePolicyInfo) t_CiFreePolicyInfo;
typedef decltype(&CiCheckSignedFile) t_CiCheckSignedFile;
typedef decltype(&CiVerifyHashInCatalog) t_CiVerifyHashInCatalog;
typedef decltype(&CiGetCertPublisherName) t_CiGetCertPublisherName;
typedef decltype(&CiSetTrustedOriginClaimId) t_CiSetTrustedOriginClaimId;
typedef decltype(&CiValidateFileObject) t_CiValidateFileObject;

//
// Function pointers to CI.dll exports
// 
typedef struct
{
    t_CiFreePolicyInfo CiFreePolicyInfo;

    t_CiCheckSignedFile CiCheckSignedFile;

    t_CiVerifyHashInCatalog CiVerifyHashInCatalog;

    t_CiGetCertPublisherName CiGetCertPublisherName;

    t_CiSetTrustedOriginClaimId CiSetTrustedOriginClaimId;

    t_CiValidateFileObject CiValidateFileObject;

} DOMITO_CODE_INTEGRITY;

//
// Global instance
// 
extern DOMITO_CODE_INTEGRITY G_CI;