/*
 * WebSocket SSH Client (wssshc) - C Implementation
 * WebSocket SSH client that registers with wssshd server.
 *
 * 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 <pthread.h>
#include <errno.h>

#define BUFFER_SIZE 4096
#define DEFAULT_PORT 9898

int send_all(SSL *ssl, const char *data, int len);

int global_debug = 0;
pthread_mutex_t tunnel_mutex;

typedef struct {
    int sock;  // Local SSH connection socket
    char request_id[64];
} tunnel_t;

typedef struct {
    char *server_ip;
    int port;
    char *client_id;
    char *password;
    int interval;
    int debug;
} wssshc_config_t;

void load_config_file(const char *config_path, wssshc_config_t *config) {
    FILE *file = fopen(config_path, "r");
    if (!file) return;

    char line[256];
    int in_section = 0;
    while (fgets(line, sizeof(line), file)) {
        line[strcspn(line, "\n")] = 0;
        if (line[0] == '[' && strstr(line, "wssshc")) {
            in_section = 1;
        } else if (line[0] == '[') {
            in_section = 0;
        } else if (in_section && strchr(line, '=')) {
            char *key = strtok(line, "=");
            char *value = strtok(NULL, "=");
            if (key && value) {
                while (*key == ' ') key++;
                char *end = key + strlen(key) - 1;
                while (end > key && *end == ' ') *end-- = 0;
                while (*value == ' ') value++;
                end = value + strlen(value) - 1;
                while (end > value && *end == ' ') *end-- = 0;
                if (strcmp(key, "password") == 0 && !config->password) {
                    config->password = strdup(value);
                } else if (strcmp(key, "server-ip") == 0 && !config->server_ip) {
                    config->server_ip = strdup(value);
                } else if (strcmp(key, "domain") == 0 && !config->server_ip) {
                    config->server_ip = strdup(value);
                } else if (strcmp(key, "port") == 0) {
                    config->port = atoi(value);
                } else if (strcmp(key, "id") == 0 && !config->client_id) {
                    config->client_id = strdup(value);
                } else if (strcmp(key, "interval") == 0) {
                    config->interval = atoi(value);
                }
            }
        }
    }
    fclose(file);
}

void load_config(wssshc_config_t *config) {
    // Configuration hierarchy:
    // 1. System config (/etc/wssshc.conf) - lowest priority
    // 2. User config (~/.config/wsssh/wssshc.conf) - medium priority
    // 3. Command line arguments - highest priority (handled in parse_args)

    // Load system config first (lowest priority)
    load_config_file("/etc/wssshc.conf", config);

    // Load user config (medium priority, overrides system)
    char *home = getenv("HOME");
    if (home) {
        char user_config_path[256];
        snprintf(user_config_path, sizeof(user_config_path), "%s/.config/wsssh/wssshc.conf", home);
        load_config_file(user_config_path, config);
    }
}

tunnel_t *active_tunnel = NULL;

void *tunnel_thread(void *arg);
void cleanup_tunnel();

void handle_tunnel_request(SSL *ssl, const char *request_id, int debug) {
    pthread_mutex_lock(&tunnel_mutex);
    if (active_tunnel) {
        close(active_tunnel->sock);
        free(active_tunnel);
    }
    active_tunnel = malloc(sizeof(tunnel_t));
    if (!active_tunnel) {
        perror("Memory allocation failed");
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }

    active_tunnel->sock = -1; // Not connected yet
    strcpy(active_tunnel->request_id, request_id);
    pthread_mutex_unlock(&tunnel_mutex);

    if (debug) {
        printf("[DEBUG] Tunnel %s prepared (not connected yet)\n", request_id);
    }

    // Send tunnel_ack back to server
    char ack_msg[256];
    snprintf(ack_msg, sizeof(ack_msg), "{\"type\":\"tunnel_ack\",\"request_id\":\"%s\"}", request_id);

    if (debug) {
        printf("[DEBUG] Sending tunnel_ack: %s\n", ack_msg);
    }

    char frame[512];
    frame[0] = 0x81; // FIN + text opcode
    int msg_len = strlen(ack_msg);
    frame[1] = 0x80 | msg_len; // MASK + length
    char mask_key[4];
    for (int i = 0; i < 4; i++) {
        mask_key[i] = rand() % 256;
        frame[2 + i] = mask_key[i];
    }
    for (int i = 0; i < msg_len; i++) {
        frame[6 + i] = ack_msg[i] ^ mask_key[i % 4];
    }
    int frame_len = 6 + msg_len;
    if (!send_all(ssl, frame, frame_len)) {
        fprintf(stderr, "Send tunnel_ack failed\n");
    }
}

void handle_tunnel_data(SSL *ssl, const char *request_id, const char *data_hex, int debug) {
    printf("handle_tunnel_data called with id: %s, data len: %zu\n", request_id, strlen(data_hex));
    pthread_mutex_lock(&tunnel_mutex);
    if (!active_tunnel || strcmp(active_tunnel->request_id, request_id) != 0) {
        printf("No active tunnel or id mismatch\n");
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }
    int sock_fd = active_tunnel->sock;
    pthread_mutex_unlock(&tunnel_mutex);

    // Open local SSH connection if not already open
    if (sock_fd == -1) {
        struct sockaddr_in addr;
        int sock;

        if (debug) {
            printf("[DEBUG] [TCP Tunnel] Opening local SSH connection for tunnel %s\n", request_id);
        }

        // Create socket for local SSH connection
        if ((sock = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
            perror("Local socket creation failed");
            return;
        }

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

        // Connect to local SSH
        if (connect(sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
            perror("Local SSH connection failed");
            close(sock);
            return;
        }

        pthread_mutex_lock(&tunnel_mutex);
        if (active_tunnel && strcmp(active_tunnel->request_id, request_id) == 0) {
            active_tunnel->sock = sock;
            sock_fd = sock;

            // Start tunnel thread
            pthread_t thread;
            pthread_create(&thread, NULL, tunnel_thread, ssl);
            pthread_detach(thread);

            if (debug) {
                printf("[DEBUG] [TCP Tunnel] Local SSH connection established and tunnel thread started\n");
            }
        } else {
            close(sock);
        }
        pthread_mutex_unlock(&tunnel_mutex);
    } else {
        // SSH connection already exists, verify it's still valid
        char test_buf[1];
        int result = recv(sock_fd, test_buf, 1, MSG_PEEK | MSG_DONTWAIT);
        if (result < 0 && errno != EAGAIN && errno != EWOULDBLOCK) {
            // SSH connection is broken, close it and create a new one
            if (debug) {
                printf("[DEBUG] [TCP Tunnel] Existing SSH connection is broken, creating new one\n");
            }
            close(sock_fd);

            struct sockaddr_in addr;
            int sock;

            // Create new socket for local SSH connection
            if ((sock = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
                perror("Local socket creation failed");
                pthread_mutex_lock(&tunnel_mutex);
                if (active_tunnel && strcmp(active_tunnel->request_id, request_id) == 0) {
                    active_tunnel->sock = -1;
                }
                pthread_mutex_unlock(&tunnel_mutex);
                return;
            }

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

            // Connect to local SSH
            if (connect(sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
                perror("Local SSH connection failed");
                close(sock);
                pthread_mutex_lock(&tunnel_mutex);
                if (active_tunnel && strcmp(active_tunnel->request_id, request_id) == 0) {
                    active_tunnel->sock = -1;
                }
                pthread_mutex_unlock(&tunnel_mutex);
                return;
            }

            pthread_mutex_lock(&tunnel_mutex);
            if (active_tunnel && strcmp(active_tunnel->request_id, request_id) == 0) {
                active_tunnel->sock = sock;
                sock_fd = sock;

                // Start tunnel thread
                pthread_t thread;
                pthread_create(&thread, NULL, tunnel_thread, ssl);
                pthread_detach(thread);

                if (debug) {
                    printf("[DEBUG] [TCP Tunnel] New SSH connection established and tunnel thread started\n");
                }
            } else {
                close(sock);
            }
            pthread_mutex_unlock(&tunnel_mutex);
        } else if (debug) {
            printf("[DEBUG] [TCP Tunnel] Reusing existing SSH connection for tunnel %s\n", request_id);
        }
    }

    if (debug) {
        printf("[DEBUG] Handling tunnel_data, hex len: %zu, hex data: %.100s...\n", strlen(data_hex), data_hex);
    }

    // Decode hex data
    size_t data_len = strlen(data_hex) / 2;
    char *data = malloc(data_len);
    if (!data) {
        fprintf(stderr, "Memory allocation failed for tunnel data\n");
        return;
    }

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

    if (debug) {
        printf("[DEBUG] Decoded data len: %zu, sending to local SSH: ", data_len);
        for (size_t i = 0; i < data_len && i < 50; i++) {
            if (data[i] >= 32 && data[i] < 127) {
                printf("%c", data[i]);
            } else {
                printf("\\x%02x", (unsigned char)data[i]);
            }
        }
        printf("\n");
    }

    // Send to local SSH
    ssize_t sent = send(sock_fd, data, data_len, 0);
    if (sent < 0) {
        perror("Failed to send data to local SSH");
    } else if ((size_t)sent != data_len) {
        fprintf(stderr, "Partial send to local SSH: sent %zd of %zu bytes\n", sent, data_len);
    }
    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) {
        close(active_tunnel->sock);
        free(active_tunnel);
        active_tunnel = NULL;
        if (debug) {
            printf("[DEBUG] [TCP Tunnel] Tunnel %s closed\n", request_id);
        }
    }
    pthread_mutex_unlock(&tunnel_mutex);
}

void *tunnel_thread(void *arg) {
    SSL *ssl = (SSL *)arg;
    char buffer[BUFFER_SIZE];
    int bytes_read;

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

        bytes_read = recv(sock, buffer, sizeof(buffer), 0);
        if (bytes_read <= 0) {
            if (bytes_read < 0) {
                perror("Error reading from local SSH");
                // Check if this is due to tunnel cleanup (ECONNRESET, EPIPE, etc.)
                if (errno == ECONNRESET || errno == EPIPE || errno == EBADF) {
                    if (global_debug) {
                        printf("[DEBUG] [TCP Tunnel] SSH connection closed due to tunnel reset\n");
                    }
                }
            } else {
                if (global_debug) {
                    printf("[DEBUG] [TCP Tunnel] Local SSH connection closed normally\n");
                }
            }
            break;
        }

        if (global_debug) {
            printf("[DEBUG] [TCP Tunnel] Read %d bytes from local SSH: ", bytes_read);
            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");
        }

        // Send as tunnel_response
        char *hex_data = malloc(bytes_read * 2 + 1);
        if (!hex_data) {
            fprintf(stderr, "Memory allocation failed for hex_data in tunnel_thread\n");
            break;
        }
        for (int i = 0; i < bytes_read; i++) {
            sprintf(hex_data + i * 2, "%02x", (unsigned char)buffer[i]);
        }
        hex_data[bytes_read * 2] = '\0';

        if (global_debug) {
            printf("[DEBUG] [WebSocket] Sending tunnel_response, hex len: %zu, hex: %.100s...\n", strlen(hex_data), hex_data);
        }

        char response[9000];
        int response_len = snprintf(response, sizeof(response), "{\"type\":\"tunnel_response\",\"request_id\":\"%s\",\"data\":\"%s\"}", request_id, hex_data);
        if (response_len < 0 || (size_t)response_len >= sizeof(response)) {
            perror("JSON construction failed");
            free(hex_data);
            break;
        }

        if (global_debug) {
            printf("[DEBUG] Sending tunnel_response, len: %d\n", response_len);
        }

        // Send as WebSocket frame with proper length encoding
        int msg_len = response_len;
        int header_len = 2;
        char mask_key[4];

        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) {
            fprintf(stderr, "Memory allocation failed for WebSocket frame in tunnel_thread\n");
            free(hex_data);
            break;
        }

        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
            // For simplicity, assume length fits in 32 bits
            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
        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] = response[i] ^ mask_key[i % 4];
        }

        int frame_len = header_len + msg_len;
        if (!send_all(ssl, frame, frame_len)) {
            fprintf(stderr, "Send tunnel_response failed - tunnel may hang\n");
            free(frame);
            free(hex_data);
            break;
        }

        free(frame);
        free(hex_data);
    }

    return NULL;
}

void print_usage(const char *program_name) {
    fprintf(stderr, "Usage: %s [options]\n", program_name);
    fprintf(stderr, "WebSocket SSH Client - Register with wssshd server\n\n");
    fprintf(stderr, "Protect the dolls!\n\n");
    fprintf(stderr, "Options:\n");
    fprintf(stderr, "  --config FILE        Configuration file path (overrides default hierarchy)\n");
    fprintf(stderr, "  --server-ip IP       Server IP address\n");
    fprintf(stderr, "  --port PORT          Server port (default: %d)\n", DEFAULT_PORT);
    fprintf(stderr, "  --id ID              Client identifier\n");
    fprintf(stderr, "  --password PASS      Registration password\n");
    fprintf(stderr, "  --interval SEC       Reconnection interval (default: 30)\n");
    fprintf(stderr, "  --debug              Enable debug output\n");
    fprintf(stderr, "  --help               Show this help\n");
    fprintf(stderr, "\nConfiguration hierarchy (highest priority first):\n");
    fprintf(stderr, "  1. Command line arguments\n");
    fprintf(stderr, "  2. ~/.config/wsssh/wssshc.conf (user config)\n");
    fprintf(stderr, "  3. /etc/wssshc.conf (system config)\n");
    fprintf(stderr, "\nDonations:\n");
    fprintf(stderr, "  BTC: bc1q3zlkpu95amtcltsk85y0eacyzzk29v68tgc5hx\n");
    fprintf(stderr, "  ETH: 0xdA6dAb526515b5cb556d20269207D43fcc760E51\n");
}

int parse_args(int argc, char *argv[], wssshc_config_t *config) {
    static struct option long_options[] = {
        {"config", required_argument, 0, 'c'},
        {"server-ip", required_argument, 0, 's'},
        {"port", required_argument, 0, 'p'},
        {"id", required_argument, 0, 'i'},
        {"password", required_argument, 0, 'w'},
        {"interval", required_argument, 0, 't'},
        {"debug", no_argument, 0, 'd'},
        {"help", no_argument, 0, 'h'},
        {0, 0, 0, 0}
    };

    int opt;
    char *custom_config = NULL;

    while ((opt = getopt_long(argc, argv, "c:s:p:i:w:t:dh", long_options, NULL)) != -1) {
        switch (opt) {
            case 'c':
                custom_config = optarg;
                break;
            case 's':
                if (config->server_ip) free(config->server_ip);
                config->server_ip = strdup(optarg);
                break;
            case 'p':
                config->port = atoi(optarg);
                break;
            case 'i':
                if (config->client_id) free(config->client_id);
                config->client_id = strdup(optarg);
                break;
            case 'w':
                if (config->password) free(config->password);
                config->password = strdup(optarg);
                break;
            case 't':
                config->interval = atoi(optarg);
                break;
            case 'd':
                config->debug = 1;
                break;
            case 'h':
            default:
                print_usage(argv[0]);
                return 0;
        }
    }

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

    // If --config is specified, load that file instead of the hierarchy
    if (custom_config) {
        load_config_file(custom_config, config);
    }

    return 1;
}

SSL_CTX* init_ssl_ctx() {
    SSL_CTX *ctx;

    SSL_library_init();
    OpenSSL_add_all_algorithms();
    SSL_load_error_strings();

    ctx = SSL_CTX_new(TLS_client_method());
    if (!ctx) {
        ERR_print_errors_fp(stderr);
        return NULL;
    }

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

    return ctx;
}

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 (!send_all(ssl, request, strlen(request))) {
        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_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_json_message(SSL *ssl, const char *type, const char *id, const char *password) {
    char message[1024];
    char frame[1024];
    int message_len;

    if (password) {
        message_len = snprintf(message, sizeof(message),
                 "{\"type\":\"%s\",\"id\":\"%s\",\"password\":\"%s\"}",
                 type, id, password);
    } else {
        message_len = snprintf(message, sizeof(message),
                 "{\"type\":\"%s\",\"id\":\"%s\"}",
                 type, id);
    }

    // Send as WebSocket frame with proper length encoding
    frame[0] = 0x81; // FIN + text opcode
    int header_len = 2;

    if (message_len <= 125) {
        frame[1] = 0x80 | message_len; // MASK + length
    } else if (message_len <= 65535) {
        frame[1] = 0x80 | 126; // MASK + extended length
        frame[2] = (message_len >> 8) & 0xFF;
        frame[3] = message_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] = (message_len >> 24) & 0xFF;
        frame[7] = (message_len >> 16) & 0xFF;
        frame[8] = (message_len >> 8) & 0xFF;
        frame[9] = message_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 < message_len; i++) {
        frame[header_len + i] = message[i] ^ mask_key[i % 4];
    }

    int frame_len = header_len + message_len;

    if (!send_all(ssl, frame, frame_len)) {
        fprintf(stderr, "Send failed\n");
        return 0;
    }

    return 1;
}

int connect_to_server(const wssshc_config_t *config) {
    struct sockaddr_in server_addr;
    struct hostent *he;
    int sock;
    SSL_CTX *ssl_ctx;
    SSL *ssl;
    char buffer[BUFFER_SIZE];
    int bytes_read;

    // Clean up any leftover tunnel state from previous connection
    cleanup_tunnel();

    // Resolve hostname
    if ((he = gethostbyname(config->server_ip)) == 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(config->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_ctx = init_ssl_ctx();
    if (!ssl_ctx) {
        close(sock);
        return 1;
    }

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

    if (SSL_connect(ssl) <= 0) {
        ERR_print_errors_fp(stderr);
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return 1;
    }

    // Perform WebSocket handshake
    if (!websocket_handshake(ssl, config->server_ip, config->port, "/")) {
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return 1;
    }

    // Send registration message
    if (!send_json_message(ssl, "register", config->client_id, config->password)) {
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return 1;
    }

    // Read WebSocket frame with registration response
    bytes_read = SSL_read(ssl, buffer, sizeof(buffer));
    if (bytes_read <= 0) {
        fprintf(stderr, "Failed to read registration response\n");
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return 1;
    }

    // Parse WebSocket frame with extended length support
    if (bytes_read >= 2 && ((buffer[0] & 0x8F) == 0x81 || (buffer[0] & 0x8F) == 0x82)) { // Text or binary frame
        int masked = buffer[1] & 0x80;
        int len_indicator = buffer[1] & 0x7F;
        int header_len = 2;
        size_t payload_len;

        if (config->debug) {
            printf("[DEBUG] Initial len_indicator = %d, masked = %d\n", len_indicator, masked);
            fflush(stdout);
        }

        if (len_indicator <= 125) {
            payload_len = len_indicator;
        } else if (len_indicator == 126) {
            if (bytes_read < 4) {
                fprintf(stderr, "Incomplete extended length frame\n");
                SSL_free(ssl);
                SSL_CTX_free(ssl_ctx);
                close(sock);
                return 1;
            }
            payload_len = ((unsigned char)buffer[2] << 8) | (unsigned char)buffer[3];
            header_len = 4;
        } else if (len_indicator == 127) {
            if (bytes_read < 10) {
                fprintf(stderr, "Incomplete extended length frame\n");
                SSL_free(ssl);
                SSL_CTX_free(ssl_ctx);
                close(sock);
                return 1;
            }
            uint64_t len = 0;
            for (int i = 0; i < 8; i++) {
                len = (len << 8) | (unsigned char)buffer[2 + i];
            }
            payload_len = (int)len;
            header_len = 10;
        }

        if (masked) {
            header_len += 4;
        }

        if ((size_t)bytes_read >= (size_t)header_len + payload_len) {
            char *payload = buffer + header_len;
            if (masked) {
                // Unmask the payload
                char *mask_key = buffer + header_len - 4;
                for (size_t i = 0; i < payload_len; i++) {
                    payload[i] ^= mask_key[i % 4];
                }
            }
            // Copy payload to buffer start
            memmove(buffer, payload, payload_len);
            buffer[payload_len] = '\0';
        } else {
            fprintf(stderr, "WebSocket frame incomplete: need %zu bytes, got %d\n", payload_len + header_len, bytes_read);
            SSL_free(ssl);
            SSL_CTX_free(ssl_ctx);
            close(sock);
            return 1;
        }
    } else {
        buffer[bytes_read] = '\0';
        fprintf(stderr, "Unexpected frame: 0x%02x 0x%02x\n", buffer[0], buffer[1]);
    }

    if (strstr(buffer, "registered") == NULL) {
        fprintf(stderr, "Registration failed: %s\n", buffer);
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return 2;
    }

    printf("Connected and registered as %s\n", config->client_id);

    // Keep connection alive and handle tunnel requests
    // active_tunnels = NULL; // Will implement tunnel handling

    while (1) {
        bytes_read = SSL_read(ssl, buffer, sizeof(buffer));
        if (bytes_read <= 0) {
            if (bytes_read < 0) {
                ERR_print_errors_fp(stderr);
                fprintf(stderr, "SSL read error\n");
            } else {
                if (config->debug) {
                    printf("[DEBUG] Connection closed by server\n");
                }
            }
            // Clean up tunnel resources before breaking
            cleanup_tunnel();
            break;
        }

        if (config->debug) {
            printf("[DEBUG] [WebSocket] Read %d bytes, frame: 0x%02x 0x%02x 0x%02x 0x%02x\n", bytes_read, buffer[0], buffer[1], buffer[2], buffer[3]);
        }

        // Parse WebSocket frame with extended length support
        if (bytes_read >= 2 && (buffer[0] & 0x8F) == 0x81) { // Text frame
            int masked = buffer[1] & 0x80;
            int len_indicator = buffer[1] & 0x7F;
            int header_len = 2;
            size_t payload_len = 0;

            if (len_indicator <= 125) {
                payload_len = len_indicator;
            } else if (len_indicator == 126) {
                if (bytes_read < 4) {
                    if (config->debug) {
                        printf("[DEBUG] Incomplete extended length frame\n");
                    }
                    continue;
                }
                payload_len = ((unsigned char)buffer[2] << 8) | (unsigned char)buffer[3];
                header_len = 4;
            } else if (len_indicator == 127) {
                if (bytes_read < 10) {
                    if (config->debug) {
                        printf("[DEBUG] Incomplete extended length frame\n");
                    }
                    continue;
                }
                uint64_t len = 0;
                for (int i = 0; i < 8; i++) {
                    len = (len << 8) | (unsigned char)buffer[2 + i];
                }
                payload_len = (int)len;
                header_len = 10;
            }

            if (masked) {
                header_len += 4;
            }

            if (payload_len <= sizeof(buffer) - header_len && (size_t)bytes_read >= (size_t)header_len + payload_len) {
                char *payload = buffer + header_len;
                if (masked) {
                    char *mask_key = buffer + header_len - 4;
                    for (size_t i = 0; i < payload_len; i++) {
                        payload[i] ^= mask_key[i % 4];
                    }
                }
                memmove(buffer, payload, payload_len);
                buffer[payload_len] = '\0';

                // Handle message
                if (config->debug) {
                    printf("[DEBUG] [WebSocket] Received message: %s\n", buffer);
                }

                // Parse JSON message
                if (strstr(buffer, "tunnel_request")) {
                    // Extract request_id
                    char *id_start = strstr(buffer, "\"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';
                                    if (config->debug) {
                                        printf("[DEBUG] [WebSocket] Received tunnel_request for ID: %s\n", id_start);
                                    }
                                    handle_tunnel_request(ssl, id_start, config->debug);
                                }
                            }
                        }
                    }
                } else if (strstr(buffer, "tunnel_data")) {
                    if (config->debug) {
                        printf("[DEBUG] [WebSocket] Received tunnel_data message\n");
                    }
                    // Extract request_id and data
                    char *id_start = strstr(buffer, "\"request_id\"");
                    char *data_start = strstr(buffer, "\"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(ssl, id_start, data_start, config->debug);
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                } else if (strstr(buffer, "tunnel_close")) {
                    char *id_start = strstr(buffer, "\"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';
                                    if (config->debug) {
                                        printf("[DEBUG] [WebSocket] Received tunnel_close for ID: %s\n", id_start);
                                    }
                                    handle_tunnel_close(ssl, id_start, config->debug);
                                }
                            }
                        }
                    }
                }
            } else {
                fprintf(stderr, "WebSocket frame payload too large or incomplete: payload_len=%zu, buffer_size=%zu, bytes_read=%d\n",
                        payload_len, sizeof(buffer) - header_len, bytes_read);
                break;
            }
        } else if (bytes_read >= 2 && (buffer[0] & 0x8F) == 0x88) { // Close frame
            if (config->debug) {
                printf("[DEBUG] [WebSocket] Received close frame - cleaning up and reconnecting...\n");
            }
            // Clean up tunnel resources before reconnecting
            cleanup_tunnel();
            return 0;
        } else if (bytes_read >= 2 && (buffer[0] & 0x8F) == 0x89) { // Ping frame
            if (config->debug) {
                printf("[DEBUG] [WebSocket] Received ping frame\n");
            }
            // Parse ping frame and send pong with echoed payload
            int masked = buffer[1] & 0x80;
            int len_indicator = buffer[1] & 0x7F;
            int header_len = 2;
            size_t payload_len;

            if (len_indicator <= 125) {
                payload_len = len_indicator;
            } else if (len_indicator == 126) {
                if (bytes_read < 4) {
                    if (config->debug) {
                        printf("[DEBUG] Incomplete ping extended length frame\n");
                    }
                    continue;
                }
                payload_len = ((unsigned char)buffer[2] << 8) | (unsigned char)buffer[3];
                header_len = 4;
            } else if (len_indicator == 127) {
                if (bytes_read < 10) {
                    if (config->debug) {
                        printf("[DEBUG] Incomplete ping extended length frame\n");
                    }
                    continue;
                }
                uint64_t len = 0;
                for (int i = 0; i < 8; i++) {
                    len = (len << 8) | (unsigned char)buffer[2 + i];
                }
                if (len > SIZE_MAX) {
                    fprintf(stderr, "Ping payload too large\n");
                    continue;
                }
                payload_len = (size_t)len;
                header_len = 10;
            }

            if (masked) {
                header_len += 4;
            }

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

                // Send pong frame with same payload, properly masked
                char pong_frame[14 + 1024]; // Max header + max pong payload
                if (payload_len > 1024) {
                    fprintf(stderr, "Pong payload too large: %zu bytes\n", payload_len);
                    continue;
                }
                pong_frame[0] = 0x8A; // Pong opcode
                int pong_header_len = 2;
                char pong_mask_key[4];

                size_t pong_payload_len = payload_len;

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

                // Add mask key
                for (int i = 0; i < 4; i++) {
                    pong_mask_key[i] = rand() % 256;
                    pong_frame[pong_header_len + i] = pong_mask_key[i];
                }
                pong_header_len += 4;

                // Mask payload
                for (size_t i = 0; i < pong_payload_len; i++) {
                    pong_frame[pong_header_len + i] = payload[i] ^ pong_mask_key[i % 4];
                }

                int pong_frame_len = pong_header_len + (int)pong_payload_len;
                if (config->debug) {
                    printf("[DEBUG] [WebSocket] Sending pong frame, len: %d\n", pong_frame_len);
                }
                if (!send_all(ssl, pong_frame, pong_frame_len)) {
                    fprintf(stderr, "[ERROR] [WebSocket] Send pong failed\n");
                }
            }
        }
    }

    // Cleanup
    SSL_free(ssl);
    SSL_CTX_free(ssl_ctx);
    close(sock);
    return 1;
}

void cleanup_tunnel() {
    pthread_mutex_lock(&tunnel_mutex);
    if (active_tunnel) {
        if (active_tunnel->sock >= 0) {
            // Check if SSH connection is still valid before closing
            char test_buf[1];
            int result = recv(active_tunnel->sock, test_buf, 1, MSG_PEEK | MSG_DONTWAIT);
            if (result == 0 || (result < 0 && (errno == ECONNRESET || errno == EPIPE))) {
                // SSH connection is closed or broken, safe to close
                close(active_tunnel->sock);
                if (global_debug) {
                    printf("[DEBUG] [TCP Tunnel] Closed broken SSH connection\n");
                }
            } else {
                // SSH connection appears valid, don't close it
                if (global_debug) {
                    printf("[DEBUG] [TCP Tunnel] Keeping SSH connection alive for potential reuse\n");
                }
                // Reset socket to -1 so it will be reconnected if needed
                active_tunnel->sock = -1;
            }
        }
        free(active_tunnel);
        active_tunnel = NULL;
    }
    pthread_mutex_unlock(&tunnel_mutex);
}

int main(int argc, char *argv[]) {
    wssshc_config_t config = {
        .server_ip = NULL,
        .port = DEFAULT_PORT,
        .client_id = NULL,
        .password = NULL,
        .interval = 30,
        .debug = 0
    };

    pthread_mutex_init(&tunnel_mutex, NULL);

    // Load configuration files first (system and user configs)
    load_config(&config);

    // Parse command line arguments (highest priority, overrides config files)
    if (!parse_args(argc, argv, &config)) {
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    // Validate required arguments
    if (!config.server_ip || !config.client_id || !config.password) {
        fprintf(stderr, "Error: --server-ip, --id, and --password are required\n");
        print_usage(argv[0]);
        if (config.server_ip) free(config.server_ip);
        if (config.client_id) free(config.client_id);
        if (config.password) free(config.password);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    global_debug = config.debug;
    printf("WebSocket SSH Client starting...\n");

    while (1) {
        int result = connect_to_server(&config);
        if (result == 1) {
            // Error condition - use normal retry interval
            printf("Connection lost, retrying in %d seconds...\n", config.interval);
            sleep(config.interval);
        } else if (result == 0) {
            // Close frame received - add small delay to prevent rapid reconnection
            if (config.debug) {
                printf("[DEBUG] Server initiated disconnect, reconnecting in 2 seconds...\n");
            }
            sleep(2);
        }
        // result == 2 (registration failed) also uses normal interval
    }

    // Cleanup
    free(config.server_ip);
    free(config.client_id);
    free(config.password);
    pthread_mutex_destroy(&tunnel_mutex);

    return 0;
}