// retoor <retoor@molodetz.nl>

#include "tool.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; }

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[1mScanning %s ports %d to %d...\033[0m\n", host, start_port, end_port);

    int total = end_port - start_port + 1;
    int current = 0;

    for (int port = start_port; port <= end_port; port++) {
        current++;
        fprintf(stderr, "\r  -> [%d/%d] Probing port %d... ", current, total, port);
        fflush(stderr);

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

        struct timeval tv;
        tv.tv_sec = 0;
        tv.tv_usec = 150000; // 150ms for faster feedback
        setsockopt(sockfd, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv));

        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)port);

        if (connect(sockfd, (struct sockaddr *)&addr, sizeof(addr)) == 0) {
            fprintf(stderr, "\033[32mOPEN\033[0m   \n");
            struct json_object *entry = json_object_new_object();
            json_object_object_add(entry, "port", json_object_new_int(port));
            json_object_object_add(entry, "status", json_object_new_string("OPEN"));
            
            char banner[1024];
            memset(banner, 0, sizeof(banner));
            tv.tv_sec = 1;
            tv.tv_usec = 0;
            setsockopt(sockfd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
            ssize_t n = recv(sockfd, banner, sizeof(banner)-1, 0);
            if (n > 0) {
                // Clean banner for display
                for(int i=0; i<n; i++) if(banner[i] < 32) banner[i] = ' ';
                fprintf(stderr, "     \033[2mBanner: %s\033[0m\n", banner);
                json_object_object_add(entry, "banner", json_object_new_string(banner));
            }
            json_object_array_add(results, entry);
        }
        close(sockfd);
    }
    fprintf(stderr, "\r  \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_object_add(parameters, "required", required);
    json_object_object_add(function, "parameters", parameters);
    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_object_add(parameters, "required", required);
    json_object_object_add(function, "parameters", parameters);
    json_object_object_add(root, "function", function);
    return root;
}