Tamper Aware and Self Healing Code

// Self Healing 6.cpp
//

#include "stdafx.h"

#include "sha.h"        // SHA
#include "hex.h"        // HexEncoder
#include "files.h"      // FileSink
#include "filters.h"    // StringSink
#include "gzip.h"

VOID ImageInformation( HMODULE& hModule, PVOID& pVirtualAddress,
                       PVOID& pCodeStart, SIZE_T& dwCodeSize,
                       PVOID& pCodeEnd );
                       
VOID DumpImageInformation( HMODULE hModule, PVOID pVirtualAddress,
                           PVOID pCodeStart, SIZE_T dwCodeSize,
                           PVOID pCodeEnd );
              
VOID CalculateImageHash( PVOID pCodeStart, SIZE_T dwCodeSize,
                         PBYTE pcbDigest );
                         
VOID DumpHash( PBYTE pcbDigest, SIZE_T dwSize, std::string message );

VOID ExportTextImage( const std::string& filename,
                      PVOID pCodeStart, SIZE_T dwCodeSize );
                      
VOID ImportTextImage( const std::string& filename,
                      PBYTE pBuffer, SIZE_T dwBufferSize );                      

VOID HexDump( LPCVOID pcbStartAddress,
              LPCVOID pDisplayBaseAddress = (PVOID)-1,
              DWORD dwSize = DEFAULT_DUMP_SIZE );

// These values must be Global. Place them inside
//   main(), and you get different code generation
//   after each back patch operation.

#ifdef _DEBUG
BYTE cbExpectedImageHash[ CryptoPP::SHA224::DIGESTSIZE ] =
    { 0xF3,0x04,0xD3,0xAB,0x49,0xE6,0xC9,
      0x31,0xB8,0x3A,0x18,0x01,0xEF,0xD9,
      0xC9,0x48,0x31,0x48,0x69,0xD9,0x25,
      0xD6,0x3C,0x20,0x33,0x69,0xF2,0x9E };
#else
BYTE cbExpectedImageHash[ CryptoPP::SHA224::DIGESTSIZE ] =
    { 0x0C,0x34,0x05,0xEE,0x47,0xB3,0xAB,
      0x11,0xD4,0x21,0xE3,0x95,0x86,0x09,
      0xEE,0x7D,0x0D,0x38,0xD9,0x92,0x17,
      0x6B,0xBA,0x7F,0x00,0x83,0x36,0xF8 };
#endif

BYTE cbCalculatedImageHash[ CryptoPP::SHA224::DIGESTSIZE ];

int _tmain(int argc, _TCHAR* argv[])
{
    HMODULE hModule = NULL;
    PVOID   pVirtualAddress = NULL;
    PVOID   pCodeStart = NULL;
    PVOID   pCodeEnd = NULL;    
    SIZE_T  dwCodeSize = 0;
       
    ImageInformation( hModule, pVirtualAddress, pCodeStart,
                      dwCodeSize, pCodeEnd );                   
    
    DumpImageInformation( hModule, pVirtualAddress, pCodeStart,
                          dwCodeSize, pCodeEnd );                          
    
    std::string filename = "TextImage.gz";      
    ExportTextImage( filename, pCodeStart, dwCodeSize );
    
    SIZE_T dwBufferSize = dwCodeSize;
    PBYTE pArchiveBuffer = new BYTE[ dwBufferSize ];
    if( NULL == pArchiveBuffer ) { return -1; }
    
    ImportTextImage( filename, pArchiveBuffer, dwBufferSize );    
    
    std::tcout << _T("Original TEXT Section") << std::endl;
    HexDump( pCodeStart, pCodeStart, DUMP_SIZE );
    std::tcout << std::endl;
    
    std::tcout << _T("Archived TEXT Section") << std::endl;
    HexDump( pArchiveBuffer, (LPCVOID)NULL, DUMP_SIZE );
    
    CalculateImageHash( pCodeStart, dwCodeSize, cbCalculatedImageHash );
    
    DumpHash( cbExpectedImageHash, CryptoPP::SHA224::DIGESTSIZE,
              "SHA-224 Expected Image Hash" );    
    DumpHash( cbCalculatedImageHash, CryptoPP::SHA224::DIGESTSIZE,
              "SHA-224 Calculated Image Hash" );   
    
    if( 0 == memcmp( cbExpectedImageHash, cbCalculatedImageHash,
        CryptoPP::SHA224::DIGESTSIZE ) )
    {
        std::tcout << _T("Image is verified.") << std::endl;
    }
    else
    {
        std::tcout << _T("Image has been modified.") << std::endl;
    }      

    if( NULL != pArchiveBuffer ) { delete[] pArchiveBuffer; }
    
    return 0;
}

