#include "Domito.Internal.h"

#if (NTDDI_VERSION >= NTDDI_WIN10)

static UINT32
DomitoCiDigestSize(
    ALG_ID Algorithm
)
{
    switch (Algorithm)
    {
    case CALG_SHA1:
        return MINCRYPT_SHA1_LENGTH;
    case CALG_SHA_256:
        return MINCRYPT_SHA256_LENGTH;
    case CALG_SHA_384:
        return MINCRYPT_SHA384_LENGTH;
    case CALG_SHA_512:
        return MINCRYPT_SHA512_LENGTH;
    default:
        return 0;
    }
}

static void
DomitoClearCiDigest(
    DOMITO_CI_VALIDATION_RESULT* Result
)
{
    RtlZeroMemory(Result->Hash, sizeof(Result->Hash));
    Result->HashSize = 0;
    Result->HashAlgId = 0;
    Result->SigningTime.QuadPart = 0;
}

_Use_decl_annotations_
void
DomitoFreeCiValidationResult(
    DOMITO_CI_VALIDATION_RESULT* Result
)
{
    PAGED_CODE();

    if (!Result)
    {
        return;
    }

    CiFreePolicyInfo(&Result->PolicyInfo);
    CiFreePolicyInfo(&Result->TimeStampPolicyInfo);
    RtlZeroMemory(Result, sizeof(*Result));
}

_Use_decl_annotations_
NTSTATUS
DomitoValidateFileObjectEx(
    PFILE_OBJECT FileObject,
    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 (!FileObject)
    {
        return STATUS_INVALID_PARAMETER;
    }

    if (Selection != DomitoCiSignatureSelectionNative &&
        Selection != DomitoCiSignatureSelectionFirstSuccessfulExperimental)
    {
        return STATUS_INVALID_PARAMETER;
    }

    if (Selection == DomitoCiSignatureSelectionFirstSuccessfulExperimental)
    {
#if defined(_M_AMD64) || defined(_AMD64_)
        ValidationFlags |= DOMITO_CI_VALIDATE_FLAG_FIRST_SUCCESS_EXPERIMENTAL;
#else
        return STATUS_NOT_SUPPORTED;
#endif
    }

    Result->FileObjectValidationFlags = ValidationFlags;
    Result->RequestedSigningLevel = RequestedSigningLevel;
    Result->HashSize = sizeof(Result->Hash);
    NTSTATUS status = CiValidateFileObject(
        FileObject,
        ValidationFlags,
        RequestedSigningLevel,
        &Result->PolicyInfo,
        &Result->TimeStampPolicyInfo,
        &Result->SigningTime,
        Result->Hash,
        &Result->HashSize,
        &Result->HashAlgId
    );

    if (NT_SUCCESS(status))
    {
        const UINT32 expected = DomitoCiDigestSize(Result->HashAlgId);

        if (!expected)
        {
            status = STATUS_NOT_SUPPORTED;
        }
        else if (Result->HashSize != expected || Result->HashSize > sizeof(Result->Hash))
        {
            status = STATUS_INVALID_BUFFER_SIZE;
        }
    }

    if (!NT_SUCCESS(status))
    {
        DomitoClearCiDigest(Result);
    }
    else
    {
        RtlZeroMemory(Result->Hash + Result->HashSize, sizeof(Result->Hash) - Result->HashSize);
    }

    return status;
}

_Use_decl_annotations_
NTSTATUS
DomitoValidatePortableExecutableEmbeddedSignature(
    const UCHAR* PeBytes,
    ULONG PeSize,
    ALG_ID HashAlgId,
    DOMITO_CI_VALIDATION_RESULT* Result
)
{
    PAGED_CODE();

    if (!Result)
    {
        return STATUS_INVALID_PARAMETER;
    }

    DomitoFreeCiValidationResult(Result);

    if (!PeBytes)
    {
        return STATUS_INVALID_PARAMETER;
    }

    const UINT32 expected = DomitoCiDigestSize(HashAlgId);

    if (!expected)
    {
        return STATUS_NOT_SUPPORTED;
    }

    PVOID digest = nullptr;
    ULONG digestSize = 0;
    ULONG securitySize = 0;
    LPWIN_CERTIFICATE certificate = nullptr;
    NTSTATUS status = DomitoCalculatePortableExecutableDigestEx(
        PeBytes, PeSize, HashAlgId, &digestSize, &digest, &certificate, &securitySize
    );

    if (NT_SUCCESS(status))
    {
        if (!certificate || !securitySize)
        {
            status = STATUS_INVALID_IMAGE_HASH;
        }
        else if (!digest || digestSize != expected || digestSize > sizeof(Result->Hash))
        {
            status = STATUS_INVALID_BUFFER_SIZE;
        }
        else
        {
            status = CiCheckSignedFile(
                digest,
                digestSize,
                HashAlgId,
                certificate,
                securitySize,
                &Result->PolicyInfo,
                &Result->SigningTime,
                &Result->TimeStampPolicyInfo
            );

            if (NT_SUCCESS(status))
            {
                RtlCopyMemory(Result->Hash, digest, digestSize);
                Result->HashSize = digestSize;
                Result->HashAlgId = HashAlgId;
            }
        }
    }

    if (digest)
    {
        DomitoFreePortableExecutableDigest(digest);
    }

    if (!NT_SUCCESS(status))
    {
        DomitoClearCiDigest(Result);
    }

    return status;
}
#endif // NTDDI_VERSION >= NTDDI_WIN10