/*
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 ...

*/

#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);
}