Listing Processes and Disecting a Process

#include <string>
#include <fstream>
#include <sstream>
#include <vector>
#include <windows.h>
#include <tlhelp32.h>
using namespace std;

//bad naming, but life goes on
#define RN "\r\n"
#define ID_BUTTON    1
#define ID_BUTTON_R  2
#define ID_BUTTON_HEAP 3
#define ID_MULTILINE 100
#define ID_PID       101
#define ID_LOG       102

string l_to_string(const unsigned long &val)
{
    stringstream ss;
    ss << val;
    return(ss.str());
}

long  string_to_long(string &val)
{
    istringstream iss(val);
    long retl;
    iss >> retl;
    return(retl);
}

class ProcessView
{
    private:
            vector<PROCESSENTRY32>       processes;
            vector<IMAGE_SECTION_HEADER> sections;
            IMAGE_NT_HEADERS header;
            MODULEENTRY32    current;
            string           info, name;
            int              pid,  pe_offset;

            void thread_info(THREADENTRY32 &thread);
            void module_info(MODULEENTRY32 &module);

            bool get_heaps();
            bool get_threads();
            bool get_modules();
            bool dump_process();

            void parse_header();
            void get_sec_heads();
            int get_sections();

    public:
            bool get_sys_snap();
            bool get_process(string &saveto, int _pid, bool getheap);

            void view_processes(HWND hWnd);
};

void ProcessView::thread_info(THREADENTRY32 &thread)
{
    if(thread.th32OwnerProcessID == pid)
        info += "Thread ID: " + l_to_string(thread.th32ThreadID) + RN
             +  "  Size: " + l_to_string(thread.dwSize) + RN
             +  "  Base Priority: " + l_to_string(thread.tpBasePri) + RN + RN;
}

void ProcessView::module_info(MODULEENTRY32 &module)
{
    info += "Name: ";
      for(int i = 0; i < MAX_MODULE_NAME32 + 1; i++) {
        if(module.szModule[i] == '\0')
          break;
        info += module.szModule[i]; }

      info += "\r\n";
      info += "  Size: " + l_to_string(module.dwSize) + RN
           +  "  Base Size: " + l_to_string(module.modBaseSize) + RN
           +  "  Path: ";
      for(int i = 0; i < MAX_PATH; i++) {
        if(module.szExePath[i] == '\0')
          break;
        info += module.szExePath[i]; }
      info += "\r\n\r\n";
}

bool ProcessView::get_heaps()
{
    HEAPLIST32  heaps;
    HEAPENTRY32 heap;
    HANDLE      hHeaps;

    hHeaps   = CreateToolhelp32Snapshot(TH32CS_SNAPHEAPLIST, pid);
      if(!hHeaps)
        return(false);

    heap.dwSize  = sizeof(HEAPENTRY32);
    heaps.dwSize  = sizeof(HEAPLIST32);

    if(!Heap32ListFirst(hHeaps, &heaps)) {
      CloseHandle(hHeaps);
      return(false); }

    for(;;)
    {
      if(Heap32First(&heap, pid, heaps.th32HeapID))
      {
        info += "Heap List " + l_to_string(heaps.th32HeapID) + ":" + RN;
        do {
          if(heap.dwFlags != LF32_FREE) {
            info += "  Block Size: " + l_to_string(heap.dwBlockSize) + RN; }
        } while(Heap32Next(&heap));
      }
      else
        info += "  Free.\r\n";
      if(!Heap32ListNext(hHeaps, &heaps)) {
        info += RN;
        break; }
    }

    CloseHandle(hHeaps);
    return(true);
}

bool ProcessView::get_threads()
{
    THREADENTRY32 thread;
    HANDLE        hThreads;

    hThreads = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
      if(!hThreads)
        return(false);

    thread.dwSize = sizeof(THREADENTRY32);

    if(!Thread32First(hThreads, &thread)) {
      CloseHandle(hThreads);
      return(false); }

    thread_info(thread);
    while(Thread32Next(hThreads, &thread))
      thread_info(thread);

    CloseHandle(hThreads);
    return(true);
}

bool ProcessView::get_modules()
{
    HANDLE hMods;
    MODULEENTRY32 module;

    hMods    = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, pid);
      if(!hMods)
        return(false);

    module.dwSize = sizeof(MODULEENTRY32);

    if(!Module32First(hMods, &module)) {
      CloseHandle(hMods);
      return(false); }

    module_info(module);
    current = module;
    for(int i = 0; i < MAX_PATH; i++) {
        if(module.szExePath[i] == '\0')
          break;
        name += module.szExePath[i]; }

    while(Module32Next(hMods, &module))
      module_info(module);

    CloseHandle(hMods);
    return(true);
}

