#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 custom set
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";
    }
}


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

//
// 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
);

//
// 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_ 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
);

//
// Verifies if the Authenticode signature of a give PE file matches the provided (e.g. SHA1) file digest.
// 
_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
EXTERN_C
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.
 // 
_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 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_ 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 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
);

//
// 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_ ULONG ProcessId,
    _Inout_ PUNICODE_STRING* ProcessImageName
);