#include "Domito.Internal.h"

_Use_decl_annotations_
NTSTATUS
DomitoFindModuleExportedFunctionAddress(
    PANSI_STRING ModuleName,
    PANSI_STRING FunctionName,
    PVOID* FunctionAddress
)
{
    PAGED_CODE();

    if (!FunctionAddress)
    {
        return STATUS_INVALID_PARAMETER;
    }

    *FunctionAddress = nullptr;

    if (!ModuleName || !ModuleName->Buffer || !ModuleName->Length || ModuleName->Length > ModuleName->MaximumLength ||
        !FunctionName || !FunctionName->Buffer || !FunctionName->Length ||
        FunctionName->Length > FunctionName->MaximumLength)
    {
        return STATUS_INVALID_PARAMETER;
    }

    PVOID moduleBase = nullptr;
    NTSTATUS status = DomitoFindModuleBaseAddress(ModuleName, &moduleBase);

    if (!NT_SUCCESS(status))
    {
        return status;
    }

    return DomitoFindExportedFunctionAddress(moduleBase, FunctionName, FunctionAddress);
}

_Use_decl_annotations_
void
DomitoFreeProcessImageName(
    PUNICODE_STRING ImageName
)
{
    PAGED_CODE();

    if (ImageName)
    {
        G_Memory.Free(ImageName);
    }
}

_Use_decl_annotations_
NTSTATUS
DomitoGetProcessImageNameEx(
    PEPROCESS Process,
    PUNICODE_STRING* ImageName
)
{
    PAGED_CODE();

    if (!ImageName)
    {
        return STATUS_INVALID_PARAMETER;
    }

    *ImageName = nullptr;

    if (!Process)
    {
        return STATUS_INVALID_PARAMETER;
    }

    PUNICODE_STRING processImageName = nullptr;
    NTSTATUS status = SeLocateProcessImageName(Process, &processImageName);

    if (NT_SUCCESS(status))
    {
        if (!processImageName || !processImageName->Buffer || !processImageName->Length ||
            (processImageName->Length & 1) || processImageName->Length > processImageName->MaximumLength ||
            processImageName->Length > MAXUSHORT - sizeof(WCHAR))
        {
            status = STATUS_INVALID_BUFFER_SIZE;
        }
        else
        {
            const SIZE_T size = sizeof(UNICODE_STRING) + processImageName->Length + sizeof(WCHAR);
            PUNICODE_STRING imageName = static_cast<PUNICODE_STRING>(G_Memory.Allocate(size));

            if (!imageName)
            {
                status = STATUS_INSUFFICIENT_RESOURCES;
            }
            else
            {
                imageName->Buffer = reinterpret_cast<PWCH>(imageName + 1);
                imageName->Length = processImageName->Length;
                imageName->MaximumLength = processImageName->Length + sizeof(WCHAR);
                RtlCopyMemory(imageName->Buffer, processImageName->Buffer, processImageName->Length);
                imageName->Buffer[imageName->Length / sizeof(WCHAR)] = 0;
                *ImageName = imageName;
            }
        }
    }

    if (processImageName)
    {
        ExFreePool(processImageName);
    }

    return status;
}

#if (NTDDI_VERSION >= NTDDI_WIN10)
_Use_decl_annotations_
NTSTATUS
DomitoValidateFileHandleEx(
    HANDLE FileHandle,
    UINT32 ValidationFlags,
    UINT8 RequestedSigningLevel,
    DOMITO_CI_SIGNATURE_SELECTION Selection,
    DOMITO_CI_VALIDATION_RESULT* Result
)
{
    PAGED_CODE();

    if (!Result)
    {
        return STATUS_INVALID_PARAMETER;
    }

    DomitoFreeCiValidationResult(Result);

    if (!FileHandle)
    {
        return STATUS_INVALID_PARAMETER;
    }

    PFILE_OBJECT fileObject = nullptr;
    NTSTATUS status = ObReferenceObjectByHandle(
        FileHandle, FILE_READ_DATA, *IoFileObjectType, KernelMode, reinterpret_cast<PVOID*>(&fileObject), nullptr
    );

    if (!NT_SUCCESS(status))
    {
        return status;
    }

    status = DomitoValidateFileObjectEx(fileObject, ValidationFlags, RequestedSigningLevel, Selection, Result);
    ObDereferenceObject(fileObject);
    return status;
}

