From c641d0835af52719a80d4f2b689eb8ab5794997e Mon Sep 17 00:00:00 2001 From: retoor Date: Thu, 29 Jan 2026 00:51:24 +0100 Subject: [PATCH] Update. --- src/tool_registry.c | 88 +++++++++++++----------- src/tools/tool_network.c | 142 ++++++++++++++++++++++++++++----------- 2 files changed, 149 insertions(+), 81 deletions(-) diff --git a/src/tool_registry.c b/src/tool_registry.c index fe59113..597da82 100755 --- a/src/tool_registry.c +++ b/src/tool_registry.c @@ -4,6 +4,21 @@ #include #include #include +#include + +typedef struct { + tool_t *tool; + struct json_object *args; + char *output; +} tool_thread_args_t; + +static void *tool_thread_func(void *ptr) { + tool_thread_args_t *args = (tool_thread_args_t *)ptr; + if (args->tool->vtable->execute) { + args->output = args->tool->vtable->execute(args->tool, args->args); + } + return NULL; +} tool_registry_t *tool_registry_create(void) { tool_registry_t *registry = calloc(1, sizeof(tool_registry_t)); @@ -77,50 +92,34 @@ struct json_object *tool_registry_execute(tool_registry_t *registry, if (!results) return NULL; int len = json_object_array_length(tool_calls); + if (len == 0) return results; + + pthread_t *threads = calloc((size_t)len, sizeof(pthread_t)); + tool_thread_args_t *t_args = calloc((size_t)len, sizeof(tool_thread_args_t)); + struct json_object **result_objs = calloc((size_t)len, sizeof(struct json_object *)); for (int i = 0; i < len; i++) { struct json_object *call = json_object_array_get_idx(tool_calls, i); - struct json_object *result = json_object_new_object(); + result_objs[i] = json_object_new_object(); struct json_object *id_obj = json_object_object_get(call, "id"); if (id_obj) { - json_object_object_add(result, "tool_call_id", + json_object_object_add(result_objs[i], "tool_call_id", json_object_new_string(json_object_get_string(id_obj))); } - json_object_object_add(result, "role", json_object_new_string("tool")); - - struct json_object *type_obj; - if (json_object_object_get_ex(call, "type", &type_obj)) { - if (strcmp(json_object_get_string(type_obj), "function") != 0) { - json_object_put(result); - continue; - } - } + json_object_object_add(result_objs[i], "role", json_object_new_string("tool")); struct json_object *function_obj; if (!json_object_object_get_ex(call, "function", &function_obj)) { - json_object_object_add(result, "content", - json_object_new_string("Error: missing function object")); - json_object_array_add(results, result); + json_object_object_add(result_objs[i], "content", json_object_new_string("Error: missing function")); continue; } - struct json_object *name_obj; - if (!json_object_object_get_ex(function_obj, "name", &name_obj)) { - json_object_object_add(result, "content", - json_object_new_string("Error: missing function name")); - json_object_array_add(results, result); - continue; - } - - const char *name = json_object_get_string(name_obj); + const char *name = json_object_get_string(json_object_object_get(function_obj, "name")); tool_t *tool = tool_registry_find(registry, name); if (!tool) { - fprintf(stderr, "Unknown function: %s\n", name); - json_object_object_add(result, "content", - json_object_new_string("Error: function not found.")); - json_object_array_add(results, result); + json_object_object_add(result_objs[i], "content", json_object_new_string("Error: tool not found")); continue; } @@ -135,23 +134,30 @@ struct json_object *tool_registry_execute(tool_registry_t *registry, } if (verbose && args) { - fprintf(stderr, " \033[2m[verbose] %s args: %s\033[0m\n", - tool->name, json_object_to_json_string(args)); + fprintf(stderr, " \033[2m[parallel] launching %s\033[0m\n", tool->name); } - char *output = NULL; - if (tool->vtable->execute) { - output = tool->vtable->execute(tool, args); - } + t_args[i].tool = tool; + t_args[i].args = args; + t_args[i].output = NULL; - json_object_object_add(result, "content", - json_object_new_string(output ? output : "")); - - free(output); - if (args) json_object_put(args); - - json_object_array_add(results, result); + pthread_create(&threads[i], NULL, tool_thread_func, &t_args[i]); } + for (int i = 0; i < len; i++) { + if (threads[i]) { + pthread_join(threads[i], NULL); + } + json_object_object_add(result_objs[i], "content", + json_object_new_string(t_args[i].output ? t_args[i].output : "")); + free(t_args[i].output); + if (t_args[i].args) json_object_put(t_args[i].args); + json_object_array_add(results, result_objs[i]); + } + + free(threads); + free(t_args); + free(result_objs); + return results; -} +} \ No newline at end of file diff --git a/src/tools/tool_network.c b/src/tools/tool_network.c index 82d3fdd..e415a6b 100644 --- a/src/tools/tool_network.c +++ b/src/tools/tool_network.c @@ -37,6 +37,17 @@ static tool_t network_port_scan_tool = { .vtable = &network_port_scan_vtable, .n 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 +#include + +#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; @@ -53,55 +64,106 @@ static char *network_port_scan_execute(tool_t *self, struct json_object *args) { 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); + fprintf(stderr, " \033[1mParallel scanning %s ports %d to %d...\033[0m\n", host, start_port, end_port); - int total = end_port - start_port + 1; - int current = 0; + probe_t probes[MAX_PARALLEL_PROBES]; + memset(probes, 0, sizeof(probes)); - for (int port = start_port; port <= end_port; port++) { - current++; - fprintf(stderr, "\r -> [%d/%d] Probing port %d... ", current, total, port); - fflush(stderr); + int current_port = start_port; + int active_count = 0; - int sockfd = socket(AF_INET, SOCK_STREAM, 0); - if (sockfd < 0) continue; + 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 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)current_port); - 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); + connect(s, (struct sockaddr *)&addr, sizeof(addr)); - 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 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;