// retoor <retoor@molodetz.nl>

#include "tool.h"
#include "r_config.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <errno.h>

static struct json_object *network_check_get_description(void);
static char *network_check_execute(tool_t *self, struct json_object *args);
static void network_check_print_action(const char *name, struct json_object *args);

static struct json_object *network_port_scan_get_description(void);
static char *network_port_scan_execute(tool_t *self, struct json_object *args);

static const tool_vtable_t network_check_vtable = {
    .get_description = network_check_get_description,
    .execute = network_check_execute,
    .print_action = network_check_print_action
};

static const tool_vtable_t network_port_scan_vtable = {
    .get_description = network_port_scan_get_description,
    .execute = network_port_scan_execute,
    .print_action = NULL
};

static tool_t network_check_tool = { .vtable = &network_check_vtable, .name = "network_check" };
static tool_t network_port_scan_tool = { .vtable = &network_port_scan_vtable, .name = "network_port_scan" };

tool_t *tool_network_check_create(void) { return &network_check_tool; }
tool_t *tool_network_port_scan_create(void) { return &network_port_scan_tool; }

#include <fcntl.h>
#include <errno.h>

#define MAX_PARALLEL_PROBES 50

typedef struct {
    int sockfd;
    int port;
    bool in_use;
} probe_t;

