/*
 * WebSocket SSH (wsssh) - C Implementation
 * SSH wrapper that uses WebSocket tunnels through wssshd.
 *
 * Copyright (C) 2024 Stefy Lanza <stefy@nexlab.net> and SexHack.me
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <openssl/ssl.h>
#include <openssl/err.h>
#include <getopt.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <pthread.h>
#include <sys/select.h>

#define BUFFER_SIZE 1048576
#define DEFAULT_PORT 22

typedef struct {
    char *local_port;
    int debug;
} wsssh_config_t;

typedef struct {
    int local_sock;  // Local TCP connection socket
    char request_id[37];  // UUID string
    int active;
    SSL *ssl;  // WebSocket SSL connection
} tunnel_t;

typedef struct {
    SSL *ssl;
    int debug;
} thread_args_t;

tunnel_t *active_tunnel = NULL;
pthread_mutex_t tunnel_mutex;

int send_websocket_frame(SSL *ssl, const char *data);
int send_pong_frame(SSL *ssl, const char *ping_payload, int payload_len);

int send_pong_frame(SSL *ssl, const char *ping_payload, int payload_len) {
    char frame[BUFFER_SIZE];
    frame[0] = 0x8A; // FIN + pong opcode
    int header_len = 2;

    if (payload_len <= 125) {
        frame[1] = 0x80 | payload_len; // MASK + length
    } else if (payload_len <= 65535) {
        frame[1] = 0x80 | 126; // MASK + extended length
        frame[2] = (payload_len >> 8) & 0xFF;
        frame[3] = payload_len & 0xFF;
        header_len = 4;
    } else {
        frame[1] = 0x80 | 127; // MASK + extended length
        frame[2] = 0;
        frame[3] = 0;
        frame[4] = 0;
        frame[5] = 0;
        frame[6] = (payload_len >> 24) & 0xFF;
        frame[7] = (payload_len >> 16) & 0xFF;
        frame[8] = (payload_len >> 8) & 0xFF;
        frame[9] = payload_len & 0xFF;
        header_len = 10;
    }

    // Add mask key
    char mask_key[4];
    for (int i = 0; i < 4; i++) {
        mask_key[i] = rand() % 256;
        frame[header_len + i] = mask_key[i];
    }
    header_len += 4;

    // Mask payload
    for (int i = 0; i < payload_len; i++) {
        frame[header_len + i] = ping_payload[i] ^ mask_key[i % 4];
    }

    int frame_len = header_len + payload_len;
    return SSL_write(ssl, frame, frame_len) > 0;
}

void print_trans_flag(void) {
    // Transgender pride flag colors using ANSI escape codes
    const char *colors[] = {
        "\033[48;5;117m",  // Light blue background
        "\033[48;5;218m",  // Pink background
        "\033[48;5;231m",  // White background
        "\033[48;5;218m",  // Pink background
        "\033[48;5;117m"   // Light blue background
    };
    const char *reset = "\033[0m";

    // Print 5 rows of colored blocks
    for (int i = 0; i < 5; i++) {
        printf("%s", colors[i]);
        for (int j = 0; j < 40; j++) {
            printf(" ");
        }
        printf("%s\n", reset);
    }
    printf("\n");
}

void print_usage(const char *program_name) {
    fprintf(stderr, "Usage: %s [options] user@client.domain [ssh_options...]\n", program_name);
    fprintf(stderr, "WebSocket SSH Wrapper - SSH through WebSocket tunnels\n\n");
    fprintf(stderr, "Protect the dolls!\n\n");
    fprintf(stderr, "Options:\n");
    fprintf(stderr, "  --local-port PORT    Local tunnel port (default: auto)\n");
    fprintf(stderr, "  --debug              Enable debug output\n");
    fprintf(stderr, "  --help               Show this help\n");
    fprintf(stderr, "\nDonations:\n");
    fprintf(stderr, "  BTC: bc1q3zlkpu95amtcltsk85y0eacyzzk29v68tgc5hx\n");
    fprintf(stderr, "  ETH: 0xdA6dAb526515b5cb556d20269207D43fcc760E51\n");
}

int parse_args(int argc, char *argv[], wsssh_config_t *config, int *remaining_argc, char ***remaining_argv) {
    // Manually parse arguments to separate wsssh options from SSH options
    int ssh_start = 1; // Skip argv[0]

    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "--local-port") == 0 && i + 1 < argc) {
            config->local_port = strdup(argv[i + 1]);
            i++; // Skip the argument
        } else if (strcmp(argv[i], "--debug") == 0) {
            config->debug = 1;
        } else if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0) {
            print_usage(argv[0]);
            return 0;
        } else if (argv[i][0] == '-') {
            // Unknown option, treat as SSH option
            ssh_start = i;
            break;
        } else {
            // Non-option argument, start of SSH args
            ssh_start = i;
            break;
        }
    }

    // Return remaining arguments (SSH arguments)
    *remaining_argc = argc - ssh_start;
    *remaining_argv = &argv[ssh_start];

    return 1;
}

int parse_hostname(const char *hostname, char **client_id, char **wssshd_host) {
    char *at_pos = strchr(hostname, '@');
    if (!at_pos) {
        fprintf(stderr, "Error: Invalid hostname format. Expected user@host\n");
        return 0;
    }

    char *host_part = at_pos + 1;
    char *host_copy = strdup(host_part);
    if (!host_copy) {
        return 0;
    }

    // Remove port if present
    char *colon_pos = strchr(host_copy, ':');
    if (colon_pos) {
        *colon_pos = '\0';
    }

    // Split host by dots to extract client_id
    char *dot_pos = strchr(host_copy, '.');
    if (!dot_pos) {
        fprintf(stderr, "Error: Invalid hostname format. Expected client.domain format\n");
        free(host_copy);
        return 0;
    }

    *dot_pos = '\0';
    *client_id = strdup(host_copy);
    *wssshd_host = strdup(dot_pos + 1);
    free(host_copy);

    return 1;
}

int parse_ssh_args(int argc, char *argv[], char **host, int *ssh_port, int debug) {
    *host = NULL;
    *ssh_port = DEFAULT_PORT;

    int skip_next = 0;
    for (int i = 0; i < argc; i++) {
        if (skip_next) {
            skip_next = 0;
            continue;
        }

        if ((strcmp(argv[i], "-p") == 0 || strcmp(argv[i], "--port") == 0) && i + 1 < argc) {
            *ssh_port = atoi(argv[i + 1]);
            skip_next = 1;
            if (debug) {
                printf("[DEBUG] Found SSH port: %d\n", *ssh_port);
            }
        } else if (argv[i][0] != '-' && !*host) {
            // First non-option argument should be the host
            *host = argv[i];
            if (debug) {
                printf("[DEBUG] Found SSH host: %s\n", *host);
            }
        }
    }

    if (!*host) {
        fprintf(stderr, "Error: Could not determine target host from SSH arguments\n");
        return 0;
    }

    return 1;
}

char **modify_ssh_args(int argc, char *argv[], const char *original_host, int local_port, int *new_argc) {
    // Allocate space for: ssh + -o + StrictHostKeyChecking=no + -o + UserKnownHostsFile=/dev/null + -o + LogLevel ERROR + -p + port + original args + NULL
    char **new_args = malloc((argc + 10) * sizeof(char *));
    if (!new_args) {
        return NULL;
    }

    int idx = 0;
    new_args[idx++] = "ssh";

    // Add StrictHostKeyChecking=no option
    new_args[idx++] = "-o";
    new_args[idx++] = "StrictHostKeyChecking=no";

    // Add UserKnownHostsFile=/dev/null option
    new_args[idx++] = "-o";
    new_args[idx++] = "UserKnownHostsFile=/dev/null";

    // Add LogLevel ERROR option
    new_args[idx++] = "-o";
    new_args[idx++] = "LogLevel ERROR";

    // Add port argument for local tunnel
    new_args[idx++] = "-p";
    char *port_str = malloc(16);
    if (!port_str) {
        free(new_args);
        return NULL;
    }
    snprintf(port_str, 16, "%d", local_port);
    new_args[idx++] = port_str;

    int skip_next = 0;
    for (int i = 0; i < argc; i++) {
        if (skip_next) {
            skip_next = 0;
            continue;
        }

        if (strcmp(argv[i], "-p") == 0 || strcmp(argv[i], "--port") == 0 ||
            strcmp(argv[i], "-h") == 0 || strcmp(argv[i], "--host") == 0) {
            skip_next = 1;
            continue;
        }

        // Check if this argument is the original host
        if (strcmp(argv[i], original_host) == 0) {
            // Replace user@host with user@localhost or host with localhost
            char *at_pos = strchr(original_host, '@');
            if (at_pos) {
                // Has user part
                size_t user_len = at_pos - original_host;
                char *new_host = malloc(user_len + 10); // "@localhost"
                if (!new_host) {
                    free(port_str);
                    free(new_args);
                    return NULL;
                }
                memcpy(new_host, original_host, user_len);
                strcpy(new_host + user_len, "@localhost");
                new_args[idx++] = new_host;
            } else {
                // No user part
                new_args[idx++] = "localhost";
            }
            continue;
        }

        new_args[idx++] = argv[i];
    }

    new_args[idx] = NULL;
    *new_argc = idx;
    return new_args;
}

int find_available_port() {
    struct sockaddr_in addr;
    int sock;
    int port = 0;

    if ((sock = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
        perror("Socket creation failed");
        return 0;
    }

    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
    addr.sin_port = 0;

    if (bind(sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
        perror("Bind failed");
        close(sock);
        return 0;
    }

    socklen_t len = sizeof(addr);
    if (getsockname(sock, (struct sockaddr *)&addr, &len) < 0) {
        perror("Getsockname failed");
        close(sock);
        return 0;
    }

    port = ntohs(addr.sin_port);
    close(sock);

    return port;
}

void generate_request_id(char *request_id, size_t size) {
    const char charset[] = "0123456789abcdef";
    for (size_t i = 0; i < size - 1; i++) {
        request_id[i] = charset[rand() % (sizeof(charset) - 1)];
    }
    request_id[size - 1] = '\0';
}

int websocket_handshake(SSL *ssl, const char *host, int port, const char *path) {
    char request[1024];
    char response[BUFFER_SIZE];
    int bytes_read;

    // Send WebSocket handshake
    snprintf(request, sizeof(request),
             "GET %s HTTP/1.1\r\n"
             "Host: %s:%d\r\n"
             "Upgrade: websocket\r\n"
             "Connection: upgrade\r\n"
             "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n"
             "Sec-WebSocket-Version: 13\r\n"
             "\r\n",
             path, host, port);

    if (SSL_write(ssl, request, strlen(request)) <= 0) {
        ERR_print_errors_fp(stderr);
        fprintf(stderr, "WebSocket handshake send failed\n");
        return 0;
    }

    // Read response
    bytes_read = SSL_read(ssl, response, sizeof(response) - 1);
    if (bytes_read <= 0) {
        ERR_print_errors_fp(stderr);
        fprintf(stderr, "WebSocket handshake recv failed\n");
        return 0;
    }

    response[bytes_read] = '\0';

    // Check for successful handshake
    if (strstr(response, "101 Switching Protocols") == NULL) {
        fprintf(stderr, "WebSocket handshake failed\n");
        return 0;
    }

    return 1;
}

int send_json_message(SSL *ssl, const char *type, const char *client_id, const char *request_id) {
    char message[1024];

    if (request_id) {
        snprintf(message, sizeof(message),
                 "{\"type\":\"%s\",\"client_id\":\"%s\",\"request_id\":\"%s\"}",
                 type, client_id, request_id);
    } else {
        snprintf(message, sizeof(message),
                 "{\"type\":\"%s\",\"client_id\":\"%s\"}",
                 type, client_id);
    }

    // Send as WebSocket frame
    return send_websocket_frame(ssl, message);
}

int send_all(SSL *ssl, const char *data, int len) {
    int total_sent = 0;
    while (total_sent < len) {
        int sent = SSL_write(ssl, data + total_sent, len - total_sent);
        if (sent <= 0) {
            ERR_print_errors_fp(stderr);
            return 0;
        }
        total_sent += sent;
    }
    return 1;
}

int send_websocket_frame(SSL *ssl, const char *data) {
    int msg_len = strlen(data);
    int header_len = 2;

    if (msg_len <= 125) {
        header_len = 6; // 2 + 4 for mask
    } else if (msg_len <= 65535) {
        header_len = 8; // 4 + 4 for mask
    } else {
        header_len = 14; // 10 + 4 for mask
    }

    char *frame = malloc(header_len + msg_len);
    if (!frame) {
        perror("Memory allocation failed for frame");
        return 0;
    }

    frame[0] = 0x81; // FIN + text opcode

    if (msg_len <= 125) {
        frame[1] = 0x80 | msg_len; // MASK + length
    } else if (msg_len <= 65535) {
        frame[1] = 0x80 | 126; // MASK + extended length
        frame[2] = (msg_len >> 8) & 0xFF;
        frame[3] = msg_len & 0xFF;
    } else {
        frame[1] = 0x80 | 127; // MASK + extended length
        frame[2] = 0;
        frame[3] = 0;
        frame[4] = 0;
        frame[5] = 0;
        frame[6] = (msg_len >> 24) & 0xFF;
        frame[7] = (msg_len >> 16) & 0xFF;
        frame[8] = (msg_len >> 8) & 0xFF;
        frame[9] = msg_len & 0xFF;
    }

    // Add mask key
    char mask_key[4];
    for (int i = 0; i < 4; i++) {
        mask_key[i] = rand() % 256;
        frame[header_len - 4 + i] = mask_key[i];
    }

    // Mask payload
    for (int i = 0; i < msg_len; i++) {
        frame[header_len + i] = data[i] ^ mask_key[i % 4];
    }

    int frame_len = header_len + msg_len;
    int result = send_all(ssl, frame, frame_len);
    free(frame);
    return result;
}

int parse_websocket_frame(const char *buffer, int bytes_read, char **payload, int *payload_len) {
    if (bytes_read < 2 || ((buffer[0] & 0x8F) != 0x81 && (buffer[0] & 0x8F) != 0x82)) { // Not a text or binary frame
        return 0;
    }

    int masked = buffer[1] & 0x80;
    int len_indicator = buffer[1] & 0x7F;
    int header_len = 2;

    if (len_indicator <= 125) {
        *payload_len = len_indicator;
    } else if (len_indicator == 126) {
        if (bytes_read < 4) return 0;
        *payload_len = (buffer[2] << 8) | buffer[3];
        header_len = 4;
    } else if (len_indicator == 127) {
        if (bytes_read < 10) return 0;
        *payload_len = (buffer[6] << 24) | (buffer[7] << 16) | (buffer[8] << 8) | buffer[9];
        header_len = 10;
    }

    if (masked) {
        header_len += 4;
    }

    if (bytes_read < header_len + *payload_len) {
        return 0;
    }

    *payload = (char *)buffer + header_len;
    if (masked) {
        char *mask_key = (char *)buffer + header_len - 4;
        for (int i = 0; i < *payload_len; i++) {
            (*payload)[i] ^= mask_key[i % 4];
        }
    }

    return 1;
}

void *forward_tcp_to_ws(void *arg) {
    thread_args_t *args = (thread_args_t *)arg;
    SSL *ssl = args->ssl;
    int debug = args->debug;
    char buffer[BUFFER_SIZE];
    int bytes_read;
    fd_set readfds;
    struct timeval tv;

    while (1) {
        pthread_mutex_lock(&tunnel_mutex);
        if (!active_tunnel || !active_tunnel->active) {
            pthread_mutex_unlock(&tunnel_mutex);
            break;
        }
        int sock = active_tunnel->local_sock;
        char request_id[37];
        strcpy(request_id, active_tunnel->request_id);
        pthread_mutex_unlock(&tunnel_mutex);

        // Use select to wait for data on local socket
        FD_ZERO(&readfds);
        FD_SET(sock, &readfds);
        tv.tv_sec = 0;  // 0 seconds
        tv.tv_usec = 50000;  // 50ms timeout

        int retval = select(sock + 1, &readfds, NULL, NULL, &tv);
        if (retval == -1) {
            if (debug) {
                perror("[DEBUG] select failed");
                fflush(stdout);
            }
            break;
        } else if (retval == 0) {
            // Timeout, continue loop
            continue;
        }

        if (FD_ISSET(sock, &readfds)) {
            bytes_read = recv(sock, buffer, sizeof(buffer), 0);
            if (bytes_read <= 0) {
                if (debug) {
                    printf("[DEBUG] TCP connection closed or error\n");
                    fflush(stdout);
                }
                break;
            }

            if (debug) {
                printf("[DEBUG] Forwarding %d bytes from TCP to WebSocket\n", bytes_read);
                fflush(stdout);
            }

            // Convert to hex
            char hex_data[bytes_read * 2 + 1];
            for (int i = 0; i < bytes_read; i++) {
                sprintf(hex_data + i * 2, "%02x", (unsigned char)buffer[i]);
            }
            hex_data[bytes_read * 2] = '\0';

            // Send as tunnel_data
            char message[BUFFER_SIZE];
            snprintf(message, sizeof(message),
                     "{\"type\":\"tunnel_data\",\"request_id\":\"%s\",\"data\":\"%s\"}",
                     request_id, hex_data);

            if (!send_websocket_frame(ssl, message)) {
                if (debug) {
                    printf("[DEBUG] Failed to send WebSocket frame\n");
                    fflush(stdout);
                }
                break;
            }
        }
    }

    if (debug) {
        printf("[DEBUG] TCP to WebSocket forwarding thread exiting\n");
        fflush(stdout);
    }

    free(args);
    return NULL;
}

void handle_tunnel_data(SSL *ssl __attribute__((unused)), const char *request_id, const char *data_hex, int debug __attribute__((unused))) {
    pthread_mutex_lock(&tunnel_mutex);
    if (!active_tunnel || strcmp(active_tunnel->request_id, request_id) != 0) {
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }
    int sock = active_tunnel->local_sock;
    pthread_mutex_unlock(&tunnel_mutex);

    // Decode hex data
    size_t data_len = strlen(data_hex) / 2;
    char *data = malloc(data_len);
    if (!data) return;

    for (size_t i = 0; i < data_len; i++) {
        sscanf(data_hex + i * 2, "%2hhx", &data[i]);
    }

    // Send to local socket
    send(sock, data, data_len, 0);
    free(data);
}

void handle_tunnel_close(SSL *ssl __attribute__((unused)), const char *request_id, int debug) {
    pthread_mutex_lock(&tunnel_mutex);
    if (active_tunnel && strcmp(active_tunnel->request_id, request_id) == 0) {
        active_tunnel->active = 0;
        close(active_tunnel->local_sock);
        free(active_tunnel);
        active_tunnel = NULL;
        if (debug) {
            printf("[DEBUG] Tunnel %s closed\n", request_id);
        }
    }
    pthread_mutex_unlock(&tunnel_mutex);
}

int setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id, int local_port, int debug) {
    struct sockaddr_in server_addr;
    struct hostent *he;
    int sock;
    SSL_CTX *ssl_ctx;
    SSL *ssl;
    char buffer[BUFFER_SIZE];
    int bytes_read;

    // Generate request ID
    char request_id[37];
    generate_request_id(request_id, sizeof(request_id));

    // Resolve hostname
    if ((he = gethostbyname(wssshd_host)) == NULL) {
        herror("gethostbyname");
        return -1;
    }

    // Create socket
    if ((sock = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
        perror("Socket creation failed");
        return -1;
    }

    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(wssshd_port);
    server_addr.sin_addr = *((struct in_addr *)he->h_addr);

    // Connect to server
    if (connect(sock, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("Connection failed");
        close(sock);
        return -1;
    }

    // Initialize SSL
    SSL_library_init();
    OpenSSL_add_all_algorithms();
    SSL_load_error_strings();

    ssl_ctx = SSL_CTX_new(TLS_client_method());
    if (!ssl_ctx) {
        ERR_print_errors_fp(stderr);
        close(sock);
        return -1;
    }

    // Allow self-signed certificates
    SSL_CTX_set_verify(ssl_ctx, SSL_VERIFY_NONE, NULL);

    ssl = SSL_new(ssl_ctx);
    SSL_set_fd(ssl, sock);

    if (debug) {
        printf("[DEBUG] Establishing SSL connection...\n");
        fflush(stdout);
    }
    if (SSL_connect(ssl) <= 0) {
        ERR_print_errors_fp(stderr);
        fprintf(stderr, "SSL connection failed\n");
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }
    if (debug) {
        printf("[DEBUG] SSL connection established\n");
        fflush(stdout);
    }

    // Perform WebSocket handshake
    if (debug) {
        printf("[DEBUG] Performing WebSocket handshake to %s:%d\n", wssshd_host, wssshd_port);
        fflush(stdout);
    }
    if (!websocket_handshake(ssl, wssshd_host, wssshd_port, "/")) {
        fprintf(stderr, "WebSocket handshake failed\n");
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }
    if (debug) {
        printf("[DEBUG] WebSocket handshake successful\n");
        fflush(stdout);
    }

    // Send tunnel request
    if (!send_json_message(ssl, "tunnel_request", client_id, request_id)) {
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    if (debug) {
        printf("[DEBUG] Tunnel request sent for client: %s, request_id: %s\n", client_id, request_id);
    }

    // Read acknowledgment
    bytes_read = SSL_read(ssl, buffer, sizeof(buffer));
    if (bytes_read <= 0) {
        if (debug) {
            printf("[DEBUG] No acknowledgment received\n");
        }
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    if (debug) {
        printf("[DEBUG] Read %d bytes acknowledgment\n", bytes_read);
        printf("[DEBUG] Raw acknowledgment: ");
        for (int i = 0; i < bytes_read && i < 50; i++) {
            if (buffer[i] >= 32 && buffer[i] < 127) {
                printf("%c", buffer[i]);
            } else {
                printf("\\x%02x", (unsigned char)buffer[i]);
            }
        }
        printf("\n");
        fflush(stdout);
    }

    // Check if it's a close frame
    if (bytes_read >= 2 && (buffer[0] & 0x8F) == 0x88) {
        if (debug) {
            printf("[DEBUG] Server sent close frame\n");
            if (bytes_read >= 6) {
                int close_code = (buffer[2] << 8) | buffer[3];
                printf("[DEBUG] Close code: %d\n", close_code);
                if (bytes_read > 6) {
                    printf("[DEBUG] Close reason: %.*s\n", bytes_read - 6, buffer + 6);
                }
            }
            fflush(stdout);
        }
        fprintf(stderr, "Tunnel request rejected by server\n");
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    // Parse WebSocket frame
    char *payload;
    int payload_len;
    if (!parse_websocket_frame(buffer, bytes_read, &payload, &payload_len)) {
        fprintf(stderr, "Failed to parse WebSocket frame\n");
        if (debug) {
            printf("[DEBUG] Frame parsing failed. First few bytes: ");
            for (int i = 0; i < bytes_read && i < 10; i++) {
                printf("%02x ", (unsigned char)buffer[i]);
            }
            printf("\n");
            fflush(stdout);
        }
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    // Null terminate payload
    payload[payload_len] = '\0';

    if (debug) {
        printf("[DEBUG] Received payload: '%s' (length: %d)\n", payload, payload_len);
        printf("[DEBUG] Looking for tunnel_ack in: ");
        for (int i = 0; i < payload_len && i < 50; i++) {
            if (payload[i] >= 32 && payload[i] < 127) {
                printf("%c", payload[i]);
            } else {
                printf("\\x%02x", (unsigned char)payload[i]);
            }
        }
        printf("\n");
        fflush(stdout);
    }

    // Check for tunnel acknowledgment
    if (strstr(payload, "tunnel_ack") == NULL) {
        fprintf(stderr, "Tunnel request denied or failed: %s\n", payload);
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    if (debug) {
        printf("[DEBUG] Tunnel established, local port: %d\n", local_port);
    }

    // Create tunnel structure
    active_tunnel = malloc(sizeof(tunnel_t));
    if (!active_tunnel) {
        perror("Memory allocation failed");
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    strcpy(active_tunnel->request_id, request_id);
    active_tunnel->local_sock = -1;
    active_tunnel->active = 1;
    active_tunnel->ssl = ssl;

    // Start listening on local port
    int listen_sock = socket(AF_INET, SOCK_STREAM, 0);
    if (listen_sock < 0) {
        perror("Local socket creation failed");
        free(active_tunnel);
        active_tunnel = NULL;
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    struct sockaddr_in local_addr;
    memset(&local_addr, 0, sizeof(local_addr));
    local_addr.sin_family = AF_INET;
    local_addr.sin_port = htons(local_port);
    local_addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);

    if (bind(listen_sock, (struct sockaddr *)&local_addr, sizeof(local_addr)) < 0) {
        perror("Local bind failed");
        close(listen_sock);
        free(active_tunnel);
        active_tunnel = NULL;
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    if (listen(listen_sock, 1) < 0) {
        perror("Local listen failed");
        close(listen_sock);
        free(active_tunnel);
        active_tunnel = NULL;
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    if (debug) {
        printf("[DEBUG] Listening on localhost:%d\n", local_port);
        fflush(stdout);
    }

    // Return success - tunnel is set up and listening
    return listen_sock;
}

int main(int argc, char *argv[]) {
    wsssh_config_t config = {
        .local_port = NULL,
        .debug = 0
    };

    // Easter egg: --support option (only when it's the only argument)
    if (argc == 2 && strcmp(argv[1], "--support") == 0) {
        print_trans_flag();
        return 0;
    }

    // Handle --help before parsing
    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0) {
            print_usage(argv[0]);
            return 0;
        }
    }

    pthread_mutex_init(&tunnel_mutex, NULL);

    // Parse wsssh arguments
    int remaining_argc;
    char **remaining_argv;
    if (!parse_args(argc, argv, &config, &remaining_argc, &remaining_argv)) {
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    // Need at least one SSH argument (the target)
    if (remaining_argc == 0) {
        fprintf(stderr, "Error: No target specified\n");
        print_usage(argv[0]);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    // Parse SSH arguments to extract host and port
    char *ssh_host = NULL;
    int ssh_port = DEFAULT_PORT;
    if (!parse_ssh_args(remaining_argc, remaining_argv, &ssh_host, &ssh_port, config.debug)) {
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    // Parse the SSH host to extract client_id and wssshd_host
    char *client_id = NULL;
    char *wssshd_host = NULL;
    int wssshd_port = ssh_port;  // The -p port becomes the WebSocket server port

    if (!parse_hostname(ssh_host, &client_id, &wssshd_host)) {
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    if (config.debug) {
        printf("[DEBUG] SSH Host: %s\n", ssh_host);
        printf("[DEBUG] Client ID: %s\n", client_id);
        printf("[DEBUG] WSSSHD Host: %s\n", wssshd_host);
        printf("[DEBUG] WSSSHD Port: %d\n", wssshd_port);
    }

    // Find available local port
    int local_port = config.local_port ? atoi(config.local_port) : find_available_port();
    if (local_port == 0) {
        fprintf(stderr, "Error: Could not find available local port\n");
        free(client_id);
        free(wssshd_host);
        free(config.local_port);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    if (config.debug) {
        printf("[DEBUG] Using local port: %d\n", local_port);
    }

    // Modify SSH arguments
    int new_ssh_argc;
    char **new_ssh_args = modify_ssh_args(remaining_argc, remaining_argv, ssh_host, local_port, &new_ssh_argc);
    if (!new_ssh_args) {
        fprintf(stderr, "Error: Failed to modify SSH arguments\n");
        free(client_id);
        free(wssshd_host);
        free(config.local_port);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    if (config.debug) {
        printf("[DEBUG] Modified SSH command:");
        for (int i = 0; new_ssh_args[i]; i++) {
            printf(" %s", new_ssh_args[i]);
        }
        printf("\n");
        fflush(stdout);
    }

    // Parent process: setup tunnel FIRST
    if (config.debug) {
        printf("[DEBUG] Parent process setting up tunnel...\n");
        fflush(stdout);
    }

    int listen_sock = setup_tunnel(wssshd_host, wssshd_port, client_id, local_port, config.debug);
    if (listen_sock < 0) {
        fprintf(stderr, "Error: Failed to establish tunnel\n");
        free(client_id);
        free(wssshd_host);
        free(config.local_port);
        for (int i = 1; i < new_ssh_argc; i++) {
            if (new_ssh_args[i] != remaining_argv[i-1]) {
                free(new_ssh_args[i]);
            }
        }
        free(new_ssh_args);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    // NOW fork to run SSH in background (after tunnel is ready)
    if (config.debug) {
        printf("[DEBUG] About to fork SSH process...\n");
        fflush(stdout);
    }
    pid_t pid = fork();
    if (pid < 0) {
        perror("fork failed");
        free(client_id);
        free(wssshd_host);
        free(config.local_port);
        for (int i = 1; i < new_ssh_argc; i++) {
            if (new_ssh_args[i] != remaining_argv[i-1]) {
                free(new_ssh_args[i]);
            }
        }
        free(new_ssh_args);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    if (pid == 0) {
        // Child process: run SSH
        if (config.debug) {
            printf("[DEBUG] Child process starting SSH...\n");
            fflush(stdout);
        }
        execvp("ssh", new_ssh_args);

        // If we reach here, execvp failed
        perror("execvp failed");
        free(client_id);
        free(wssshd_host);
        for (int i = 1; i < new_ssh_argc; i++) {
            if (new_ssh_args[i] != remaining_argv[i-1]) {
                free(new_ssh_args[i]);
            }
        }
        free(new_ssh_args);
        exit(1);
    }

    // Parent process: accept SSH connection and start forwarding
    if (config.debug) {
        printf("[DEBUG] Waiting for SSH connection on localhost:%d...\n", local_port);
        fflush(stdout);
    }

    struct sockaddr_in client_addr;
    socklen_t client_len = sizeof(client_addr);
    active_tunnel->local_sock = accept(listen_sock, (struct sockaddr *)&client_addr, &client_len);
    if (active_tunnel->local_sock < 0) {
        perror("Local accept failed");
        kill(pid, SIGTERM);
        waitpid(pid, NULL, 0);
        close(listen_sock);
        free(active_tunnel);
        active_tunnel = NULL;
        free(client_id);
        free(wssshd_host);
        free(config.local_port);
        for (int i = 1; i < new_ssh_argc; i++) {
            if (new_ssh_args[i] != remaining_argv[i-1]) {
                free(new_ssh_args[i]);
            }
        }
        free(new_ssh_args);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    close(listen_sock); // No longer needed

    if (config.debug) {
        printf("[DEBUG] Local SSH connection accepted! Starting data forwarding...\n");
        fflush(stdout);
    }

    // Start forwarding thread
    thread_args_t *thread_args = malloc(sizeof(thread_args_t));
    if (!thread_args) {
        perror("Memory allocation failed for thread args");
        kill(pid, SIGTERM);
        waitpid(pid, NULL, 0);
        close(active_tunnel->local_sock);
        free(active_tunnel);
        active_tunnel = NULL;
        free(client_id);
        free(wssshd_host);
        free(config.local_port);
        for (int i = 1; i < new_ssh_argc; i++) {
            if (new_ssh_args[i] != remaining_argv[i-1]) {
                free(new_ssh_args[i]);
            }
        }
        free(new_ssh_args);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }
    thread_args->ssl = active_tunnel->ssl;  // Need to store SSL in tunnel struct
    thread_args->debug = config.debug;

    pthread_t thread;
    pthread_create(&thread, NULL, forward_tcp_to_ws, thread_args);
    pthread_detach(thread);

    // Main tunnel loop - handle WebSocket messages
    char buffer[BUFFER_SIZE];
    int bytes_read;
    fd_set readfds;
    struct timeval tv;
    int ssl_fd = SSL_get_fd(active_tunnel->ssl);

    while (active_tunnel && active_tunnel->active) {
        // Use select to wait for data on SSL socket with timeout
        FD_ZERO(&readfds);
        FD_SET(ssl_fd, &readfds);
        tv.tv_sec = 0;  // 0 seconds
        tv.tv_usec = 50000;  // 50ms timeout

        int retval = select(ssl_fd + 1, &readfds, NULL, NULL, &tv);
        if (retval == -1) {
            if (config.debug) {
                perror("[DEBUG] select on SSL fd failed");
                fflush(stdout);
            }
            break;
        } else if (retval == 0) {
            // Timeout, continue loop
            continue;
        }

        if (FD_ISSET(ssl_fd, &readfds)) {
            bytes_read = SSL_read(active_tunnel->ssl, buffer, sizeof(buffer));
            if (bytes_read <= 0) {
                if (config.debug) {
                    printf("[DEBUG] WebSocket connection closed\n");
                    fflush(stdout);
                }
                break;
            }

            // Check frame type
            unsigned char frame_type = buffer[0] & 0x8F;
    
            // Handle close frame
            if (frame_type == 0x88) {
                if (config.debug) {
                    printf("[DEBUG] Received close frame from server\n");
                    fflush(stdout);
                }
                break;
            }
    
            // Handle ping frame
            if (frame_type == 0x89) {
                if (config.debug) {
                    printf("[DEBUG] Received ping frame, sending pong\n");
                    fflush(stdout);
                }
                // Parse the ping frame to get payload
                char *ping_payload;
                int ping_payload_len;
                if (parse_websocket_frame(buffer, bytes_read, &ping_payload, &ping_payload_len)) {
                    // Send pong with same payload
                    if (!send_pong_frame(active_tunnel->ssl, ping_payload, ping_payload_len)) {
                        if (config.debug) {
                            printf("[DEBUG] Failed to send pong frame\n");
                            fflush(stdout);
                        }
                    }
                } else {
                    if (config.debug) {
                        printf("[DEBUG] Failed to parse ping frame\n");
                        fflush(stdout);
                    }
                }
                continue;
            }
    
            // Handle pong frame
            if (frame_type == 0x8A) {
                if (config.debug) {
                    printf("[DEBUG] Received pong frame\n");
                    fflush(stdout);
                }
                // Just acknowledge, no response needed
                continue;
            }
    
            // Handle regular frames (text/binary)
            char *payload;
            int payload_len;
            if (!parse_websocket_frame(buffer, bytes_read, &payload, &payload_len)) {
                if (config.debug) {
                    printf("[DEBUG] Failed to parse WebSocket frame\n");
                    fflush(stdout);
                }
                continue;
            }
    
            payload[payload_len] = '\0';
    
            if (config.debug) {
                printf("[DEBUG] Received: %s\n", payload);
                fflush(stdout);
            }
    
            // Handle messages
            if (strstr(payload, "tunnel_data")) {
                if (config.debug) {
                    printf("[DEBUG] Received tunnel_data message\n");
                    fflush(stdout);
                }
                // Extract request_id and data
                char *id_start = strstr(payload, "\"request_id\"");
                char *data_start = strstr(payload, "\"data\"");
                if (id_start && data_start) {
                    char *colon = strchr(id_start, ':');
                    if (colon) {
                        char *open_quote = strchr(colon, '"');
                        if (open_quote) {
                            id_start = open_quote + 1;
                            char *close_quote = strchr(id_start, '"');
                            if (close_quote) {
                                *close_quote = '\0';
                                char *data_colon = strchr(data_start, ':');
                                if (data_colon) {
                                    char *data_quote = strchr(data_colon, '"');
                                    if (data_quote) {
                                        data_start = data_quote + 1;
                                        char *data_end = strchr(data_start, '"');
                                        if (data_end) {
                                            *data_end = '\0';
                                            handle_tunnel_data(active_tunnel->ssl, id_start, data_start, config.debug);
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            } else if (strstr(payload, "tunnel_close")) {
                if (config.debug) {
                    printf("[DEBUG] Received tunnel_close message\n");
                    fflush(stdout);
                }
                char *id_start = strstr(payload, "\"request_id\"");
                if (id_start) {
                    char *colon = strchr(id_start, ':');
                    if (colon) {
                        char *open_quote = strchr(colon, '"');
                        if (open_quote) {
                            id_start = open_quote + 1;
                            char *close_quote = strchr(id_start, '"');
                            if (close_quote) {
                                *close_quote = '\0';
                                handle_tunnel_close(active_tunnel->ssl, id_start, config.debug);
                            }
                        }
                    }
                }
            } else {
                if (config.debug) {
                    printf("[DEBUG] Received unknown message type: %s\n", payload);
                    fflush(stdout);
                }
            }
        }
    }

    if (config.debug) {
        printf("[DEBUG] Tunnel loop ended, waiting for SSH to finish...\n");
        fflush(stdout);
    }

    // Wait for SSH to finish
    int status;
    waitpid(pid, &status, 0);

    if (config.debug) {
        printf("[DEBUG] SSH process finished with status: %d\n", status);
        fflush(stdout);
    }

    // Wait a few seconds before exiting cleanly
    if (config.debug) {
        printf("[DEBUG] Waiting 3 seconds before exit...\n");
        fflush(stdout);
    }
    sleep(3);

    // Cleanup
    if (active_tunnel) {
        close(active_tunnel->local_sock);
        SSL_free(active_tunnel->ssl);
        free(active_tunnel);
        active_tunnel = NULL;
    }

    free(client_id);
    free(wssshd_host);
    free(config.local_port);
    // Free allocated strings in new_ssh_args
    for (int i = 0; i < new_ssh_argc; i++) {
        // Free strings that were allocated with malloc/strdup
        if (i == 4) { // The port string
            free(new_ssh_args[i]);
        } else if (i == 5) { // The user@localhost string (if allocated)
            // Check if this was allocated (contains @localhost)
            if (strstr(new_ssh_args[i], "@localhost")) {
                free(new_ssh_args[i]);
            }
        }
    }
    free(new_ssh_args);
    pthread_mutex_destroy(&tunnel_mutex);

    return 0;
}