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

#include "Domito.Internal.h"


DOMITO_COMMON G_Common = {};

static STRING G_FN_CiFreePolicyInfo = RTL_CONSTANT_STRING("CiFreePolicyInfo");
static STRING G_FN_CiCheckSignedFile = RTL_CONSTANT_STRING("CiCheckSignedFile");
static STRING G_FN_CiVerifyHashInCatalog = RTL_CONSTANT_STRING("CiVerifyHashInCatalog");
static STRING G_FN_CiGetCertPublisherName = RTL_CONSTANT_STRING("CiGetCertPublisherName");
static STRING G_FN_CiSetTrustedOriginClaimId = RTL_CONSTANT_STRING("CiSetTrustedOriginClaimId");
static STRING G_FN_CiValidateFileObject = RTL_CONSTANT_STRING("CiValidateFileObject");

DECLARE_GLOBAL_CONST_UNICODE_STRING(G_QipRoutineName, L"ZwQueryInformationProcess");
DECLARE_GLOBAL_CONST_UNICODE_STRING(G_IdetdRoutineName, L"RtlImageDirectoryEntryToData");


_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
NTSTATUS
DomitoInit()
{
    //
    // Do those first since the follow-up code depends on them
    // 

    G_Common.ZwQueryInformationProcess =
        (t_ZwQueryInformationProcess)MmGetSystemRoutineAddress((PUNICODE_STRING)&G_QipRoutineName);
    G_Common.RtlImageDirectoryEntryToData =
        (t_RtlImageDirectoryEntryToData)MmGetSystemRoutineAddress((PUNICODE_STRING)&G_IdetdRoutineName);

    STRING ciModuleName = RTL_CONSTANT_STRING("\\SystemRoot\\system32\\CI.dll");
    PVOID driverBaseAddress = NULL, functionAddress = NULL;

    if (NT_SUCCESS(DomitoFindModuleBaseAddress(&ciModuleName, &driverBaseAddress)))
    {
        if (NT_SUCCESS(DomitoFindExportedFunctionAddress(driverBaseAddress, &G_FN_CiFreePolicyInfo, &functionAddress)))
        {
            G_CI.CiFreePolicyInfo = (t_CiFreePolicyInfo)functionAddress;
        }

        if (NT_SUCCESS(DomitoFindExportedFunctionAddress(driverBaseAddress, &G_FN_CiCheckSignedFile, &functionAddress)))
        {
            G_CI.CiCheckSignedFile = (t_CiCheckSignedFile)functionAddress;
        }

        if (NT_SUCCESS(DomitoFindExportedFunctionAddress(driverBaseAddress, &G_FN_CiVerifyHashInCatalog, &functionAddress)))
        {
            G_CI.CiVerifyHashInCatalog = (t_CiVerifyHashInCatalog)functionAddress;
        }

        if (NT_SUCCESS(DomitoFindExportedFunctionAddress(driverBaseAddress, &G_FN_CiGetCertPublisherName, &functionAddress)))
        {
            G_CI.CiGetCertPublisherName = (t_CiGetCertPublisherName)functionAddress;
        }

        if (NT_SUCCESS(DomitoFindExportedFunctionAddress(driverBaseAddress, &G_FN_CiSetTrustedOriginClaimId, &functionAddress)))
        {
            G_CI.CiSetTrustedOriginClaimId = (t_CiSetTrustedOriginClaimId)functionAddress;
        }

        if (NT_SUCCESS(DomitoFindExportedFunctionAddress(driverBaseAddress, &G_FN_CiValidateFileObject, &functionAddress)))
        {
            G_CI.CiValidateFileObject = (t_CiValidateFileObject)functionAddress;
        }
    }
    
    return STATUS_SUCCESS; // TODO: unused currently
}

_IRQL_requires_max_(PASSIVE_LEVEL)
void
DomitoShutdown()
{

}