static char *network_port_scan_execute(tool_t *self, struct json_object *args) {
    (void)self;
    struct json_object *host_obj = NULL, *start_obj = NULL, *end_obj = NULL;
    if (!json_object_object_get_ex(args, "host", &host_obj)) return strdup("Error: host missing");
    
    json_object_object_get_ex(args, "start_port", &start_obj);
    json_object_object_get_ex(args, "end_port", &end_obj);

    const char *host = json_object_get_string(host_obj);
    int start_port = start_obj ? json_object_get_int(start_obj) : 0;
    int end_port = end_obj ? json_object_get_int(end_obj) : start_port;

    struct hostent *server = gethostbyname(host);
    if (!server) return strdup("Error: dns resolution failed");

    struct json_object *results = json_object_new_array();
    fprintf(stderr, "  \033[1mParallel scanning %s ports %d to %d...\033[0m\n", host, start_port, end_port);

    probe_t probes[MAX_PARALLEL_PROBES];
    memset(probes, 0, sizeof(probes));

    int current_port = start_port;
    int active_count = 0;

    while (current_port <= end_port || active_count > 0) {
        // Fill slots
        while (active_count < MAX_PARALLEL_PROBES && current_port <= end_port) {
            int s = socket(AF_INET, SOCK_STREAM, 0);
            if (s >= 0) {
                int flags = fcntl(s, F_GETFL, 0);
                fcntl(s, F_SETFL, flags | O_NONBLOCK);

                struct sockaddr_in addr;
                memset(&addr, 0, sizeof(addr));
                addr.sin_family = AF_INET;
                memcpy(&addr.sin_addr.s_addr, server->h_addr, (size_t)server->h_length);
                addr.sin_port = htons((unsigned short)current_port);

                connect(s, (struct sockaddr *)&addr, sizeof(addr));

                for (int i = 0; i < MAX_PARALLEL_PROBES; i++) {
                    if (!probes[i].in_use) {
                        probes[i].sockfd = s;
                        probes[i].port = current_port;
                        probes[i].in_use = true;
                        active_count++;
                        break;
                    }
                }
            }
            current_port++;
        }

        if (active_count > 0) {
            fd_set writefds;
            FD_ZERO(&writefds);
            int max_fd = 0;
            for (int i = 0; i < MAX_PARALLEL_PROBES; i++) {
                if (probes[i].in_use) {
                    FD_SET(probes[i].sockfd, &writefds);
                    if (probes[i].sockfd > max_fd) max_fd = probes[i].sockfd;
                }
            }

            struct timeval tv = {0, 100000}; // 100ms
            int sel = select(max_fd + 1, NULL, &writefds, NULL, &tv);

            if (sel >= 0) {
                for (int i = 0; i < MAX_PARALLEL_PROBES; i++) {
                    if (probes[i].in_use) {
                        bool done = false;
                        if (FD_ISSET(probes[i].sockfd, &writefds)) {
                            int error = 0;
                            socklen_t len = sizeof(error);
                            getsockopt(probes[i].sockfd, SOL_SOCKET, SO_ERROR, &error, &len);
                            if (error == 0) {
                                fprintf(stderr, "  -> Port %d: \033[32mOPEN\033[0m\n", probes[i].port);
                                struct json_object *entry = json_object_new_object();
                                json_object_object_add(entry, "port", json_object_new_int(probes[i].port));
                                json_object_object_add(entry, "status", json_object_new_string("OPEN"));
                                json_object_array_add(results, entry);
                            }
                            done = true;
                        } else {
                            // Check if timeout for this specific probe? 
                            // Simplified: if select says nothing for 1s total, we'll close it eventually.
                        }

                        // Close sockets that stay too long? 
                        // For brevity, we'll just use a shorter overall timeout loop or check time.
                        // Here we'll just close if select reported it.
                        if (done) {
                            close(probes[i].sockfd);
                            probes[i].in_use = false;
                            active_count--;
                        }
                    }
                }
            }
            
            // Cleanup stalled probes (rudimentary timeout)
            static int timeout_counter = 0;
            if (++timeout_counter > 20) { // ~2 seconds
                for (int i = 0; i < MAX_PARALLEL_PROBES; i++) {
                    if (probes[i].in_use) {
                        close(probes[i].sockfd);
                        probes[i].in_use = false;
                        active_count--;
                    }
                }
                timeout_counter = 0;
            }
        }
    }

    fprintf(stderr, "  \033[32mScan 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 struct json_object *network_port_scan_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("network_port_scan"));
    json_object_object_add(function, "description", json_object_new_string("Scans a range of TCP ports on a host and attempts service banner grabbing."));
    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 *start = json_object_new_object();
    json_object_object_add(start, "type", json_object_new_string("integer"));
    json_object_object_add(properties, "start_port", start);
    struct json_object *end = json_object_new_object();
    json_object_object_add(end, "type", json_object_new_string("integer"));
    json_object_object_add(properties, "end_port", end);
    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("start_port"));
    json_object_array_add(required, json_object_new_string("end_port"));
    json_object_object_add(parameters, "required", required);
    json_object_object_add(parameters, "additionalProperties", json_object_new_boolean(0));

    json_object_object_add(function, "parameters", parameters);

    r_config_handle cfg = r_config_get_instance();
    if (r_config_use_strict(cfg)) {
        json_object_object_add(function, "strict", json_object_new_boolean(1));
    }

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

static void network_check_print_action(const char *name, struct json_object *args) {
    (void)name;
    struct json_object *host;
    if (json_object_object_get_ex(args, "host", &host)) {
        fprintf(stderr, "  -> Checking network: %s\n", json_object_get_string(host));
    }
}

static char *network_check_execute(tool_t *self, struct json_object *args) {
    (void)self;
    struct json_object *host_obj, *port_obj;
    if (!json_object_object_get_ex(args, "host", &host_obj)) {
        return strdup("Error: missing 'host' argument");
    }
    const char *host = json_object_get_string(host_obj);
    int port = 0;
    if (json_object_object_get_ex(args, "port", &port_obj)) {
        port = json_object_get_int(port_obj);
    }

    struct addrinfo hints, *res;
    memset(&hints, 0, sizeof(hints));
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;

    if (getaddrinfo(host, NULL, &hints, &res) != 0) {
        return strdup("Error: DNS lookup failed");
    }

    char ipstr[INET6_ADDRSTRLEN];
    void *addr;
    if (res->ai_family == AF_INET) {
        struct sockaddr_in *ipv4 = (struct sockaddr_in *)res->ai_addr;
        addr = &(ipv4->sin_addr);
    } else {
        struct sockaddr_in6 *ipv6 = (struct sockaddr_in6 *)res->ai_addr;
        addr = &(ipv6->sin6_addr);
    }
    inet_ntop(res->ai_family, addr, ipstr, sizeof(ipstr));

    char result[1024];
    if (port > 0) {
        int sockfd = socket(res->ai_family, res->ai_socktype, res->ai_protocol);
        if (sockfd < 0) {
            freeaddrinfo(res);
            return strdup("Error: could not create socket");
        }

        // Set non-blocking for timeout
        int flags = fcntl(sockfd, F_GETFL, 0);
        fcntl(sockfd, F_SETFL, flags | O_NONBLOCK);

        struct sockaddr_in serv_addr;
        if (res->ai_family == AF_INET) {
            memcpy(&serv_addr, res->ai_addr, sizeof(serv_addr));
            serv_addr.sin_port = htons(port);
        }

        int conn_res = connect(sockfd, (struct sockaddr *)&serv_addr, sizeof(serv_addr));
        if (conn_res < 0) {
            if (errno == EINPROGRESS) {
                fd_set set;
                struct timeval tv;
                FD_ZERO(&set);
                FD_SET(sockfd, &set);
                tv.tv_sec = 2; // 2 second timeout
                tv.tv_usec = 0;
                int select_res = select(sockfd + 1, NULL, &set, NULL, &tv);
                if (select_res > 0) {
                    int valopt;
                    socklen_t lon = sizeof(int);
                    getsockopt(sockfd, SOL_SOCKET, SO_ERROR, (void*)(&valopt), &lon);
                    if (valopt == 0) conn_res = 0;
                }
            }
        }

        close(sockfd);
        snprintf(result, sizeof(result), "Host: %s (IP: %s), Port %d: %s", 
                 host, ipstr, port, (conn_res == 0) ? "OPEN" : "CLOSED/TIMEOUT");
    } else {
        snprintf(result, sizeof(result), "Host: %s (IP: %s), DNS: OK", host, ipstr);
    }

    freeaddrinfo(res);
    return strdup(result);
}

static struct json_object *network_check_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("network_check"));
    json_object_object_add(function, "description", json_object_new_string("Check DNS and port connectivity for a host."));
    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(host, "description", json_object_new_string("Hostname or IP to check."));
    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(port, "description", json_object_new_string("Port to check (optional)."));
    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(parameters, "additionalProperties", json_object_new_boolean(0));

    json_object_object_add(function, "parameters", parameters);

    r_config_handle cfg = r_config_get_instance();
    if (r_config_use_strict(cfg)) {
        json_object_object_add(function, "strict", json_object_new_boolean(1));
    }

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