An Analysis of the Windows PE Checksum Algorithm

// PEChecksum.cpp : implementation file
//

#include "stdafx.h"
#include "PEChecksum.h"
#include "PEChecksumDlg.h"

#ifdef _DEBUG
#define new DEBUG_NEW
#endif

// CPEChecksum dialog

CPEChecksum::CPEChecksum(CWnd* pParent /*=NULL*/)
: CDialog(CPEChecksum::IDD, pParent)
{
    m_hIcon = AfxGetApp()->LoadIcon(IDR_MAINFRAME);
}

void CPEChecksum::DoDataExchange(CDataExchange* pDX)
{
    CDialog::DoDataExchange(pDX);
}

BEGIN_MESSAGE_MAP(CPEChecksum, CDialog)
    ON_WM_PAINT()
    ON_WM_QUERYDRAGICON()
    //}}AFX_MSG_MAP
    ON_BN_CLICKED(IDC_BTN_BROWSE, &CPEChecksum::OnBnClickedBrowse)
    ON_BN_CLICKED(IDC_BTN_RETRIEVE, &CPEChecksum::OnBnClickedRetrieve)
    ON_BN_CLICKED(IDC_BTN_WRITE_NEW_VALUE, &CPEChecksum::OnBnClickedBtnWriteNewValue)
END_MESSAGE_MAP()


// CPEChecksum message handlers

BOOL CPEChecksum::OnInitDialog()
{
    CDialog::OnInitDialog();

    // Set the icon for this dialog. The framework does this automatically
    // when the application's main window is not a dialog
    SetIcon(m_hIcon, TRUE); // Set big icon
    SetIcon(m_hIcon, FALSE); // Set small icon

    CEdit* pEdit = static_cast<CEdit*>
        ( GetDlgItem( IDC_EDIT_BROWSE ) );

    if( NULL != pEdit )
    {
        pEdit->SetWindowText( _T("C:\\WINDOWS\\NOTEPAD.EXE") );
    }

    return TRUE; // return TRUE unless you set the focus to a control
}

// If you add a minimize button to your dialog, you will need the code below
// to draw the icon. For MFC applications using the document/view model,
// this is automatically done for you by the framework.

void CPEChecksum::OnPaint()
{
    if( IsIconic() )
    {
        CPaintDC dc( this ); // device context for painting

        SendMessage(WM_ICONERASEBKGND, reinterpret_cast<WPARAM>(dc.GetSafeHdc()), 0);

        // Center icon in client rectangle
        int cxIcon = GetSystemMetrics(SM_CXICON);
        int cyIcon = GetSystemMetrics(SM_CYICON);
        CRect rect;
        GetClientRect(&rect);
        int x = (rect.Width() - cxIcon + 1) / 2;
        int y = (rect.Height() - cyIcon + 1) / 2;

        // Draw the icon
        dc.DrawIcon(x, y, m_hIcon);
    }
    else
    {
        CDialog::OnPaint();
    }
}

// The system calls this function to obtain the cursor to display while the user drags
// the minimized window.
HCURSOR CPEChecksum::OnQueryDragIcon()
{
    return static_cast<HCURSOR>(m_hIcon);
}

void CPEChecksum::OnBnClickedBrowse()
{
    CString szFile;
    CEdit* pEdit = NULL;

    if( FALSE == SelectFile( szFile ) ) { return; }

    ////////////////////////////////////////////////////
    // Set Filename
    pEdit = static_cast<CEdit*>
        ( GetDlgItem( IDC_EDIT_BROWSE ) );

    if( NULL == pEdit ) { return; }
    pEdit->SetWindowText( szFile );

    ////////////////////////////////////////////////////
    // Set Existing Checksum
    pEdit = static_cast<CEdit*>
        ( GetDlgItem( IDC_EDIT_CS_EXISTING ) );

    if( NULL == pEdit ) { return; }
    pEdit->SetWindowText( _T("") );

    ////////////////////////////////////////////////////
    // Set Calculated Checksum
    pEdit = static_cast<CEdit*>
        ( GetDlgItem( IDC_EDIT_CS_CALCULATED ) );

    if( NULL == pEdit ) { return; }
    pEdit->SetWindowText( _T("") );

    ////////////////////////////////////////////////////
    // Set Calculated CRC
    pEdit = static_cast<CEdit*>
        ( GetDlgItem( IDC_EDIT_NEW_VALUE ) );

    if( NULL == pEdit ) { return; }
    pEdit->SetWindowText( _T("0x00") );
}
BOOL CPEChecksum::RetrieveFilename( CString& szFilename )
{
    CEdit* pEdit = static_cast<CEdit*>
        ( GetDlgItem( IDC_EDIT_BROWSE ) );

    // Sanity Check
    if( NULL == pEdit ) { return FALSE; }

    pEdit->GetWindowText( szFilename );
    szFilename.Trim();

    return TRUE;
}

