#include <ntifs.h>
#include <ntintsafe.h>
#include <ntimage.h>
#include <bcrypt.h>

#include "Domito.h"
#include "ci.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
NTSYSAPI NTSTATUS NTAPI ZwQuerySystemInformation(
    ULONG SystemInformationClass,
    PVOID SystemInformation,
    ULONG SystemInformationLength,
    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
    );


_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
NTSTATUS
DomitoFindModuleBaseAddress(
    _In_ STRING ModuleName,
    _In_ PFN_DOMITO_ALLOCATE_ROUTINE Allocator,
    _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)Allocator(
        bufferSize
    );

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

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

    if (!NT_SUCCESS(status))
    {
        ExFreePool(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;
        }
    }

    ExFreePool(moduleInfo);

    return status;
}

_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
NTSTATUS
DomitoFindExportedFunctionAddress(
    _In_ PVOID ModuleBase,
    _In_ STRING FunctionName,
    _Inout_opt_ PVOID * FunctionAddress
)
{
    NTSTATUS status = STATUS_NOT_FOUND;
    ULONG exportSize;

    DECLARE_CONST_UNICODE_STRING(routineName, L"RtlImageDirectoryEntryToData");

    const t_RtlImageDirectoryEntryToData fp_RtlImageDirectoryEntryToData =
        (t_RtlImageDirectoryEntryToData)MmGetSystemRoutineAddress((PUNICODE_STRING)&routineName);

    if (fp_RtlImageDirectoryEntryToData == NULL)
    {
        return STATUS_NOT_IMPLEMENTED;
    }

    // Retrieve the export directory information
    const PIMAGE_EXPORT_DIRECTORY exportDirectory = (PIMAGE_EXPORT_DIRECTORY)fp_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;
}

_IRQL_requires_max_(DISPATCH_LEVEL)
UINT32
DomitoGetPortableExecutableDigestKind(
    _In_ PUCHAR pPeBytes,
    _In_ PIMAGE_DATA_DIRECTORY pImgDataDirectory
)
{
    if (!pImgDataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY].VirtualAddress)
    {
        return CALG_SHA1;
    }

    const PVOID pBase = pPeBytes + pImgDataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY].VirtualAddress;
    const LPWIN_CERTIFICATE pCert = (WIN_CERTIFICATE*)pBase;
    PUCHAR pMatch = NULL;

    if (NT_SUCCESS(DomitoMemorySearchPattern(
        (const PUCHAR)"\x30\x21\x30\x09\x06\x05\x2b\x0e\x03\x02\x1a\x05\x00\x04\x14",
        0x00,
        15,
        pBase,
        pCert->dwLength,
        (PVOID*)&pMatch
    )))
    {
        return CALG_SHA1;
    }

    if (NT_SUCCESS(DomitoMemorySearchPattern(
        (const PUCHAR)"\x30\xcc\x30\x0d\x06\x09\x60\x86\x48\x01\x65\x03\x04\x02\xcc\x05\x00\x04\xcc",
        0xcc,
        19,
        pBase,
        pCert->dwLength,
        (PVOID*)&pMatch
    )))
    {
        if (pMatch == NULL)
        {
            return CALG_SHA1;
        }

        if (pMatch[1] == 0x31 && pMatch[14] == 0x01 && pMatch[18] == 0x20)
        {
            return CALG_SHA256;
        }
        else if (pMatch[1] == 0x41 && pMatch[14] == 0x02 && pMatch[18] == 0x30)
        {
            return CALG_SHA384;
        }
        else if (pMatch[1] == 0x51 && pMatch[14] == 0x03 && pMatch[18] == 0x40)
        {
            return CALG_SHA512;
        }
    }

    return CALG_SHA1;
}

