#pragma once

#include <ci.h>


/********************************************************************************
 * Memory management, misc.                                                     *
 ********************************************************************************/

 //
 // Custom allocator for function that allocate pool memory.
 // 
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;


/********************************************************************************
 * Cryptography                                                                 *
 ********************************************************************************/

 //
 // 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";
    }
}


/********************************************************************************
 * Library functions                                                            *
 ********************************************************************************/

 // 
 // Finds the base address of a driver module.
 // 
_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
NTSTATUS
DomitoFindModuleBaseAddress(
    _In_ STRING ModuleName,
    _In_ PFN_DOMITO_ALLOCATE_ROUTINE Allocator,
    _Inout_opt_ PVOID * ModuleBase
);

//
// Finds the address of an exported function by name.
// 
_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
NTSTATUS
DomitoFindExportedFunctionAddress(
    _In_ PVOID ModuleBase,
    _In_ 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
);

//
// Extracts the CALG_ID from a signed PE that was used to
// calculate the message digest when it was signed
//
_IRQL_requires_max_(DISPATCH_LEVEL)
EXTERN_C
UINT32
DomitoGetPortableExecutableDigestKind(
    _In_ PUCHAR pPeBytes,
    _In_ PIMAGE_DATA_DIRECTORY pImgDataDirectory
);

//
// Reads from the beginning of a file until the end or the buffer size is reached.
// 
_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
);

//
// Extracts Authenticode signing information and calculates the file digest of a PE file.
// 
_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
NTSTATUS
DomitoCalculatePortableExecutableDigest(
    _In_ PFN_DOMITO_ALLOCATE_ROUTINE Allocator,
    _In_ PUCHAR pPeBytes,
    _In_ ULONG PeSize,
    _Out_ PUINT32 pDigestCalgOut,
    _Out_ PULONG pDigestSizeOut,
    _Out_ PVOID* pDigestOut,
    _Out_ LPWIN_CERTIFICATE* pCertOut,
    _Out_ PULONG pSizeOfSecurityDirectory
);

//
// Gets the name of the main image of the process identified by PID.
// 
_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
NTSTATUS
DomitoGetProcessImageName(
    _In_ PFN_DOMITO_ALLOCATE_ROUTINE Allocator,
    _In_ ULONG ProcessId,
    _Inout_ PUNICODE_STRING* ProcessImageName
);