Update.
This commit is contained in:
parent
df124bb10b
commit
c641d0835a
@ -4,6 +4,21 @@
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
#include <stdio.h>
|
||||
#include <pthread.h>
|
||||
|
||||
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;
|
||||
|
||||
pthread_create(&threads[i], NULL, tool_thread_func, &t_args[i]);
|
||||
}
|
||||
|
||||
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);
|
||||
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;
|
||||
}
|
||||
@ -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 <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;
|
||||
@ -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[1mParallel scanning %s ports %d to %d...\033[0m\n", host, start_port, end_port);
|
||||
|
||||
fprintf(stderr, " \033[1mScanning %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 total = end_port - start_port + 1;
|
||||
int current = 0;
|
||||
int current_port = start_port;
|
||||
int active_count = 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));
|
||||
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)port);
|
||||
addr.sin_port = htons((unsigned short)current_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"));
|
||||
connect(s, (struct sockaddr *)&addr, sizeof(addr));
|
||||
|
||||
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));
|
||||
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);
|
||||
}
|
||||
close(sockfd);
|
||||
done = true;
|
||||
} else {
|
||||
// Check if timeout for this specific probe?
|
||||
// Simplified: if select says nothing for 1s total, we'll close it eventually.
|
||||
}
|
||||
fprintf(stderr, "\r \033[32mScan complete.\033[0m \n");
|
||||
|
||||
// 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;
|
||||
|
||||
Loading…
Reference in New Issue
Block a user