/*
PE Packer Decompression Benchmarker v2.0
(c)2000-2006 Bitsum Tehcnologies
by Jeremy Collake
support@bitsum.com
http://www.bitsum.com

Version 1.0: (2000)
Initial release.

Version 2.0: (01-2006)
Updated to use TCHARs.

Version 2.01
Modified to include module dependency load times.

Notes: This code was written some 6 years ago ...


Copyright and license notice: You MAY NOT use this code in your application (or copy it) without
permission from the author. I intend to allow for mostly free use of this code. 
By simply getting permission, you will save yourself from possible lawsuits. Email Jeremy Collake
if you intend to use any portion of this code.	
	

*/

#include<stdio.h>
#include<windows.h>
#include<winuser.h>
#include<imagehlp.h>
#include<winnt.h>
#include<tchar.h>
#include"resource.h"

#define BENCHMARK_XTIMES 100

unsigned int GetNamedFileSize(TCHAR *ptszFileName);
bool BackupFile(TCHAR *ptszFileName);
bool RestoreFile(TCHAR *ptszFileName);
bool PutBreakPointAtEntry(TCHAR *ptszAppName, DWORD *rvaEntry);
unsigned int TestExecutableLoadTime(TCHAR *ptszAppName, DWORD rvaEntry);
int CALLBACK DlgProc(HWND hwndDlg, UINT uMsg, WPARAM wParam, LPARAM lParam);
int CALLBACK DlgProcWaiting(HWND hwndDlg, UINT uMsg, WPARAM wParam, LPARAM lParam);

const TCHAR *g_ptszCaption=_T("PE Packer Decompression Benchmarker v2.01 alpha");
const TCHAR *g_ptszAuthor=_T("Jeremy Collake, jeremy@bitsum.com. http://www.bitsum.com");
HINSTANCE g_hInstance;
HWND g_hMain;
HWND g_hwndWaitingDlg;

unsigned int GetNamedFileSize(TCHAR *ptszFileName)
{
    HANDLE hFile;
    unsigned int nSize=0;
    hFile=CreateFile(ptszFileName,GENERIC_READ,FILE_SHARE_READ,0,OPEN_EXISTING,FILE_ATTRIBUTE_NORMAL,0);
    if(hFile==INVALID_HANDLE_VALUE)
    {
        MessageBox(g_hMain, _T("Error: Cannot obtain size of file.. is it here? is it write accessible at the moment?"), g_ptszCaption, MB_ICONSTOP);
        return 0;
    }
    nSize=GetFileSize(hFile,NULL);
    CloseHandle(hFile);
    return nSize;
}

bool BackupFile(TCHAR *ptszFileName)
{
    TCHAR tszBackup[MAX_PATH];
    _tcscpy_s(tszBackup,sizeof(tszBackup),ptszFileName);
    _tcscat_s(tszBackup,sizeof(tszBackup),_T(".backup"));
    return CopyFile(ptszFileName,tszBackup,FALSE);
}

bool RestoreFile(TCHAR *ptszFileName)
{
    TCHAR tszBackup[MAX_PATH];
    _tcscpy_s(tszBackup,sizeof(tszBackup),ptszFileName);
    _tcscat_s(tszBackup,sizeof(tszBackup),_T(".backup"));
    BOOL bR=CopyFile(tszBackup,ptszFileName,FALSE);
    if(bR)
    {
        DeleteFile(tszBackup);
    }
    return bR;
}