VOID ImportTextImage( const std::string& filename,
                      PBYTE pBuffer, SIZE_T dwBufferSize )
{

    try {
        std::string RecoveredTextSection;
        CryptoPP::FileSource( filename.c_str(), true,
            new CryptoPP::Gunzip(
                new CryptoPP::StringSink( RecoveredTextSection )
            ) // Gunzip
        ); // FileSource
    
        if( RecoveredTextSection.length() > dwBufferSize )
        {
            std::tcerr << _T("ImportTextImage: Executing Buffer Overflow");
        }
        
        memcpy( pBuffer, RecoveredTextSection.c_str(), dwBufferSize );       
    }   
    
    catch( CryptoPP::Exception& e )
    {
        std::cerr << e.what() << std:: endl;
    }
    
    catch( ... )
    {
        std::tcerr << _T("Caught Unknown Exception");
        std::tcerr << std:: endl;
    }
}
                      
VOID ExportTextImage( const std::string& filename,
                      PVOID pCodeStart, SIZE_T dwCodeSize )
{
    try {             
        CryptoPP::Gzip zipper(
            new CryptoPP::FileSink (filename.c_str(), true ),
        CryptoPP::Gzip::MAX_DEFLATE_LEVEL ); // Gzip
        
        zipper.Put( (byte*)pCodeStart, dwCodeSize );
        zipper.MessageEnd( );
    }
    
    catch( CryptoPP::Exception& e )
    {
        std::cerr << e.what() << std:: endl;
    }
    
    catch( ... )
    {
        std::tcerr << _T("Caught Unknown Exception");
        std::tcerr << std:: endl;
    }
}

VOID DumpHash( PBYTE pcbDigest, SIZE_T dwSize, std::string message )
{               
    CryptoPP::HexEncoder encoder;
    std::string sink;
    
    encoder.Attach( new CryptoPP::StringSink (sink) );
    encoder.Put( pcbDigest, dwSize );
    encoder.MessageEnd();
    
    std::cout << std::endl;
    
    if( 0 != message.length() )
    {
        std::cout << message << std::endl;
    }
    
    std::cout << sink << std::endl << std::endl;
}   

VOID CalculateImageHash( PVOID pCodeStart, SIZE_T dwCodeSize,
                         PBYTE pcbDigest )
{
    CryptoPP::SHA224 hash;
    
    hash.Update( (PBYTE)pCodeStart, dwCodeSize );
    hash.Final( pcbDigest );    
}                

VOID DumpImageInformation( HMODULE hModule, PVOID pVirtualAddress,
                           PVOID pCodeStart, SIZE_T dwCodeSize,
                           PVOID pCodeEnd )
{
    std::tcout << _T("****************************************************");
    std::tcout << std::endl;
    std::tcout << _T("************* Memory Image Information *************");
    std::tcout << std::endl;
    std::tcout << _T("****************************************************");
    std::tcout << std::endl << std::endl;
    
    std::tcout << _T("         hModule: ");
    std::tcout << HEXADECIMAL_OUTPUT(8);
    std::tcout << hModule << std::endl;
            
    std::tcout << _T(" Virtual Address: ");
    std::tcout << HEXADECIMAL_OUTPUT(8);
    std::tcout << pVirtualAddress << std::endl;  
                    
    std::tcout << _T("     .text Start: ");
    std::tcout << HEXADECIMAL_OUTPUT(8);
    std::tcout << pCodeStart << std::endl;

    std::tcout << _T("      .text Size: ");
    std::tcout << HEXADECIMAL_OUTPUT(8);
    std::tcout << dwCodeSize << std::endl;    

    std::tcout << _T("       .text End: ");
    std::tcout << HEXADECIMAL_OUTPUT(8);
    std::tcout << pCodeEnd << std::endl;
               
    std::tcout << std::endl;
}                           