bool ProcessView::dump_process()
{
    IMAGE_IMPORT_DESCRIPTOR imp_des;
    IMAGE_DOS_HEADER dos;   //no one cares, except for the 0x3C
    DWORD            read;

    if(!Toolhelp32ReadProcessMemory(pid, (void*)current.modBaseAddr, &dos, sizeof(dos), &read))
      return(false);
    if(read != sizeof(dos)) //why, why, why?!
      return(false);

    if(!Toolhelp32ReadProcessMemory(pid, (void*)(current.modBaseAddr + dos.e_lfanew), &header, sizeof(header), &read))
      return(false);
    if(read != sizeof(header)) //deep breaths
      return(false);

    pe_offset = dos.e_lfanew;
    parse_header();
    get_sec_heads();

    /*get_imports(); removed due to laziness*/

    get_sections();
    return(true);
}

void ProcessView::parse_header()
{
    info += "PE Header crap: \r\n";
    if(header.OptionalHeader.MajorImageVersion >0)
      info += "  Version: " + l_to_string(header.OptionalHeader.MajorImageVersion)
           +  "." + l_to_string(header.OptionalHeader.MinorImageVersion) + RN;
    if(header.OptionalHeader.MajorLinkerVersion > 0)
      info += "  Linker Version: " + l_to_string(header.OptionalHeader.MajorLinkerVersion)
           +  "." + l_to_string(header.OptionalHeader.MinorLinkerVersion) + RN;
    if(header.OptionalHeader.MajorOperatingSystemVersion > 0)
      info += "  OS Version: " + l_to_string(header.OptionalHeader.MajorOperatingSystemVersion)
           +  "." + l_to_string(header.OptionalHeader.MinorOperatingSystemVersion) + RN;
    if(header.OptionalHeader.MajorSubsystemVersion > 0)
      info += "  Subsystem Version: " + l_to_string(header.OptionalHeader.MajorSubsystemVersion)
           +  "." + l_to_string(header.OptionalHeader.MinorSubsystemVersion) + RN;

    info += "  Size: " + l_to_string(current.modBaseSize)
         +  "\r\n  Time Stamp: " + l_to_string(header.FileHeader.TimeDateStamp)
         +  "\r\n  n-Sections: " + l_to_string(header.FileHeader.NumberOfSections)
         +  "\r\n  n-Symbols: "  + l_to_string(header.FileHeader.NumberOfSymbols) + RN + RN;
}

void ProcessView::get_sec_heads()
{
    IMAGE_SECTION_HEADER sec;
    DWORD address, read;

    address = (unsigned long)current.modBaseAddr + pe_offset + sizeof(header);

    for(int i = 0; i < header.FileHeader.NumberOfSections; i++)
    {
      if(!Toolhelp32ReadProcessMemory(pid, (void*)address, &sec, sizeof(sec), &read))
        break;
      if(read != 40)
        break;

      info += "Section: ";
      for(int j = 0; j < 8; j++) {
        if(sec.Name[j] == '\0')
          break;
        info += sec.Name[j]; }

      info += "\r\n  Size: " + l_to_string(sec.SizeOfRawData)
           +  "\r\n  @: " + l_to_string(sec.PointerToRawData) + RN;

      if(sec.Characteristics & IMAGE_SCN_CNT_CODE)
        info += "  Contains code.\r\n";
      if(sec.Characteristics & IMAGE_SCN_MEM_EXECUTE)
        info += "  Is executable.\r\n";
      if(sec.Characteristics & IMAGE_SCN_CNT_INITIALIZED_DATA)
        info += "  Contains initialized data.\r\n";
      if(sec.Characteristics & IMAGE_SCN_CNT_UNINITIALIZED_DATA)
        info += "  Contains uninitialized data.\r\n";

      sections.push_back(sec);
      address += 40;
    }
}

int ProcessView::get_sections()
{
    /*removed due to laziness*/
}

bool ProcessView::get_sys_snap()
{
    string temp;
    PROCESSENTRY32 entry;
    HANDLE hOne, hAll;

    processes.clear(); //on refresh

    hAll = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if(hAll == INVALID_HANDLE_VALUE)
      return(false);

    entry.dwSize = sizeof(PROCESSENTRY32);
    if(!Process32First(hAll, &entry)) {
      CloseHandle(hOne);
      CloseHandle(hAll);
      return(false); }

    while(Process32Next(hAll, &entry))
    {
      hOne = OpenProcess(PROCESS_VM_READ, 0, entry.th32ProcessID);

      if(hOne == INVALID_HANDLE_VALUE)
        return(false);

      processes.push_back(entry);
    }

    CloseHandle(hOne);
    CloseHandle(hAll);
    return(true);
}