bool PutBreakPointAtEntry(TCHAR *ptszAppName, DWORD *rvaEntry)
{
    HANDLE hFile;
    HANDLE hMapping;
    LPVOID pView;
    PIMAGE_NT_HEADERS pNTHeader;
    PIMAGE_DOS_HEADER pDOS;
    PIMAGE_SECTION_HEADER pSectionHeader;
    unsigned char *pEntry;

    hFile=CreateFile(ptszAppName,GENERIC_READ|GENERIC_WRITE,FILE_SHARE_READ,0,OPEN_EXISTING,FILE_ATTRIBUTE_NORMAL,0);
    if(hFile==INVALID_HANDLE_VALUE)
    {
        MessageBox(g_hMain,_T("Error: Could not open file!"),g_ptszCaption,MB_ICONSTOP);
        return false;
    }
    hMapping=CreateFileMapping(hFile,0,PAGE_READWRITE,0,GetFileSize(hFile,NULL),NULL);
    if(hMapping==INVALID_HANDLE_VALUE)
    {
        MessageBox(g_hMain,_T("Error: Could not map file!"),g_ptszCaption,MB_ICONSTOP);
        CloseHandle(hFile);
        return false;
    }
    pView=MapViewOfFile(hMapping,FILE_MAP_ALL_ACCESS,0,0,0);
    if(!pView)
    {
        MessageBox(g_hMain,_T("Error: Could not create view of file!"), g_ptszCaption,MB_ICONSTOP);
        CloseHandle(hMapping);
        CloseHandle(hFile);
    }
    pDOS=(PIMAGE_DOS_HEADER)pView;
    pNTHeader=(PIMAGE_NT_HEADERS)(pDOS->e_lfanew+(DWORD)pDOS);
    if(pDOS->e_magic!=IMAGE_DOS_SIGNATURE
        ||
        IsBadReadPtr((void *)pNTHeader,sizeof(_IMAGE_OPTIONAL_HEADER))
        ||
        pNTHeader->Signature!=IMAGE_NT_SIGNATURE)
    {
        MessageBox(g_hMain,_T("Error: Not a valid portable executable!"),g_ptszCaption,MB_ICONSTOP);
        UnmapViewOfFile((LPCVOID)pView);
        CloseHandle(hMapping);
        CloseHandle(hFile);
        return false;
    }
    if(!pNTHeader->OptionalHeader.AddressOfEntryPoint)
    {
        MessageBox(g_hMain,_T("Error: No entry point."),g_ptszCaption,MB_ICONSTOP);
        UnmapViewOfFile((LPCVOID)pView);
        CloseHandle(hMapping);
        CloseHandle(hFile);
        return false;
    }

    pSectionHeader=IMAGE_FIRST_SECTION(pNTHeader);
    unsigned int nN=0;
    *rvaEntry=pNTHeader->OptionalHeader.AddressOfEntryPoint;
    while(pSectionHeader->VirtualAddress<=pNTHeader->OptionalHeader.AddressOfEntryPoint
        && nN<pNTHeader->FileHeader.NumberOfSections)
    {
        pSectionHeader++;
        nN++;
    }
    pSectionHeader--;
    pEntry=(unsigned char *)
        (DWORD)pNTHeader->OptionalHeader.AddressOfEntryPoint
        -(DWORD)pSectionHeader->VirtualAddress
        +(DWORD)pSectionHeader->PointerToRawData
        +(DWORD)pView;

    if(IsBadWritePtr(pEntry,3))
    {
        MessageBox(g_hMain,_T("Error: Could not find physical offset of entry point."),g_ptszCaption,MB_ICONSTOP);
        UnmapViewOfFile((LPCVOID)pView);
        CloseHandle(hMapping);
        CloseHandle(hFile);
        return false;
    }
    *pEntry=0xCC;       // finally, put the breakpoint in

    UnmapViewOfFile(pView);
    CloseHandle(hMapping);
    CloseHandle(hFile);
    return true;

}
unsigned int TestExecutableLoadTime(TCHAR *ptszAppName, DWORD rvaEntry)
{
    STARTUPINFO sinfo;
    PROCESS_INFORMATION pinfo;
    DEBUG_EVENT dinfo;
    unsigned int nStart=0, nElapsed=-2;
    DWORD byteswritten;
    unsigned char pExitProcessCode[8];  // we put code to exitprocess here
    FARPROC EP;
    HMODULE hKernel;

    LARGE_INTEGER lStart;
    LARGE_INTEGER lFinish;
    lStart.QuadPart=0;
    lFinish.QuadPart=0;
    // I went overboard here - but I figured it was cleaner to let the process terminate
    // itself. Therefore, we slap a jmp to ExitProcess in after the breakpoint, then
    // continue execution.
    hKernel=GetModuleHandle(_T("KERNEL32.DLL"));
    pExitProcessCode[0]=0xB8;   // mov eax,
    EP=GetProcAddress(hKernel,"ExitProcess");
    // fuck damned typecasting
    _asm
    {
        lea eax,pExitProcessCode[1]
        mov ecx,EP
            mov [eax],ecx
    };
    pExitProcessCode[5]=0xFF;   // jmp eax
    pExitProcessCode[6]=0xE0;   //

    GetStartupInfo(&sinfo);
    if(!CreateProcess(NULL,ptszAppName,0,0,false,DEBUG_ONLY_THIS_PROCESS,NULL,NULL,&sinfo,&pinfo))
    {
        return -1;
    }
    do
    {
        if(!WaitForDebugEvent(&dinfo,INFINITE))
        {
            break;
        }
        switch(dinfo.dwDebugEventCode)
        {
        case CREATE_PROCESS_DEBUG_EVENT:
            rvaEntry+=(DWORD)dinfo.u.CreateProcessInfo.lpBaseOfImage;
            CloseHandle(dinfo.u.CreateProcessInfo.hFile);
            QueryPerformanceCounter(&lStart);
            ContinueDebugEvent(dinfo.dwProcessId,dinfo.dwThreadId,DBG_CONTINUE);
            break;
        case EXIT_PROCESS_DEBUG_EVENT:
            CloseHandle(pinfo.hThread);
            CloseHandle(pinfo.hProcess);
            ContinueDebugEvent(dinfo.dwProcessId,dinfo.dwThreadId,DBG_CONTINUE);
            break;
        case EXCEPTION_DEBUG_EVENT:
            if(dinfo.u.Exception.ExceptionRecord.ExceptionCode==EXCEPTION_BREAKPOINT)
            {
                if(dinfo.u.Exception.ExceptionRecord.ExceptionAddress==(void *)rvaEntry)
                {
                    QueryPerformanceCounter(&lFinish);
                    //write jmp to ExitProcess
                    WriteProcessMemory(pinfo.hProcess,(BYTE *)dinfo.u.Exception.ExceptionRecord.ExceptionAddress+1,(LPVOID)&pExitProcessCode,7,&byteswritten);
                    ContinueDebugEvent(dinfo.dwProcessId,dinfo.dwThreadId,DBG_CONTINUE);
                    break;
                }
                ContinueDebugEvent(dinfo.dwProcessId,dinfo.dwThreadId,DBG_CONTINUE);
                break;
            }
        default:
            ContinueDebugEvent(dinfo.dwProcessId,dinfo.dwThreadId,DBG_EXCEPTION_NOT_HANDLED);
        }

    } while(dinfo.dwDebugEventCode!=EXIT_PROCESS_DEBUG_EVENT);

    if(lFinish.QuadPart)
    {
        LARGE_INTEGER lDifference;
        LARGE_INTEGER lFreq;
        // get counts per second
        QueryPerformanceFrequency(&lFreq);
        lDifference.QuadPart=lFinish.QuadPart-lStart.QuadPart;
        LARGE_INTEGER lElapsed;
        // get elpased time in ms
        lElapsed.QuadPart=(lDifference.QuadPart*1000)/lFreq.QuadPart;
        if(!lElapsed.HighPart)
        {
            nElapsed=lElapsed.LowPart;
        }
        else
        {
            MessageBox(g_hMain,_T("Error: Elapsed perfomance count exceeds resolution."),g_ptszCaption,MB_ICONSTOP);
            nElapsed=-2;
        }

    }
    else
    {
        nElapsed=-2;
    }
    return nElapsed;        // return elapsed time
}

