/*
    g++ main.cpp -o main.exe --std=c++11
*/
#include <windows.h>
#include <tlhelp32.h>
#include <iostream>
#include <vector>
#include <fstream>

using namespace std;

typedef vector<char> Memory;
typedef vector<Memory> Memories;

/*
    OpenProcessToken:
        BOOL WINAPI OpenProcessToken(
          _In_  HANDLE  ProcessHandle,
          _In_  DWORD   DesiredAccess,
          _Out_ PHANDLE TokenHandle
        );

        DesiredAccess:
            Options:
                TOKEN_ADJUST_DEFAULT
                TOKEN_ADJUST_GROUPS
                TOKEN_ADJUST_PRIVILEGES
                TOKEN_ADJUST_SESSIONID
                TOKEN_ASSIGN_PRIMARY
                TOKEN_DUPLICATE
                TOKEN_EXECUTE
                TOKEN_IMPERSONATE
                TOKEN_QUERY
                TOKEN_QUERY_SOURCE
                TOKEN_READ
                TOKEN_WRITE
                TOKEN_ALL_ACCESS

            Reference:
                https://docs.microsoft.com/en-us/windows/desktop/SecAuthZ/access-rights-for-access-token-objects

        Reference:
            https://msdn.microsoft.com/en-us/library/Aa379295(v=VS.85).aspx

    LookupPrivilegeValue:
        BOOL LookupPrivilegeValueA(
            LPCSTR lpSystemName,
            LPCSTR lpName,
            PLUID  lpLuid
        );

        Reference:
            https://docs.microsoft.com/en-us/windows/desktop/api/winbase/nf-winbase-lookupprivilegevaluea

    AdjustTokenPrivileges:
        BOOL WINAPI AdjustTokenPrivileges(
          _In_      HANDLE            TokenHandle,
          _In_      BOOL              DisableAllPrivileges,
          _In_opt_  PTOKEN_PRIVILEGES NewState,
          _In_      DWORD             BufferLength,
          _Out_opt_ PTOKEN_PRIVILEGES PreviousState,
          _Out_opt_ PDWORD            ReturnLength
        );

        Reference:
            https://msdn.microsoft.com/en-us/library/windows/desktop/aa375202(v=vs.85).aspx

    TOKEN_PRIVILEGES:
        PrivilegeCount:
            This must be set to the number of entries in the Privileges array.

        Privileges:
            Options:
                SE_PRIVILEGE_ENABLED
                SE_PRIVILEGE_ENABLED_BY_DEFAULT
                SE_PRIVILEGE_REMOVED
                SE_PRIVILEGE_USED_FOR_ACCESS

        Reference:
            https://docs.microsoft.com/en-us/windows/desktop/api/winnt/ns-winnt-_token_privileges

    MEMORY_BASIC_INFORMATION:
        BaseAddress

        AllocationBase

        AllocationProtect:
            The memory protection option when the region was initially allocated.
            This member can be one of the memory protection constants or 0 if the caller does not have access.

        RegionSize:

        State:
            Options:
                MEM_COMMIT
                MEM_FREE
                MEM_RESERVE

        Protect:
            The access protection of the pages in the region. This member is one of the values listed for the AllocationProtect member.

        Type:
            Options:
                MEM_IMAGE
                MEM_MAPPED
                MEM_PRIVATE

        Reference:
            https://docs.microsoft.com/en-us/windows/desktop/api/winnt/ns-winnt-_memory_basic_information
*/

string getLastErrorAsString() {
    //Get the error message, if any.
    DWORD errorMessageID = ::GetLastError();

    if(errorMessageID == 0)
        return std::string(); //No error message has been recorded

    LPSTR messageBuffer = nullptr;
    size_t size = FormatMessageA(
        FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
        NULL,
        errorMessageID,
        MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
        (LPSTR)&messageBuffer,
        0,
        NULL
    );

    string message(messageBuffer, size);

    //Free the buffer.
    LocalFree(messageBuffer);

    return message;
}

/*
    Enables Debug Privilages to the current process
*/
bool enableDebugPriv() {
    HANDLE hToken;
    LUID luid;
    TOKEN_PRIVILEGES tkp;

    if(OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken)) {
        if(LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &luid)) {
            tkp.PrivilegeCount = 1;
            tkp.Privileges[0].Luid = luid;
            tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

            if(AdjustTokenPrivileges(hToken, false, &tkp, sizeof(tkp), NULL, NULL) == 0)
                return false;
        }
        else
            return false;

        CloseHandle(hToken);
    }
    else
        return false;

    return true;
}

/*
    Collect all memory areas
*/
Memories getAllMemories(HANDLE process) {
    unsigned char *p;
    MEMORY_BASIC_INFORMATION info;
    Memory memory;
    Memories memories;

    #if defined(_M_X64) || defined(__amd64__)
        size_t bytes_read;
    #else
        DWORD bytes_read;
    #endif

    for(p = NULL; VirtualQueryEx(process, p, &info, sizeof(info)) == sizeof(info); p += info.RegionSize) {
        if(info.AllocationProtect >= PAGE_READONLY) {
            memory.resize(info.RegionSize);

            if(ReadProcessMemory(process, p, &memory[0], info.RegionSize, &bytes_read)) {
                memory.resize(bytes_read);

                memories.push_back(memory);
            }
        }
    }

    return memories;
}

/*
    Get Process by name
*/
HANDLE getProcessByName(string name) {
    HANDLE snapshot, process;
    PROCESSENTRY32 entry;

    // Patch entry size
    entry.dwSize = sizeof(PROCESSENTRY32);

    // Enable debug mode
    if(!enableDebugPriv())
        return NULL;

    // Create snapshot
    snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if(snapshot == INVALID_HANDLE_VALUE)
        return NULL;

    if(Process32First(snapshot, &entry) == TRUE) {
        // Loop Processes
        while(Process32Next(snapshot, &entry) == TRUE) {
            // Check executable name
            if(strcmp(entry.szExeFile, name.c_str()) == 0) {
                // Open process
                return OpenProcess(PROCESS_ALL_ACCESS, FALSE, entry.th32ProcessID);
            }
        }
    }

    CloseHandle(snapshot);

    return NULL;
}

void saveMemories(string process_name, Memories &m) {
    int i;
    string folder, file_path;
    ofstream outfile;

    //Init
    folder = "memories/" + process_name;

    // Create folder
    CreateDirectory(folder.c_str(), NULL);

    // Loop through memories
    for(i = 0; i < m.size(); i++) {
        // Generate file path
        file_path = folder + "/" + to_string(i) + ".data";

        // Open file
        outfile.open(file_path.c_str(), std::ofstream::binary);

        if(!outfile.fail()) {
            // Save data
            outfile.write((char*)&m[i][0], m[i].size());

            // Close file
            outfile.close();
        }
    }
}

int main() {
    string process_name;
    Memories m;
    HANDLE p;

    // Init
    //process_name = "main.exe";
    process_name = "explorer.exe";

    // Get process
    p = getProcessByName(process_name);
    if(p) {
        // Collect memories
        m = getAllMemories(p);

        // Save memories
        saveMemories(process_name, m);

        // Show total memory pages collected
        cout << "Saved " << m.size() << " memories" << endl;

        // Close process
        CloseHandle(p);
    }
    else
        cout << getLastErrorAsString() << endl;

    return 0;
}