bool ProcessView::get_process(string &saveto, int _pid, bool getheap)
{
    SYSTEMTIME st;
    pid = _pid;

    GetSystemTime(&st);
    info += l_to_string(st.wHour) + ":" + l_to_string(st.wMinute)
         + "  " + l_to_string(st.wDay) + "\\" + l_to_string(st.wMonth)
         + "\\" + l_to_string(st.wYear) + RN + RN;

    if(getheap)
      if(!get_heaps())
        info += "Could not get heap list.\r\n";
    if(!get_threads())
      info += "Could not get thread list.\r\n";
    if(!get_modules())
      info += "Could not get module list. :(\r\n";
    if(!dump_process())
      info += "Could not dump process! :'(\r\n";

    ofstream file(saveto.c_str());

      if(!file.is_open()) {
        return(false); }
      for(int i = 0; i < info.size(); i++)
        file.put(info[i]);

    file.close();
    name.clear();
    info.clear();
    return(true);
}

void ProcessView::view_processes(HWND hWnd)
{
    string temp;
    int regs[4];

    /*untested...*/
    asm volatile("cpuid": "=a" (regs[0]), "=b" (regs[1]),
                 "=c" (regs[2]), "=d" (regs[3]): "a" (0x80000001));
    if(regs[3] & (1 << 29))
      temp += "\"Long Mode\" Detected!\r\n"\
              "Warning: This program is for 32bit apps.\r\n\r\n\r\n";

    for(int i = 0; i < processes.size(); i++)
    {
      for(int j = 0; j < MAX_PATH; j++) {
        if(processes[i].szExeFile[j] == '\0')
          break;
        temp += processes[i].szExeFile[j]; }

      temp += "\r\nPID: ";
      temp += l_to_string(processes[i].th32ProcessID) + "\r\n";

      temp += "Parent PID: ";
      temp += l_to_string(processes[i].th32ParentProcessID) + "\r\n";

      temp += "Base Priority: ";
      temp += l_to_string(processes[i].pcPriClassBase) + "\r\n";

      temp += "n-Threads: ";
      temp += l_to_string(processes[i].cntThreads) + "\r\n\r\n";
    }
    SetWindowText(hWnd, temp.c_str());
}

ProcessView pview;
HWND        hBox, hPID, hLog;
HFONT       hBFont;
HINSTANCE   hInt;
bool        getheap = false;

LRESULT CALLBACK WndProc(HWND, UINT, WPARAM, LPARAM);

int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLine, int nCmdShow)
{
    WNDCLASSEX WndClass;
    HWND hWnd;
    MSG  Msg;

    hInt = hInstance;

    WndClass.cbSize        = sizeof(WNDCLASSEX);
    WndClass.style         = CS_HREDRAW | CS_VREDRAW;
    WndClass.lpfnWndProc   = WndProc;
    WndClass.cbClsExtra    = 0;
    WndClass.cbWndExtra    = 0;
    WndClass.hInstance     = hInstance;
    WndClass.hIcon         = LoadIcon(NULL, IDI_APPLICATION);
    WndClass.hCursor       = LoadCursor(NULL, IDC_ARROW);
    WndClass.hbrBackground = GetSysColorBrush(COLOR_3DFACE);
    WndClass.lpszMenuName  = NULL;
    WndClass.lpszClassName = "ProcessViewer";
    WndClass.hIconSm       = LoadIcon(NULL, IDI_APPLICATION);

    if(!RegisterClassEx(&WndClass)) {
      MessageBox(NULL, "Error: Registering window class.", NULL, MB_ICONSTOP | MB_OK);
      return(1); }

    hWnd = CreateWindowEx(WS_EX_STATICEDGE, "ProcessViewer", "ProcessView",
               WS_OVERLAPPEDWINDOW | WS_MINIMIZEBOX, CW_USEDEFAULT, CW_USEDEFAULT,
               310, 240, NULL, NULL, hInstance, NULL);

    if(!hWnd) {
      MessageBox(0, "Window Creation Failed!", "Error!", MB_ICONSTOP | MB_OK);
      return(1); }

    ShowWindow(hWnd, nCmdShow);
    UpdateWindow(hWnd);

    while(GetMessage(&Msg, NULL, 0, 0))
    {
      TranslateMessage(&Msg);
      DispatchMessage(&Msg);
    }
    return(Msg.wParam);
}