_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
NTSTATUS
DomitoFindModuleBaseAddress(
    _In_ PANSI_STRING ModuleName,
    _Inout_opt_ PVOID * ModuleBase
)
{
    ULONG bufferSize = 0;
    PSYSTEM_MODULE_INFORMATION moduleInfo = NULL;

    const ULONG SystemModuleInformation = 11;

    // Query the required buffer size for module information
    NTSTATUS status = ZwQuerySystemInformation(
        SystemModuleInformation,
        &bufferSize,
        0,
        &bufferSize
    );

    if (status != STATUS_INFO_LENGTH_MISMATCH)
    {
        return status;
    }

    // Allocate memory for the module information
    moduleInfo = (PSYSTEM_MODULE_INFORMATION)G_Memory.Allocate(
        bufferSize
    );

    if (moduleInfo == NULL)
    {
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    // Retrieve the module information
    status = ZwQuerySystemInformation(
        SystemModuleInformation,
        moduleInfo,
        bufferSize,
        NULL
    );

    if (!NT_SUCCESS(status))
    {
        G_Memory.Free(moduleInfo);
        return status;
    }

    STRING currentImageName;

    status = STATUS_NOT_FOUND;
    // Iterate through the loaded modules and find the desired module
    for (ULONG i = 0; i < moduleInfo->Count; i++)
    {
        RtlInitAnsiString(&currentImageName, moduleInfo->Module[i].ImageName);

        if (0 == RtlCompareString(ModuleName, &currentImageName, TRUE))
        {
            status = STATUS_SUCCESS;

            // Found the module, store the base address
            if (ModuleBase)
            {
                *ModuleBase = moduleInfo->Module[i].Base;
            }
            break;
        }
    }

    G_Memory.Free(moduleInfo);

    return status;
}

_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
NTSTATUS
DomitoFindExportedFunctionAddress(
    _In_ PVOID ModuleBase,
    _In_ PANSI_STRING FunctionName,
    _Inout_opt_ PVOID * FunctionAddress
)
{
    NTSTATUS status = STATUS_NOT_FOUND;
    ULONG exportSize;
    
    if (G_Common.RtlImageDirectoryEntryToData == NULL)
    {
        return STATUS_NOT_IMPLEMENTED;
    }

    // Retrieve the export directory information
    const PIMAGE_EXPORT_DIRECTORY exportDirectory = (PIMAGE_EXPORT_DIRECTORY)G_Common.RtlImageDirectoryEntryToData(
        ModuleBase,
        TRUE,
        IMAGE_DIRECTORY_ENTRY_EXPORT,
        &exportSize
    );

    if (exportDirectory == NULL)
    {
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    STRING currentFunctionName;

    const PULONG functionAddresses = (PULONG)((ULONG_PTR)ModuleBase + exportDirectory->AddressOfFunctions);
    const PULONG functionNames = (PULONG)((ULONG_PTR)ModuleBase + exportDirectory->AddressOfNames);
    const PUSHORT functionOrdinals = (PUSHORT)((ULONG_PTR)ModuleBase + exportDirectory->AddressOfNameOrdinals);

    for (ULONG i = 0; i < exportDirectory->NumberOfNames; i++)
    {
        const char* functionName = (const char*)((ULONG_PTR)ModuleBase + functionNames[i]);
        const USHORT functionOrdinal = functionOrdinals[i];
        UNREFERENCED_PARAMETER(functionOrdinal);

        const ULONG functionRva = functionAddresses[i];
        const PVOID functionAddress = (PVOID)((ULONG_PTR)ModuleBase + functionRva);

        RtlInitAnsiString(&currentFunctionName, functionName);

        if (0 == RtlCompareString(FunctionName, &currentFunctionName, TRUE))
        {
            if (FunctionAddress)
            {
                status = STATUS_SUCCESS;
                *FunctionAddress = functionAddress;
            }
            break;
        }
    }

    return status;
}

_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(DISPATCH_LEVEL)
NTSTATUS
DomitoMemorySearchPattern(
    _In_ PCUCHAR pcPattern,
    _In_ UCHAR uWildcard,
    _In_ SIZE_T puLen,
    _In_ PVOID pcBase,
    _In_ SIZE_T puSize,
    _Outptr_result_maybenull_ PVOID * ppMatch
)
{
    ASSERT(ppMatch != NULL && pcPattern != NULL && pcBase != NULL);

    if (ppMatch == NULL || pcPattern == NULL || pcBase == NULL)
    {
        return STATUS_INVALID_PARAMETER;
    }

    *ppMatch = NULL;

    for (SIZE_T i = 0; i < puSize - puLen; i++)
    {
        BOOLEAN found = TRUE;
        for (SIZE_T j = 0; j < puLen; j++)
        {
            if (pcPattern[j] != uWildcard && pcPattern[j] != ((PCUCHAR)pcBase)[i + j])
            {
                found = FALSE;
                break;
            }
        }

        if (found)
        {
            *ppMatch = (PUCHAR)pcBase + i;
            return STATUS_SUCCESS;
        }
    }

    return STATUS_NOT_FOUND;
}

_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
NTSTATUS
DomitoReadFile(
    _In_ HANDLE FileHandle,
    _Out_ PVOID Buffer,
    _In_ ULONG BufferSize
)
{
    NTSTATUS status = STATUS_SUCCESS;
    IO_STATUS_BLOCK ioStatusBlock;

    // Read the file into memory using ZwReadFile
    if (!NT_SUCCESS(status = ZwReadFile(
        FileHandle,
        NULL,
        NULL,
        NULL,
        &ioStatusBlock,
        Buffer,
        BufferSize,
        NULL,
        NULL
    )))
    {
        return status;
    }

    // Check if the file was read successfully
    if (!NT_SUCCESS(ioStatusBlock.Status))
    {
        return ioStatusBlock.Status;
    }

    // File read successfully
    return status;
}

_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
#pragma code_seg("PAGED")
NTSTATUS
DomitoGetProcessImageName(
    _In_ ULONG ProcessId,
    _Inout_ PUNICODE_STRING * ProcessImageName
)
{
    NTSTATUS status = STATUS_UNSUCCESSFUL;
    ULONG returnedLength;
    HANDLE hProcess = NULL;

    PAGED_CODE(); // this eliminates the possibility of the IDLE Thread/Process

    if (ProcessId == 0 || ProcessId == 4)
    {
        return STATUS_INVALID_PARAMETER_1;
    }

    if (ProcessImageName == NULL)
    {
        return STATUS_INVALID_PARAMETER_2;
    }

    CLIENT_ID cid;
    cid.UniqueProcess = (HANDLE)ProcessId;
    cid.UniqueThread = NULL;

    OBJECT_ATTRIBUTES objAttr;
    InitializeObjectAttributes(
        &objAttr,
        NULL,
        OBJ_KERNEL_HANDLE,
        NULL,
        NULL
    );

    if (!NT_SUCCESS(status = ZwOpenProcess(
        &hProcess,
        PROCESS_ALL_ACCESS,
        &objAttr,
        &cid
    )))
    {
        return status;
    }

    if (G_Common.ZwQueryInformationProcess == NULL)
    {
        status = STATUS_NOT_IMPLEMENTED;
        goto cleanUp;
    }

    /* Query the actual size of the process path */
    status = G_Common.ZwQueryInformationProcess(
        hProcess,
        ProcessImageFileName,
        NULL, // buffer
        0,    // buffer size
        &returnedLength
    );

    if (STATUS_INFO_LENGTH_MISMATCH != status)
    {
        goto cleanUp;
    }

    *ProcessImageName = (PUNICODE_STRING)G_Memory.Allocate(returnedLength);

    if (*ProcessImageName == NULL)
    {
        status = STATUS_INSUFFICIENT_RESOURCES;
        goto cleanUp;
    }

    /* Retrieve the process path from the handle to the process */
    if (!NT_SUCCESS(status = G_Common.ZwQueryInformationProcess(
        hProcess,
        ProcessImageFileName,
        *ProcessImageName,
        returnedLength,
        &returnedLength
    )))
    {
        G_Memory.Free(*ProcessImageName);
    }

cleanUp:

    if (hProcess)
    {
        ZwClose(hProcess);
    }

    return status;
}
#pragma code_seg()