#include <windows.h>
#include <stdio.h>
#include"child_process.h"

void dispatch_child_event(DEBUG_EVENT const &event, child_process const &child) {
    char *file_name;
    char buffer[512];

    switch ( event.dwDebugEventCode ) {
    case LOAD_DLL_DEBUG_EVENT:
        file_name = child.get_string(event.u.LoadDll.lpImageName);
        if ( event.u.LoadDll.fUnicode)
            printf("Loading %S\
", (wchar_t *)file_name);
        else
            printf("Loading %s\
", file_name);            

        break;

    case EXCEPTION_DEBUG_EVENT:
        switch (event.u.Exception.ExceptionRecord.ExceptionCode)
        {
            case EXCEPTION_ACCESS_VIOLATION:
            {                
                if ( event.u.Exception.dwFirstChance)
                    break;
                EXCEPTION_RECORD const &r = event.u.Exception.ExceptionRecord;
                printf("Access Violation %x at %0#p\
",
                        r.ExceptionCode,
                        r.ExceptionAddress);
                break;
            }

            case EXCEPTION_BREAKPOINT:
                printf("Breakpoint reached\
");
                break;

            case EXCEPTION_DATATYPE_MISALIGNMENT:
                if ( !event.u.Exception.dwFirstChance)
                    printf("Misaligned data exception.\
");
                break;

            case EXCEPTION_SINGLE_STEP:
                printf("Single Step...\
");
                break;

            case DBG_CONTROL_C:
                if ( !event.u.Exception.dwFirstChance)
                    printf("Control+C pressed\
");
                break;    
            break;
        }

    case CREATE_THREAD_DEBUG_EVENT:
        printf("Client created a thread\
");
        break;

    case CREATE_PROCESS_DEBUG_EVENT:
        printf("Create-Process\
");
        break;

    case EXIT_THREAD_DEBUG_EVENT:
        printf("Thread exited.\
");
        break;

    case UNLOAD_DLL_DEBUG_EVENT:
        printf("DLL being unloaded\
");
        break;

    case OUTPUT_DEBUG_STRING_EVENT: {
        OUTPUT_DEBUG_STRING_INFO const &d = event.u.DebugString;
        char *string = child.get_debug_string(d.lpDebugStringData,
                                        d.nDebugStringLength);
        if ( d.fUnicode)
            printf("Debug string: %S\
", string);
        else
            printf("Debug string: %s\
", string);
        break;
    }
    }
}

int main(int argc, char **argv) {

    DEBUG_EVENT event;

    if ( argc < 2 ) {
        fprintf(stderr,"Usage: Trace [executable|PID]");
        return EXIT_FAILURE;
    }

    child_process child(argv[1]);

    do {
        WaitForDebugEvent(&event, INFINITE);

        dispatch_child_event(event, child);

        ContinueDebugEvent( event.dwProcessId,
                            event.dwThreadId,
                            DBG_CONTINUE );

    } while ( event.dwDebugEventCode != EXIT_PROCESS_DEBUG_EVENT);

    return 0;
}

#ifndef CHILD_PROCESS_H_INC_
#define CHILD_PROCESS_H_INC_

#include <windows.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include
#include <io.h>

#include"syserror.h"

struct no_spawn {
    no_spawn() { system_error("Spawning Program"); }
};

class child_process {
    HANDLE process_;
    HANDLE thread_;

    mutable char buffer[FILENAME_MAX];

public:
    child_process(char const *filename);

    char *get_string(void *string_name, DWORD num = 0) const;
    char *get_debug_string(void *string, DWORD num) const;  

    HANDLE process() { return process_; }
    HANDLE thread() { return thread_; }

    ~child_process() { CloseHandle(process()); }
};

#endif

#include"child_process.h"

static BOOL find_image(char const *name, char *buffer) {
    // Try to find an image file named by the user.
    // First search for the exact file name in the current
    // directory.  If that's not found, look for same base name
    // with".com",".exe" and".bat" appended, in that order.
    // If we can't find it in the current directory, repeat
    // the entire process on directories specified in the
    // PATH environment variable.
    //
#define elements(array) (sizeof(array)/sizeof(array[0]))

    static char *extensions[] = {".com",".exe",".bat",".cmd"};
    int i;
    char temp[FILENAME_MAX];

    if (-1 != _access(name, 0)) {
        strcpy(buffer, name);
        return TRUE;
    }

    for (i=0; i<elements(extensions); i++) {
        strcpy(temp, name);
        strcat(temp, extensions[i]);
        if ( -1 != _access(temp, 0)) {
            strcpy(buffer, temp);
            return TRUE;
        }
    }

    _searchenv(name,"PATH", buffer);
    if ( buffer[0] != '\\0')
        return TRUE;

    for ( i=0; i<elements(extensions); i++) {
        strcpy(temp, name);
        strcat(temp, extensions[i]);
        _searchenv(temp,"PATH", buffer);
        if ( buffer[0] != '\\0')
            return TRUE;
    }

    return FALSE;
}

child_process::child_process(char const *filename) {
    if (isdigit(filename[0])) {
        DWORD id = atoi(filename);
        process_ = OpenProcess(PROCESS_ALL_ACCESS, false, atoi(filename));
        DebugActiveProcess(id);
    }
    else {
        char buf[FILENAME_MAX];

        PROCESS_INFORMATION pi = {0};
        STARTUPINFO si = {0};
        si.cb = sizeof(si);

        if (!find_image(filename, buf))
            throw no_spawn();

        BOOL new_process_ = CreateProcess(buf, NULL, NULL, NULL, FALSE,
            DEBUG_ONLY_THIS_PROCESS,
            NULL, NULL,
            &si, &pi);

        if (!new_process_)
            throw no_spawn();

        CloseHandle(pi.hThread);
        process_ = pi.hProcess;
        thread_ = pi.hThread;
    }
}

char *child_process::get_string(void *string_name, DWORD num) const {
// string_name is a pointer to a pointer to a string, with the pointer and the
// string itself located in another process_.  We use Readprocess_Memory to read
// the first pointer, then the string itself into our process_ address space.  
// We then return a pointer (in our address space) to the string we read in.
//
    char *ptr;
    SIZE_T bytes_read;

    if ( 0 == num )
        num = sizeof(buffer);

    if ( string_name == NULL )
        return NULL;

    ReadProcessMemory(process_,
        string_name,
        &ptr,
        sizeof(ptr),
        &bytes_read);

    if (NULL == ptr )
        return NULL;

    ReadProcessMemory(process_,
        ptr,
        buffer,
        num,
        &bytes_read);

    return buffer;
}

char *child_process::get_debug_string(void *string, DWORD num) const {

    static char buffer[FILENAME_MAX];
    SIZE_T bytes_read;

    if ( string == NULL )
        return NULL;

    ReadProcessMemory(process_,
        string,
        buffer,
        num,
        &bytes_read);
    return buffer;
}