BOOL CPEChecksum::SelectFile( CString& szFile )
{
    BROWSEINFO bi = {0};
    ITEMIDLIST* list = NULL;
    TCHAR selection[ MAX_PATH+1 ] = {0};

    ZeroMemory(&bi, sizeof(bi));

    bi.pszDisplayName = selection;
    bi.lpszTitle = _T("Please select a file.");
    bi.ulFlags = BIF_BROWSEINCLUDEFILES /*| BIF_SHAREABLE | BIF_USENEWUI*/;

    // list is NULL if user cancelled Browse
    list = SHBrowseForFolder(&bi);

    // Sanity Check
    if( NULL == list ) { return FALSE; }

    // Extract the path/file from the list structure
    SHGetPathFromIDList(list, selection);

    if( TRUE == PathIsDirectory(selection)) { return FALSE; }

    szFile = selection;

    // Clean Up the Shell's Interface Pointer
    LPMALLOC pMem = NULL;
    if( SHGetMalloc( &pMem ) == NOERROR ) {
        pMem->Free( list );
    }
    pMem->Release();

    return TRUE;
}

void CPEChecksum::OnBnClickedRetrieve()
{
    CString szFilename;

    if( FALSE == RetrieveFilename( szFilename ) )
    { return; }

    // Sanity Check
    if( 0 == szFilename.GetLength() )
    { return; }

    DWORD32 dwExistingChecksum = 0, dwChecksum = 0, dwCRC = 0;

    if( TRUE == GenerateChecksum( szFilename, dwExistingChecksum, dwChecksum ) )
    {
        CString szFormatter;
        CEdit* pEdit = static_cast<CEdit*>
            ( GetDlgItem( IDC_EDIT_CS_EXISTING ) );

        if( NULL != pEdit )
        {
            szFormatter.Format( _T("0x%08X"), dwExistingChecksum );
            pEdit->SetWindowText( szFormatter );
        }

        pEdit = static_cast<CEdit*>
            ( GetDlgItem( IDC_EDIT_CS_CALCULATED ) );

        if( NULL != pEdit )
        {
            szFormatter.Format( _T("0x%08X"), dwChecksum );
            pEdit->SetWindowText( szFormatter );
        }
    }
}

BOOL CPEChecksum::RetrieveAndGenerateChecksum( const CString& szFilename,
                                              DWORD32& dwExistingChecksum,
                                              DWORD32& dwChecksum, DWORD32& dwCRC )
{
    if( FALSE == GenerateChecksum( szFilename,
        dwExistingChecksum, dwChecksum ) )
    {
        return FALSE;
    }

    return TRUE;
}

BOOL CPEChecksum::GenerateChecksum( const CString& szFilename,
                                   DWORD32& dwExistingChecksum,
                                   DWORD32& dwChecksum )
{
    HANDLE hFile = INVALID_HANDLE_VALUE;
    HANDLE hFileMapping = NULL;
    PVOID pBaseAddress = NULL;
    DWORD dwFileLength = 0;
    DWORD dwHeaderSum; // Checksum as stated by Header
    DWORD dwCheckSum; // Calculated Checksum

#ifdef _DEBUG
    CString szTracer;
#endif

    try
    {
        /////////////////////////////////////////////////////////////
        hFile = CreateFile( szFilename, GENERIC_READ, FILE_SHARE_READ,
            NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, 0 );
        if ( INVALID_HANDLE_VALUE == hFile ||
            NULL == hFile) { throw GetLastError(); }

#ifdef _DEBUG
        szTracer.Format( _T("File Handle: 0x%06X\n"), hFile );
        OutputDebugString( szTracer );
#endif

        /////////////////////////////////////////////////////////////
        hFileMapping = CreateFileMapping(hFile, NULL,
            PAGE_READONLY, 0, 0, NULL);
        if ( NULL == hFileMapping )
        {
            throw GetLastError();
        }

#ifdef _DEBUG
        szTracer.Format( _T("File Handle: 0x%06X\n"), hFile );
        OutputDebugString( szTracer );
#endif

        /////////////////////////////////////////////////////////////
        pBaseAddress = MapViewOfFile( hFileMapping,
            FILE_MAP_READ, 0, 0, 0);
        if ( NULL == pBaseAddress )
        {
            throw GetLastError();
        }

#ifdef _DEBUG
        szTracer.Format( _T("File Handle: 0x%06X\n"), hFile );
        OutputDebugString( szTracer );
#endif

        /////////////////////////////////////////////////////////////
        DWORD dwSize = 0;
        LARGE_INTEGER liSize = { 0, 0 };
        if( TRUE == GetFileSizeEx( hFile, &liSize ) )
        {
            dwSize = liSize.LowPart;
        }

#ifdef _DEBUG
        szTracer.Format( _T("File Handle: 0x%06X\n"), hFile );
        OutputDebugString( szTracer );
#endif

        SetLastError( ERROR_SUCCESS );

        /////////////////////////////////////////////////////////////
        PIMAGE_NT_HEADERS pNTHeaders = CheckSumMappedFile(
            pBaseAddress, dwSize, &dwHeaderSum, &dwCheckSum );

        if( NULL != pNTHeaders )
        {
            dwExistingChecksum = dwHeaderSum;
            dwChecksum = dwCheckSum;
        }
    }

    catch(...)
    {
        if( ERROR_SUCCESS != GetLastError() )
        {
            LONG lLastError = GetLastError();
            CString szFormatter;
            CString szTitle;
            AfxGetApp()->GetMainWnd()->GetWindowText( szTitle );

            szFormatter.Format( _T("Error retrieving Checksum.\n")
                _T("GetLastError() returned %d."), lLastError );

            MessageBox( szFormatter, szTitle, MB_OK | MB_ICONWARNING );
        }

        return FALSE;
    }

    /////////////////////////////////////////////////////////////
    UnmapViewOfFile( pBaseAddress );
    CloseHandle( hFile );

    return TRUE;
}

