Added DomitoCalculatePortableExecutableDigest

Domito.sln.DotSettings

@@ -1,6 +1,7 @@
 <wpf:ResourceDictionary xml:space="preserve" xmlns:x="http://schemas.microsoft.com/winfx/2006/xaml" xmlns:s="clr-namespace:System;assembly=mscorlib" xmlns:ss="urn:shemas-jetbrains-com:settings-storage-xaml" xmlns:wpf="http://schemas.microsoft.com/winfx/2006/xaml/presentation">
 	<s:Boolean x:Key="/Default/UserDictionary/Words/=BCRYPT/@EntryIndexedValue">True</s:Boolean>
 	<s:Boolean x:Key="/Default/UserDictionary/Words/=CALG/@EntryIndexedValue">True</s:Boolean>
+	<s:Boolean x:Key="/Default/UserDictionary/Words/=COFF/@EntryIndexedValue">True</s:Boolean>
 	<s:Boolean x:Key="/Default/UserDictionary/Words/=Domito/@EntryIndexedValue">True</s:Boolean>
 	<s:Boolean x:Key="/Default/UserDictionary/Words/=LPWIN/@EntryIndexedValue">True</s:Boolean>
 	<s:Boolean x:Key="/Default/UserDictionary/Words/=PDOMITO/@EntryIndexedValue">True</s:Boolean>

include/Domito.h

@@ -164,3 +164,19 @@ DomitoReadFile(
     _Out_ PVOID Buffer,
     _In_ ULONG BufferSize
 );
+
+_Success_(return == STATUS_SUCCESS)
+_Must_inspect_result_
+_IRQL_requires_max_(PASSIVE_LEVEL)
+EXTERN_C
+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
+);

src/Domito.cpp

@@ -355,3 +355,284 @@ DomitoReadFile(
     // 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;
+}