int CALLBACK DlgProcWaiting(HWND hwndDlg, UINT uMsg, WPARAM wParam, LPARAM lParam)
{
    if(uMsg==WM_INITDIALOG)
    {
        g_hwndWaitingDlg=hwndDlg;
        SetDlgItemText(hwndDlg,IDC_WAITING,_T("Please wait... (and don't be messing with your mouse<g>)"));
        return 1;
    }
    else if(uMsg==WM_CLOSE)
    {
        EndDialog(hwndDlg,0);
        return 0;
    }
    return 0;
};

int CALLBACK DlgProc(HWND hwndDlg, UINT uMsg, WPARAM wParam, LPARAM lParam)
{
    HICON hIcon;
    TCHAR tszAppName[MAX_PATH];
    TCHAR tszms[MAX_PATH+20];
    static TCHAR *ptszFilter=_T("Portable Executables\0*.exe;*.dll;*.sys;*.ocx;*.bpl;*.scr\0All Files\0*.*\0\0");
    static OPENFILENAME ofn;
    unsigned int nResults[BENCHMARK_XTIMES];
    DWORD rvaEntry=0;
    unsigned int nSize;
    float fAvgMsToLoad;
    switch(uMsg)
    {
    case WM_INITDIALOG:
        SetWindowText(hwndDlg,g_ptszCaption);
        hIcon=LoadIcon(g_hInstance,MAKEINTRESOURCE(IDI_ICON1));
        SendMessage(hwndDlg,WM_SETICON,ICON_BIG,(DWORD)hIcon);
        SendMessage(hwndDlg,WM_SETICON,ICON_SMALL,(DWORD)hIcon);
        g_hMain=hwndDlg;
        return false;
    case WM_CLOSE:
    case WM_DESTROY:
        EndDialog(hwndDlg,0);
        ExitProcess(0);
        return true;
    case WM_COMMAND:
        switch(wParam)
        {
        case IDC_GO:
            if(!GetDlgItemText(hwndDlg,IDC_FILE,tszAppName,255))
            {
                MessageBox(hwndDlg,_T("Give me a file to benchmark first!"), g_ptszCaption, MB_ICONSTOP);
                break;
            }
            if(!(nSize=GetNamedFileSize(tszAppName)))
            {
                break;
            }
            BackupFile(tszAppName);
            if(PutBreakPointAtEntry(tszAppName,&rvaEntry))
            {
                if(MessageBox(hwndDlg,_T("Now, go compress the selected executable with the PE compressor of your choice. Select OK when done, or CANCEL to abort."),g_ptszCaption,MB_OKCANCEL)==IDOK)
                {
                    if(GetNamedFileSize(tszAppName)==nSize)
                    {
                        if(MessageBox(hwndDlg,_T("The size of the file hasn't changed! Is this correct?"),g_ptszCaption,MB_YESNO|MB_ICONQUESTION)==IDNO)
                        {
                            RestoreFile(tszAppName);
                            break;
                        }
                    }
                    unsigned int nMsToLoad=0;
                    EnableWindow(GetDlgItem(hwndDlg,IDC_GO),FALSE);
                    EnableWindow(GetDlgItem(hwndDlg,IDC_BROWSE),FALSE);
                    EnableWindow(GetDlgItem(hwndDlg,IDC_EXIT),FALSE);
                    ShowWindow(CreateDialog(g_hInstance,MAKEINTRESOURCE(IDD_DIALOG2),g_hMain,&DlgProcWaiting),SW_SHOW);
                    UpdateWindow(g_hwndWaitingDlg);
                    SetDlgItemInt(g_hwndWaitingDlg,IDC_TESTS,0,false);
                    SetDlgItemInt(g_hwndWaitingDlg,IDC_LASTRESULT,0,false);

                    do
                    {
                        // do once to get cached... stabilize...
                        TestExecutableLoadTime(tszAppName,rvaEntry);
                        // no do test loop..
                        unsigned int nN;
                        for(nN=0;nN<BENCHMARK_XTIMES;nN++)
                        {
                            SetDlgItemInt(g_hwndWaitingDlg,IDC_TESTS,nN+1,false);
                            if((nMsToLoad=TestExecutableLoadTime(tszAppName,rvaEntry))==-2)
                            {
                                MessageBox(hwndDlg,_T("ERROR: Executable crashed or breakpoint was removed!"),g_ptszCaption,MB_ICONSTOP);
                                break;
                            }
                            else if(nMsToLoad==-1)
                            {
                                MessageBox(hwndDlg,_T("Error: Cannot create process!"),g_ptszCaption,MB_ICONSTOP);
                                break;
                            }
                            nResults[nN]=nMsToLoad;
                            SetDlgItemInt(g_hwndWaitingDlg,IDC_LASTRESULT,nMsToLoad,false);
                        }
                        if(nN<BENCHMARK_XTIMES)
                        {
                            break;
                        }
                        for(fAvgMsToLoad=0,nN=0;nN<BENCHMARK_XTIMES;nN++)
                        {
                            fAvgMsToLoad+=nResults[nN];
                        }
                        fAvgMsToLoad/=BENCHMARK_XTIMES;
                        _stprintf(tszms,_T("It took %s an average of %.02f milliseconds to decompress/reconstruct in memory. Do you wish to benchmark this same program/compressor again?"),tszAppName,fAvgMsToLoad);

                    } while(MessageBox(hwndDlg,tszms,_T("Benchmark complete!"),MB_YESNO)==IDYES);
                    EnableWindow(GetDlgItem(hwndDlg,IDC_GO),TRUE);
                    EnableWindow(GetDlgItem(hwndDlg,IDC_BROWSE),TRUE);
                    EnableWindow(GetDlgItem(hwndDlg,IDC_EXIT),TRUE);
                }
                RestoreFile(tszAppName);
                PostMessage(g_hwndWaitingDlg,WM_CLOSE,0,0);
            }
            return true;
        case IDC_BROWSE:
            tszAppName[0]=0;
            memset(&ofn,0,sizeof(OPENFILENAME));
            ofn.lStructSize=sizeof(OPENFILENAME);
            ofn.hInstance=g_hInstance;
            ofn.hwndOwner=hwndDlg;
            ofn.lpstrFilter=ptszFilter;
            ofn.lpstrTitle=_T("Select portable executable to benchmark..");
            ofn.lpstrFile=tszAppName;
            ofn.nMaxFile=MAX_PATH;
            ofn.lpstrFileTitle=NULL;
            ofn.Flags=OFN_HIDEREADONLY;
            if(GetOpenFileName(&ofn))
            {
                SetDlgItemText(hwndDlg,IDC_FILE,tszAppName);
            }
            return true;
        case IDC_EXIT:
            EndDialog(hwndDlg,0);
            return true;
        }
        break;
    }
    return false;
}

int WINAPI WinMain(HINSTANCE hInst, HINSTANCE, LPSTR lpCmdLine, int)
{
    g_hInstance=hInst;
    return DialogBoxParam(hInst,MAKEINTRESOURCE(IDD_DIALOG1),0,&DlgProc,0);
}