void CPEChecksum::OnBnClickedBtnWriteNewValue()
{
    CString szFilename;
    DWORD32 dwNewValue = 0;

    if( FALSE == RetrieveFilename( szFilename ) )
    {
        return;
    }

    if( FALSE == RetrieveNewValue( dwNewValue ) )
    {
        CString szTitle;
        AfxGetApp()->GetMainWnd()->GetWindowText( szTitle );

        MessageBox( _T("Unable to convert New Value."),
            szTitle, MB_OK | MB_ICONWARNING );

        return;
    }

    WriteChecksum( szFilename, dwNewValue );

    OnBnClickedRetrieve();
}

BOOL CPEChecksum::RetrieveNewValue( DWORD32& dwNewValue )
{
    CEdit* pEdit = static_cast<CEdit*>
        ( GetDlgItem( IDC_EDIT_NEW_VALUE ) );

    if( NULL == pEdit )
    {
        return FALSE;
    }

    CString szNewValue;
    pEdit->GetWindowText( szNewValue );
    szNewValue.Trim();

    // NULL or Empty = 0x00
    if( 0 == szNewValue.GetLength() )
    {
        szNewValue = _T("0x00");
    }

    // No Negatives...
    if( _T('-') == szNewValue[ 0 ] )
    {
        return FALSE;
    }

    // 4294967295 is Max value of a DWORD32
    // string (0xFFFFFFFF)
    if( 10 < szNewValue.GetLength() )
    {
        return FALSE;
    }

    // dwNewValue = _tstoi( szNewValue );
    if( 0 == _stscanf_s(szNewValue, _T("%X"), &dwNewValue ) )
    {
        return FALSE;
    }

    return TRUE;
}

BOOL CPEChecksum::WriteChecksum( const CString& szFilename, DWORD32& dwChecksum)
{
    HANDLE hFile = INVALID_HANDLE_VALUE;
    HANDLE hFileMapping = NULL;
    PVOID pBaseAddress = NULL;

    try {

        /////////////////////////////////////////////////////////////
        hFile = CreateFile( szFilename, FILE_READ_DATA | FILE_WRITE_DATA,
            0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, 0 );
        if ( INVALID_HANDLE_VALUE == hFile ||
            NULL == hFile) { throw GetLastError(); }

        /////////////////////////////////////////////////////////////
        hFileMapping = CreateFileMapping(hFile, NULL,
            PAGE_READWRITE, 0, 0, NULL);
        if ( NULL == hFileMapping )
        {
            throw GetLastError();
        }

        /////////////////////////////////////////////////////////////
        pBaseAddress = MapViewOfFile( hFileMapping,
            FILE_MAP_ALL_ACCESS, 0, 0, 0);
        if ( NULL == pBaseAddress )
        {
            throw GetLastError();
        }

        /////////////////////////////////////////////////////////////
        PIMAGE_DOS_HEADER pDOSHeader = NULL;
        pDOSHeader = static_cast<PIMAGE_DOS_HEADER>( pBaseAddress );
        if( pDOSHeader->e_magic != IMAGE_DOS_SIGNATURE )
        {
            throw GetLastError();
        }

        /////////////////////////////////////////////////////////////
        PIMAGE_NT_HEADERS pNTHeader = NULL;
        pNTHeader = reinterpret_cast<PIMAGE_NT_HEADERS>(
            (PBYTE)pBaseAddress + pDOSHeader->e_lfanew );

        if( pNTHeader->Signature != IMAGE_NT_SIGNATURE )
        {
            throw GetLastError();
        }

        /////////////////////////////////////////////////////////////
        SetLastError( ERROR_SUCCESS );

        /////////////////////////////////////////////////////////////
        DWORD* pChecksum = &(pNTHeader->OptionalHeader.CheckSum);
        *pChecksum = dwChecksum;
    }

    catch( ... )
    {
        if( ERROR_SUCCESS != GetLastError() )
        {
            LONG lLastError = GetLastError();
            CString szFormatter;
            CString szTitle;
            AfxGetApp()->GetMainWnd()->GetWindowText( szTitle );

            szFormatter.Format( _T("Error writing new Checksum.\n")
                _T("GetLastError() returned %d."), lLastError );

            MessageBox( szFormatter, szTitle, MB_OK | MB_ICONWARNING );
        }

        UnmapViewOfFile( hFileMapping );
        CloseHandle( hFile );

        return FALSE;
    }

    UnmapViewOfFile( hFileMapping );
    CloseHandle( hFile );

    return TRUE;
}