_Use_decl_annotations_
NTSTATUS
DomitoValidateFilePathEx(
    const UNICODE_STRING* Path,
    UINT32 ValidationFlags,
    UINT8 RequestedSigningLevel,
    DOMITO_CI_SIGNATURE_SELECTION Selection,
    DOMITO_CI_VALIDATION_RESULT* Result
)
{
    PAGED_CODE();

    if (!Result)
    {
        return STATUS_INVALID_PARAMETER;
    }

    DomitoFreeCiValidationResult(Result);

    if (!Path || !Path->Buffer || Path->Length < sizeof(WCHAR) || (Path->Length & 1) ||
        Path->Length > Path->MaximumLength || Path->Buffer[0] != L'\\')
    {
        return STATUS_INVALID_PARAMETER;
    }

    for (USHORT i = 0; i < Path->Length / sizeof(WCHAR); ++i)
    {
        if (!Path->Buffer[i])
        {
            return STATUS_INVALID_PARAMETER;
        }
    }

    OBJECT_ATTRIBUTES attributes;
    InitializeObjectAttributes(
        &attributes, const_cast<PUNICODE_STRING>(Path), OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, nullptr, nullptr
    );
    HANDLE fileHandle = nullptr;
    IO_STATUS_BLOCK iosb = {};
    NTSTATUS status = ZwOpenFile(
        &fileHandle,
        FILE_READ_DATA | SYNCHRONIZE,
        &attributes,
        &iosb,
        FILE_SHARE_READ,
        FILE_NON_DIRECTORY_FILE | FILE_SYNCHRONOUS_IO_NONALERT
    );

    if (!NT_SUCCESS(status))
    {
        return status;
    }

    status = DomitoValidateFileHandleEx(fileHandle, ValidationFlags, RequestedSigningLevel, Selection, Result);
    ZwClose(fileHandle);
    return status;
}

#if defined(_M_AMD64) || defined(_AMD64_)
static BOOLEAN
DomitoCiOwnedSpan(
    const MINCRYPT_CHAIN_INFO* Chain,
    const void* Data,
    SIZE_T Size
)
{
    const ULONG_PTR chainBaseAddress = reinterpret_cast<ULONG_PTR>(Chain);
    const ULONG_PTR blobAddress = reinterpret_cast<ULONG_PTR>(Data);
    // Empty blobs may be NULL; they are never dereferenced by this accessor.
    if (!Size)
    {
        return TRUE;
    }

    return blobAddress >= chainBaseAddress && blobAddress - chainBaseAddress >= 56 &&
           blobAddress - chainBaseAddress <= Chain->Size && Size <= Chain->Size - (blobAddress - chainBaseAddress);
}
#endif

