feat: add upstream SSH tunnel targets and Python C API for connect/peek with protocol sniffing

sock.h

@@ -1,3 +1,6 @@
+#ifndef PGS_SOCK_H
+#define PGS_SOCK_H
+#include "protocol.h"
 #include <arpa/inet.h>
 #include <errno.h>
 #include <fcntl.h>
@@ -12,16 +15,19 @@
 #include <sys/socket.h>
 #include <sys/types.h>
 #include <unistd.h>
-
+#include "py.h"
 #define MAX_EVENTS 8096
 #define BUFFER_SIZE 1024
 
 typedef struct {
+  PROTOCOL_STATUS status;
+  PROTOCOL_NAME protocol_name;
   int client_fd;
   int upstream_fd;
   char *buffer;
   size_t buffer_size;
   size_t buffer_offset;
+  char *http_intercepted_headers;
 } connection_t;
 
 int listen_fd = 0;
@@ -124,7 +130,7 @@ char *sock_read(int fd, char *buf, size_t size) {
   size_t bytes_to_read = size > left_in_buffer ? left_in_buffer : size;
   ssize_t bytes_read = 0;
   char *buffer;
-  buffer[size];
+  buffer[size] = 0;
   if (bytes_to_read) {
     bytes_read = recv(fd, buffer, size, 0);
     buffer[bytes_read] = 0;
@@ -170,6 +176,18 @@ int forward_data(int from_fd, int to_fd) {
   return (int)bytes_read;
 }
 
+ssize_t sock_send_all(int fd, char * content, ssize_t length) {
+    ssize_t bytes_total_sent = 0;
+    while(bytes_total_sent < length) {
+      ssize_t bytes_sent = send(fd, content + bytes_total_sent, length, 0);
+      if(bytes_sent <= 0){
+        return -1;
+      }
+      bytes_total_sent += bytes_sent; 
+    }
+    return bytes_total_sent;
+}
+
 bool handle_connect(struct epoll_event event, int epoll_fd) {
   struct sockaddr_in client_addr;
   socklen_t client_len = sizeof(client_addr);
@@ -187,7 +205,8 @@ bool handle_connect(struct epoll_event event, int epoll_fd) {
   client_event.data.fd = client_fd;
   connections[client_fd]->upstream_fd = -1;
   connections[client_fd]->client_fd = client_fd;
-
+  connections[client_fd]->status = PS_SNIFF;
+  connections[client_fd]->protocol_name = PS_NONE;
   printf("New connection: client_fd=%d\n", client_fd);
 
   epoll_ctl(epoll_fd, EPOLL_CTL_ADD, client_fd, &client_event);
@@ -200,34 +219,103 @@ void handle_close(int epoll_fd, connection_t *conn) {
   close_connection(epoll_fd, conn);
 }
 
+char * sock_read_until(int fd, char * until, int until_len) {
+
+  char * result = (char *)malloc(BUFFER_SIZE);
+  result[0] = 0;
+  int bytes;
+  int bytes_total = 0;
+  char buffer[2];
+  while((bytes = recv(fd, buffer, 1,0)) > 0){
+    result[bytes_total] = buffer[0];
+    bytes_total += bytes;
+    if(bytes_total >= until_len){
+      if(!strncmp(buffer + bytes_total - until_len, until, until_len)){
+        break;
+      }
+    }
+  }
+  result[bytes_total] = 0;
+  return result;
+}
+
+void http_intercept_headers(int epoll_fd, connection_t *conn) {
+  char * headers = sock_read_until(conn->client_fd, "\r\n\r\n", 4);
+
+  char * intercepted_headers = py_http_intercept_headers(conn->client_fd,headers);
+  if(!intercepted_headers){
+    close_connection(epoll_fd, conn);
+    return;
+  } 
+  sock_send_all(conn->client_fd, intercepted_headers, strlen(intercepted_headers));
+  printf("Intercepted headers: %s\n", intercepted_headers);
+  conn->http_intercepted_headers = intercepted_headers;
+  conn->status = PS_NONE;
+  connections[conn->upstream_fd]->status = PS_HTTP_READ_HEADER;
+}
+
+
+void http_intercept_body(int epoll_fd, connection_t *conn) {
+   printf("HIERR\n"); 
+}
+
 void handle_stream(struct epoll_event event, int epoll_fd, connection_t *conn) {
   if (conn->upstream_fd == -1) {
-    conn->upstream_fd = prepare_upstream();
-    int upstream_fd = py_route(conn->client_fd, conn->upstream_fd);
-
-    if (upstream_fd == -1) {
+    conn->upstream_fd = py_on_connect(conn->client_fd);
+    if(conn->upstream_fd == -1){
       close_connection(epoll_fd, conn);
       return;
     }
-    set_nonblocking(upstream_fd);
+    conn->protocol_name = protocol_sniff(conn->client_fd);
+    if(conn->protocol_name == PN_ERROR){
+      close_connection(epoll_fd, conn);
+      return;
+    }else if(conn->protocol_name == PN_HTTP){
+      conn->status = PS_HTTP_READ_HEADER;
+    }else if(conn->protocol_name == PN_HTTP_CHUNKED){
+      conn->status = PS_HTTP_READ_HEADER;
+    }else if(conn->protocol_name == PN_HTTP_KEEP_ALIVE){
+      conn->status = PS_HTTP_READ_HEADER;
+    }else if(conn->protocol_name == PN_HTTP_REQUEST){
+      conn->status = PS_HTTP_READ_HEADER;
+    }else if(conn->protocol_name == PN_HTTP_WEBSOCKET){
+      conn->status = PS_HTTP_READ_HEADER;
+    }else if(conn->protocol_name == PN_SSH){
+      conn->status = PS_STREAM;
+    }else if(conn->protocol_name == PN_RAW){
+      conn->status = PS_STREAM;
+    }else {
+      conn->status = PS_ERROR;
+    }
+
+    if (conn->upstream_fd == -1 || conn->status == PS_ERROR) {
+      close_connection(epoll_fd, conn);
+      return;
+    }
+    set_nonblocking(conn->upstream_fd);
     struct epoll_event upstream_event;
     upstream_event.events = EPOLLIN | EPOLLOUT | EPOLLERR | EPOLLHUP;
-    upstream_event.data.ptr = connections[upstream_fd];
-    upstream_event.data.fd = upstream_fd;
+    upstream_event.data.ptr = connections[conn->upstream_fd];
+    upstream_event.data.fd = conn->upstream_fd;
+    connections[conn->client_fd]->upstream_fd = conn->upstream_fd;
+    connections[conn->upstream_fd]->client_fd = conn->client_fd;
+    connections[conn->upstream_fd]->upstream_fd = conn->upstream_fd;
 
-    connections[conn->client_fd]->upstream_fd = upstream_fd;
-    connections[upstream_fd]->client_fd = conn->client_fd;
-    connections[upstream_fd]->upstream_fd = upstream_fd;
-
-    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, upstream_fd, &upstream_event);
+    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, conn->upstream_fd, &upstream_event);
 
     printf("Connected: client_fd=%d, upstream_fd=%d\n", conn->client_fd,
            conn->upstream_fd);
+
     return;
   }
 
   if (event.data.fd == conn->client_fd) {
-
+    if(conn->status = PS_HTTP_READ_HEADER){
+      http_intercept_headers(epoll_fd, conn);
+    }
+    if(conn->status == PS_HTTP_READ_BODY){
+      http_intercept_body(epoll_fd, conn);
+    }
     if (forward_data(conn->client_fd, conn->upstream_fd) < 1) {
       close_connection(epoll_fd, conn);
     }
@@ -290,4 +378,5 @@ void serve(int port) {
       }
     }
   }
-}
+}
+#endif