LRESULT CALLBACK WndProc(HWND hWnd, UINT Msg, WPARAM wParam, LPARAM lParam)
{
    switch(Msg)
    {
      case WM_CREATE:
        if(!pview.get_sys_snap()) {
          MessageBox(NULL, "Couldn't view current process images!", NULL, MB_OK);
          PostQuitMessage(0); }

        hBox = CreateWindowEx(WS_EX_CLIENTEDGE, "EDIT", "",
            WS_CHILD | WS_VISIBLE | ES_MULTILINE| WS_VSCROLL | WS_HSCROLL,
            100, 0, 200, 220, hWnd, reinterpret_cast<HMENU>(ID_MULTILINE), hInt, NULL);

        CreateWindowEx(WS_EX_CLIENTEDGE, "BUTTON", "Disect Image",
            WS_CHILD | WS_VISIBLE | BS_DEFPUSHBUTTON,
            0, 0, 100, 50, hWnd, reinterpret_cast<HMENU>(ID_BUTTON), hInt, NULL);

        hPID = CreateWindowEx(WS_EX_CLIENTEDGE, "EDIT", "",
            WS_CHILD | WS_VISIBLE,
            0, 55, 100, 22, hWnd, reinterpret_cast<HMENU>(ID_PID), hInt, NULL);
        SetWindowText(hPID, "PID goes here");

        hLog = CreateWindowEx(WS_EX_CLIENTEDGE, "EDIT", "",
            WS_CHILD | WS_VISIBLE,
            0, 78, 100, 22, hWnd, reinterpret_cast<HMENU>(ID_LOG), hInt, NULL);
        SetWindowText(hLog, "Save where?");

        CreateWindowEx(WS_EX_CLIENTEDGE, "BUTTON", "Include Heap?",
            WS_CHILD | WS_VISIBLE | BS_AUTOCHECKBOX | BS_LEFTTEXT,
            0, 102, 100, 22, hWnd, reinterpret_cast<HMENU>(ID_BUTTON_HEAP), hInt, NULL);

        CreateWindowEx(WS_EX_CLIENTEDGE, "BUTTON", "Refresh",
            WS_CHILD | WS_VISIBLE | BS_DEFPUSHBUTTON,
            0, 165, 100, 50, hWnd, reinterpret_cast<HMENU>(ID_BUTTON_R), hInt, NULL);

        hBFont = (HFONT__*)GetStockObject(DEFAULT_GUI_FONT);
        SendMessage(hBox, WM_SETFONT, (WPARAM)hBFont, MAKELPARAM(FALSE, 0));
        pview.view_processes(hBox);

        break;
      case WM_COMMAND:

        if((LOWORD(wParam) == ID_BUTTON) && (HIWORD(wParam) == BN_CLICKED))
        {
          string pid_str, log_str;
          char   line[22];
          SendMessage(hPID, EM_GETLINE, 0, (LPARAM)&line);
          pid_str = line;
          SendMessage(hLog, EM_GETLINE, 0, (LPARAM)&line);
          log_str = line;

          if(pid_str.find_first_not_of("0123456789") != string::npos) {
            MessageBox(NULL, "PID can only be numbers.", NULL, MB_OK);
            break; }

          if(log_str.find_first_of("\\/:*?<>|") != string::npos) {
            MessageBox(NULL, "File name cannot contain: \\/:*?<>|", NULL, MB_OK);
            break; }

          if(!pview.get_process(log_str, string_to_long(pid_str), getheap)) {
            MessageBox(NULL, "Problem disecting process.", NULL, MB_OK);
            break; }

          ShellExecute(NULL, "open", log_str.c_str(), NULL, NULL, SW_SHOWNORMAL);
          break;
        }
        else if((LOWORD(wParam) == ID_BUTTON_R) && (HIWORD(wParam) == BN_CLICKED))
        {
          if(!pview.get_sys_snap()) {
            MessageBox(NULL, "Couldn't get snapshot of processes!", NULL, MB_OK);
            break; }
          pview.view_processes(hBox);
          break;
        }

        getheap = IsDlgButtonChecked(hWnd, ID_BUTTON_HEAP);
        break;
      case WM_CLOSE:
        DestroyWindow(hWnd);
        break;
      case WM_DESTROY:
        PostQuitMessage(0);
        break;
      default:
        return DefWindowProc(hWnd, Msg, wParam, lParam);
    }
    return(0);
}