_Use_decl_annotations_
NTSTATUS
DomitoGetCiChainView(
    const MINCRYPT_POLICY_INFO* Policy,
    DOMITO_CI_CHAIN_LAYOUT Layout,
    DOMITO_CI_CHAIN_VIEW* View
)
{
    PAGED_CODE();

    if (!View)
    {
        return STATUS_INVALID_PARAMETER;
    }

    *View = {};

    if (!Policy || Layout != DomitoCiChainLayoutObservedX64)
    {
        return STATUS_INVALID_PARAMETER;
    }
#if defined(_M_AMD64) || defined(_AMD64_)
    C_ASSERT(sizeof(MINCRYPT_CHAIN_ELEMENT) == 120);
    C_ASSERT(FIELD_OFFSET(MINCRYPT_CHAIN_INFO, ChainElements) == 40);
    C_ASSERT(FIELD_OFFSET(MINCRYPT_CHAIN_INFO, NumberOfChainElement) == 48);

    if (Policy->Size != sizeof(*Policy))
    {
        return STATUS_NOT_SUPPORTED;
    }

    const MINCRYPT_CHAIN_INFO* chain = Policy->ChainInfo;

    if (!chain)
    {
        return STATUS_NOT_FOUND;
    }

    if (chain->Size < 56)
    {
        return STATUS_INVALID_BUFFER_SIZE;
    }

    const ULONG count = chain->NumberOfChainElement;

    if (!count)
    {
        return STATUS_NOT_FOUND;
    }

    if (count > (chain->Size - 56) / sizeof(MINCRYPT_CHAIN_ELEMENT) ||
        !DomitoCiOwnedSpan(chain, chain->ChainElements, static_cast<SIZE_T>(count) * sizeof(MINCRYPT_CHAIN_ELEMENT)))
    {
        return STATUS_INVALID_BUFFER_SIZE;
    }

    for (ULONG i = 0; i < count; ++i)
    {
        const MINCRYPT_CHAIN_ELEMENT& element = chain->ChainElements[i];

        if (element.HashSize > sizeof(element.Hash) || !element.Certificate.Size ||
            !DomitoCiOwnedSpan(chain, element.Certificate.Data, element.Certificate.Size) ||
            !DomitoCiOwnedSpan(chain, element.Subject.Data, element.Subject.Size) ||
            !DomitoCiOwnedSpan(chain, element.Issuer.Data, element.Issuer.Size))
        {
            return STATUS_INVALID_BUFFER_SIZE;
        }
    }

    ULONG rootId;
    RtlCopyMemory(&rootId, reinterpret_cast<const UCHAR*>(chain) + 52, sizeof(rootId));
    *View = {chain->ChainElements, count, rootId};
    return STATUS_SUCCESS;
#else
    return STATUS_NOT_SUPPORTED;
#endif
}

static MINCRYPT_ALLOCATE_ROUTINE DomitoAllocatePublisher;

static PVOID NTAPI
DomitoAllocatePublisher(
    SIZE_T Size
)
{
#pragma warning(suppress : 4996)
    return ExAllocatePoolWithTag(PagedPool, Size, DOMITO_POOL_TAG);
}

_Use_decl_annotations_
void
DomitoFreeCertificatePublisherName(
    PUNICODE_STRING PublisherName
)
{
    PAGED_CODE();

    if (!PublisherName)
    {
        return;
    }

    if (PublisherName->Buffer)
    {
        ExFreePoolWithTag(PublisherName->Buffer, DOMITO_POOL_TAG);
    }

    *PublisherName = {};
}

_Use_decl_annotations_
NTSTATUS
DomitoGetCertificatePublisherName(
    const MINCERT_BLOB* Certificate,
    PUNICODE_STRING PublisherName
)
{
    PAGED_CODE();

    if (!PublisherName)
    {
        return STATUS_INVALID_PARAMETER;
    }

    *PublisherName = {};

    if (!Certificate || !Certificate->Data || !Certificate->Size)
    {
        return STATUS_INVALID_PARAMETER;
    }

    ANSI_STRING moduleName = RTL_CONSTANT_STRING("\\SystemRoot\\system32\\CI.dll");
    ANSI_STRING exportName = RTL_CONSTANT_STRING("CiGetCertPublisherName");
    PVOID publisherRoutine = nullptr;
    NTSTATUS status = DomitoFindModuleExportedFunctionAddress(&moduleName, &exportName, &publisherRoutine);

    if (!NT_SUCCESS(status))
    {
        return status;
    }

    if (!publisherRoutine)
    {
        return STATUS_PROCEDURE_NOT_FOUND;
    }

    typedef NTSTATUS(NTAPI * GET_PUBLISHER)(MINCERT_BLOB*, PMINCRYPT_ALLOCATE_ROUTINE, PUNICODE_STRING);
    const GET_PUBLISHER getPublisher = reinterpret_cast<GET_PUBLISHER>(publisherRoutine);
    status = getPublisher(const_cast<MINCERT_BLOB*>(Certificate), DomitoAllocatePublisher, PublisherName);

    if (NT_SUCCESS(status) && (PublisherName->Length > PublisherName->MaximumLength || (PublisherName->Length & 1) ||
                               !PublisherName->Length || !PublisherName->Buffer))
    {
        status = STATUS_INVALID_BUFFER_SIZE;
    }

    if (!NT_SUCCESS(status))
    {
        DomitoFreeCertificatePublisherName(PublisherName);
    }

    return status;
}
#endif