_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)
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
)
{
    PIMAGE_DATA_DIRECTORY pImgDataDirectory;
    BOOLEAN is64Bit;
    PUCHAR pHash = nullptr;
    NTSTATUS status;
    BCRYPT_ALG_HANDLE hbAlg = nullptr;
    BCRYPT_HASH_HANDLE hbHash = nullptr;
    UINT32 hashLength = 0;
    ULONG resultLength;
    ULONG fileOffset = 0;

    const PIMAGE_DOS_HEADER phDos = (IMAGE_DOS_HEADER*)pPeBytes;
    if (phDos->e_magic != IMAGE_DOS_SIGNATURE)
    {
        // Not an executable
        return STATUS_INVALID_IMAGE_NOT_MZ;
    }

    const PIMAGE_NT_HEADERS phNt = (PIMAGE_NT_HEADERS)(pPeBytes + phDos->e_lfanew);
    if (phNt->Signature != IMAGE_NT_SIGNATURE)
    {
        // Not a PE image
        return STATUS_INVALID_IMAGE_FORMAT;
    }

    //
    //	Check if PE is 32 or 64 bits
    //
    switch (phNt->OptionalHeader.Magic)
    {
    case IMAGE_NT_OPTIONAL_HDR32_MAGIC:
        is64Bit = FALSE;
        break;
    case IMAGE_NT_OPTIONAL_HDR64_MAGIC:
        is64Bit = TRUE;
        break;
    default:
        return STATUS_INVALID_IMAGE_FORMAT;			// Unsupported architecture
    }

    //
    // TODO: Not sure if  16 * 512 * 512 is right. Do something better!
    //
    ULONG copySize = phDos->e_lfanew + sizeof(IMAGE_FILE_HEADER) + 4 + 0x40;
    const PUCHAR pBuf = (PUCHAR)Allocator(16 * 512 * 512);
    if (!pBuf)
    {
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    //
    // Fetch the data directory from the PE.
    // Note that we can also use RtlImageDirectoryEntryToData to fetch
    // only the security directory entry but why not do it manually
    // when we have the full PE image anyway
    //
    if (is64Bit == TRUE)
    {
        pImgDataDirectory = ((PIMAGE_NT_HEADERS64)phNt)->OptionalHeader.DataDirectory;
    }
    else
    {
        pImgDataDirectory = ((PIMAGE_NT_HEADERS32)phNt)->OptionalHeader.DataDirectory;
    }

    //
    // Check if the PE file contains a signature
    //
    const BOOLEAN hasEmbeddedSig = pImgDataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY].VirtualAddress != 0;

    //
    // Boring bcrypt boilerplate code
    //
    const UINT32 peDigestKind = DomitoGetPortableExecutableDigestKind(pPeBytes, pImgDataDirectory);

    if (!NT_SUCCESS(status = BCryptOpenAlgorithmProvider(
        &hbAlg,
        DOMITO_CALG_TO_BCRYPT_ALGORITHM(peDigestKind), // bcrypt doesn't understand standard WinCrypt CALG IDs
        MS_PRIMITIVE_PROVIDER,
        0
    )))
    {
        goto cleanup;
    }

    if (!NT_SUCCESS(status = BCryptGetProperty(
        hbAlg,
        BCRYPT_HASH_LENGTH,
        (PUCHAR)&hashLength, // put the length of the hash into hashLength
        sizeof(hashLength),
        &resultLength,
        0
    )))
    {
        goto cleanup;
    }

    //
    // Allocate a buffer to store the resulting hash
    //
    pHash = (PUCHAR)Allocator(hashLength);
    if (!pHash)
    {
        status = STATUS_INSUFFICIENT_RESOURCES;
        goto cleanup;
    }
    RtlZeroMemory(pHash, hashLength);

    //
    // Create a handle to the resulting hash
    //
    if (!NT_SUCCESS(status = BCryptCreateHash(
        hbAlg,
        &hbHash,
        NULL,
        0,
        NULL,
        0,
        0
    )))
    {
        goto cleanup;
    }

    //
    // Ok, if we haven't BSODed yet, we're ready to continue
    // parsing the PE/COFF and hash the needed values
    //
    memcpy(pBuf, pPeBytes, copySize);
    fileOffset += copySize;

    //
    // Hash everything up to the Checksum then skip it
    //
    status = BCryptHashData(
        hbHash,
        pBuf,
        copySize,
        0
    );
    if (!NT_SUCCESS(status))
    {
        goto cleanup;
    }

    fileOffset += 4; // Skipping the checksum field here

    //
    // Reach the security directory information.
    // For x64 PEs it's 10 bytes further
    //
    copySize = 0x3C;
    if (is64Bit == TRUE)
    {
        copySize += 0x10;
    }

    memcpy(pBuf, pPeBytes + fileOffset, copySize);
    fileOffset += copySize;

    //
    // Again, hash everything up to here then skip the ignored field.
    //
    status = BCryptHashData(
        hbHash,
        pBuf,
        copySize,
        0
    );
    if (!NT_SUCCESS(status))
    {
        goto cleanup;
    }

    fileOffset += 8; // Skipping an other ignored field here

    //
    // Now hash everything else in the file up to the certificate data.
    //
    ULONG remaining = hasEmbeddedSig
        ? pImgDataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY].VirtualAddress - fileOffset
        : PeSize - fileOffset;

    while (remaining > 0)
    {
        const ULONG chunkSize = min(remaining, 4096);
        ULONG readBytes;
        memset(pBuf, 0, chunkSize);

        if (fileOffset + chunkSize > PeSize)
        {
            readBytes = PeSize - chunkSize;
            memcpy(pBuf, pPeBytes + fileOffset, readBytes);
            fileOffset += readBytes;
            break;
        }
        else
        {
            memcpy(pBuf, pPeBytes + fileOffset, chunkSize);
            readBytes = chunkSize;
            fileOffset += readBytes;
        }

        if (!NT_SUCCESS(status = BCryptHashData(
            hbHash,
            pBuf,
            readBytes,
            0
        )))
        {
            goto cleanup;
        }

        remaining -= readBytes;
    }

    //
    // Finish up the hash here
    //
    if (!NT_SUCCESS(status = BCryptFinishHash(
        hbHash,
        pHash,
        hashLength,
        0
    )))
    {
        goto cleanup;
    }

    //
    // Supply the results to the caller
    //
    *pDigestCalgOut = peDigestKind;
    *pDigestSizeOut = hashLength;
    *pDigestOut = pHash;
    if (hasEmbeddedSig == TRUE)
    {
        *pCertOut = (LPWIN_CERTIFICATE)(pPeBytes + pImgDataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY].VirtualAddress);
        *pSizeOfSecurityDirectory = pImgDataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY].Size;
    }
    else
    {
        *pCertOut = NULL;
        *pSizeOfSecurityDirectory = 0;
    }

    status = STATUS_SUCCESS;

cleanup:
    if (pBuf)
    {
        ExFreePool(pBuf);
    }

    if (hbHash)
    {
        BCryptDestroyHash(hbHash);
    }

    if (hbAlg)
    {
        BCryptCloseAlgorithmProvider(hbAlg, 0);
    }

    return status;
}