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

#include "Domito.Internal.h"
#include "Domito.MinCrypt.h"

DOMITO_CODE_INTEGRITY G_CI = {};


_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
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
)
{
    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)G_Memory.Allocate(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)G_Memory.Allocate(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)
    {
        G_Memory.Free(pBuf);
    }

    if (hbHash)
    {
        BCryptDestroyHash(hbHash);
    }

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

    return status;
}

_Success_(return == STATUS_SUCCESS)
_Must_inspect_result_
_IRQL_requires_max_(PASSIVE_LEVEL)
#pragma code_seg("PAGED")
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
)
{
    PAGED_CODE();

    NTSTATUS status = STATUS_SUCCESS;
    PVOID certDirectory = nullptr;
    KAPC_STATE systemContext = {};

    do
    {
        if (!G_CI.CiFreePolicyInfo || !G_CI.CiCheckSignedFile || !G_CI.CiVerifyHashInCatalog)
        {
            status = STATUS_NOT_IMPLEMENTED;
            break;
        }

        SigningTime->QuadPart = 0;

        G_CI.CiFreePolicyInfo(PolicyInfo);
        G_CI.CiFreePolicyInfo(TimeStampPolicyInfo);

        if (HashSize != MINCRYPT_SHA1_LENGTH && HashSize != MINCRYPT_SHA256_LENGTH)
        {
            status = STATUS_INVALID_IMAGE_HASH;
            break;
        }

        if (SecurityDirectory->Size != 0u &&
            SecurityDirectory->VirtualAddress != 0u)
        {
            certDirectory = G_Memory.Allocate(SecurityDirectory->Size);
            if (certDirectory == NULL)
            {
                status = STATUS_INSUFFICIENT_RESOURCES;
                break;
            }

            LARGE_INTEGER offset = {};
            IO_STATUS_BLOCK ioStatusBlock = {};

            offset.LowPart = SecurityDirectory->VirtualAddress;

            status = ZwReadFile(
                FileHandle,
                NULL,
                NULL,
                NULL,
                &ioStatusBlock,
                certDirectory,
                SecurityDirectory->Size,
                &offset,
                NULL
            );

            if (status == STATUS_PENDING)
            {
                ZwWaitForSingleObject(
                    FileHandle,
                    FALSE,
                    NULL
                );

                MemoryBarrier();
                status = ioStatusBlock.Status;
            }

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

            KeStackAttachProcess(PsInitialSystemProcess, &systemContext);
            {
                status = G_CI.CiCheckSignedFile(
                    Hash,
                    HashSize,
                    HashAlgId,
                    certDirectory,
                    SecurityDirectory->Size,
                    PolicyInfo,
                    SigningTime,
                    TimeStampPolicyInfo
                );
            }
            KeUnstackDetachProcess(&systemContext);

            if (NT_SUCCESS(status))
            {
                break;
            }

            if (status != STATUS_INVALID_IMAGE_HASH)
            {
                break;
            }
        }

        KeStackAttachProcess(PsInitialSystemProcess, &systemContext);
        {
            status = G_CI.CiVerifyHashInCatalog(
                Hash,
                HashSize,
                HashAlgId,
                FALSE,
                0,
                0x2007F,
                PolicyInfo,
                NULL,
                SigningTime,
                TimeStampPolicyInfo
            );

            if (status == STATUS_INVALID_IMAGE_HASH)
            {
                status = G_CI.CiVerifyHashInCatalog(
                    Hash,
                    HashSize,
                    HashAlgId,
                    TRUE,
                    0,
                    0x2007F,
                    PolicyInfo,
                    NULL,
                    SigningTime,
                    TimeStampPolicyInfo
                );
            }
        }
        KeUnstackDetachProcess(&systemContext);

    } while (FALSE);

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

    return status;
}
#pragma code_seg()