VOID ImageInformation( HMODULE& hModule, PVOID& pVirtualAddress,
                       PVOID& pCodeStart, SIZE_T& dwCodeSize,
                       PVOID& pCodeEnd )
{
    const UINT PATH_SIZE = 2 * MAX_PATH;
    TCHAR szFilename[ PATH_SIZE ] = { 0 };    

    __try {

        /////////////////////////////////////////////////
        /////////////////////////////////////////////////
        if( 0 == GetModuleFileName( NULL, szFilename, PATH_SIZE ) )
        {
            std::tcerr << _T("Error Retrieving Process Filename");
            std::tcerr << std::endl;
            __leave;
        }

        hModule = GetModuleHandle( szFilename );
        if( NULL == hModule )
        {
            std::tcerr << _T("Error Retrieving Process Module Handle");
            std::tcerr << std::endl;
            __leave;
        }

        /////////////////////////////////////////////////////////////
        /////////////////////////////////////////////////////////////
        PIMAGE_DOS_HEADER pDOSHeader = NULL;
        pDOSHeader = static_cast<PIMAGE_DOS_HEADER>( (PVOID)hModule );
        if( pDOSHeader->e_magic != IMAGE_DOS_SIGNATURE )
        {
            std::tcerr << _T("Error - File is not EXE Format");
            std::tcerr << std::endl;
            __leave;
        }

        /////////////////////////////////////////////////////////////
        /////////////////////////////////////////////////////////////
        PIMAGE_NT_HEADERS pNTHeader = NULL;
        pNTHeader = reinterpret_cast<PIMAGE_NT_HEADERS>(
            (PBYTE)hModule + pDOSHeader->e_lfanew );
        if( pNTHeader->Signature != IMAGE_NT_SIGNATURE )
        {
            std::tcerr << _T("Error - File is not PE Format") << std::endl;
            __leave;
        }

        /////////////////////////////////////////////////////////////
        /////////////////////////////////////////////////////////////
        PIMAGE_FILE_HEADER pFileHeader = NULL;
        pFileHeader = reinterpret_cast<PIMAGE_FILE_HEADER>(
            (PBYTE)&pNTHeader->FileHeader );

        /////////////////////////////////////////////////////////////
        /////////////////////////////////////////////////////////////
        PIMAGE_OPTIONAL_HEADER pOptionalHeader = NULL;
        pOptionalHeader = reinterpret_cast<PIMAGE_OPTIONAL_HEADER>(
            (PBYTE)&pNTHeader->OptionalHeader );

        /////////////////////////////////////////////////////////////
        /////////////////////////////////////////////////////////////
        if( IMAGE_NT_OPTIONAL_HDR32_MAGIC !=
            pNTHeader->OptionalHeader.Magic )
        {
            std::tcerr << _T("Error - File is not 32 bit") << std::endl;
            __leave;
        }

        /////////////////////////////////////////////////////////////
        /////////////////////////////////////////////////////////////
        PIMAGE_SECTION_HEADER pSectionHeader = NULL;
        pSectionHeader = reinterpret_cast<PIMAGE_SECTION_HEADER>(
            (PBYTE)&pNTHeader->OptionalHeader +
            pNTHeader->FileHeader.SizeOfOptionalHeader );

        /////////////////////////////////////////////////////////////
        /////////////////////////////////////////////////////////////
        const CHAR TEXT[] = ".text";
        const CHAR BSSTEXT[] = ".textbss";
        UINT nSectionCount = pNTHeader->FileHeader.NumberOfSections;
        CHAR szSectionName[ IMAGE_SIZEOF_SHORT_NAME + 1 ];
        szSectionName[ IMAGE_SIZEOF_SHORT_NAME ] = '\0';
        for( UINT i = 0; i < nSectionCount; i++ )
        {
            memcpy( szSectionName, pSectionHeader->Name,
                    IMAGE_SIZEOF_SHORT_NAME );

            if( 0 == strncmp( TEXT, szSectionName,
                              IMAGE_SIZEOF_SHORT_NAME ) )
            {
                std::tcout << std::endl;
                break;
            }

            pSectionHeader++;
        }

        /////////////////////////////////////////////////////////////
        /////////////////////////////////////////////////////////////
        if( 0 != strncmp( TEXT, szSectionName, IMAGE_SIZEOF_SHORT_NAME ) )
        {
            std::tcerr << _T("Error - Unable to locate ");
            std::cerr << TEXT;
            std::tcerr << _T(" TEXT") << std::endl;
            __leave;
        }

        /////////////////////////////////////////////////////////////
        /////////////////////////////////////////////////////////////
        pVirtualAddress = (PVOID)(pSectionHeader->VirtualAddress);
        
        dwCodeSize = pSectionHeader->Misc.VirtualSize;
        
        /////////////////////////////////////////////////////////////
        /////////////////////////////////////////////////////////////        
        pCodeStart = (PVOID)(((PBYTE)hModule) +
                     (SIZE_T)((PBYTE)pVirtualAddress) );
                     
        /////////////////////////////////////////////////////////////
        /////////////////////////////////////////////////////////////                      
        pCodeEnd = (PVOID)((PBYTE)pCodeStart + dwCodeSize );                      
    }

    __except( EXCEPTION_EXECUTE_HANDLER ) {
        std::tcerr << std::endl << _T("Caught Exception") << std::endl;
    }
}

VOID HexDump( LPCVOID pcbStartAddress, LPCVOID pDisplayBaseAddress,
              DWORD dwSize )
{
    /////////////////////////////////////////////////////////////
    /////////////////////////////////////////////////////////////
    std::tcout << HEXADECIMAL_OUTPUT(8);

    /////////////////////////////////////////////////////////////
    /////////////////////////////////////////////////////////////
    DWORD dwPosition = 0;
    const UINT BYTES_PER_LINE = 16;

    if( (PVOID)-1 == pDisplayBaseAddress )
    {
        pDisplayBaseAddress = pcbStartAddress;
    }

    do
    {
        if( 0 == dwPosition % BYTES_PER_LINE )
        {
            std::tcout << HEXADECIMAL_OUTPUT(8);
            std::tcout << (SIZE_T)((PBYTE)pDisplayBaseAddress + dwPosition);
            std::tcout << _T(": ");
        }

        std::tcout << HEXADECIMAL_OUTPUT(2);
        std::tcout << (SIZE_T)(((PBYTE)pcbStartAddress)[ dwPosition ]);
        std::tcout << _T(" ");

        if( 0 == (dwPosition+1) % BYTES_PER_LINE && 0 != dwPosition )
        {
            std::tcout << std::endl;
        }
        else
        if( 0 == (dwPosition+1) % (BYTES_PER_LINE/2) && 0 != dwPosition )
        {
            std::tcout << _T(" ");
        }

    } while( ++dwPosition < dwSize );
}