/*
 * WebSocket SSH Library - Tunnel management implementation
 *
 * 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 "tunnel.h"
#include "wssshlib.h"
#include "websocket.h"
#include "wssh_ssl.h"
#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 <fcntl.h>
#include <errno.h>

#define INITIAL_FRAME_BUFFER_SIZE 8192

// Global variables
tunnel_t *active_tunnel = NULL;
pthread_mutex_t tunnel_mutex;

frame_buffer_t *frame_buffer_init() {
    frame_buffer_t *fb = malloc(sizeof(frame_buffer_t));
    if (!fb) return NULL;

    fb->buffer = malloc(INITIAL_FRAME_BUFFER_SIZE);
    if (!fb->buffer) {
        free(fb);
        return NULL;
    }

    fb->size = INITIAL_FRAME_BUFFER_SIZE;
    fb->used = 0;
    return fb;
}

void frame_buffer_free(frame_buffer_t *fb) {
    if (fb) {
        free(fb->buffer);
        free(fb);
    }
}

int frame_buffer_resize(frame_buffer_t *fb, size_t new_size) {
    if (new_size <= fb->size) return 1;

    char *new_buffer = realloc(fb->buffer, new_size);
    if (!new_buffer) return 0;

    fb->buffer = new_buffer;
    fb->size = new_size;
    return 1;
}

int frame_buffer_append(frame_buffer_t *fb, const char *data, size_t len) {
    if (fb->used + len > fb->size) {
        size_t new_size = fb->size * 2;
        while (new_size < fb->used + len) {
            new_size *= 2;
        }
        if (!frame_buffer_resize(fb, new_size)) {
            return 0;
        }
    }

    memcpy(fb->buffer + fb->used, data, len);
    fb->used += len;
    return 1;
}

int frame_buffer_consume(frame_buffer_t *fb, size_t len) {
    if (len > fb->used) return 0;

    memmove(fb->buffer, fb->buffer + len, fb->used - len);
    fb->used -= len;
    return 1;
}

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 (SSL_write(ssl, frame, frame_len) <= 0) {
        fprintf(stderr, "Send tunnel_ack failed\n");
    }
}

void cleanup_tunnel(int debug) {
    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 (debug) {
                    printf("[DEBUG] [TCP Tunnel] Closed broken SSH connection\n");
                }
            } else {
                // SSH connection appears valid, don't close it
                if (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);
}

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);

        // Send pending data from outgoing buffer to local socket (wsscp only)
        if (active_tunnel->outgoing_buffer && active_tunnel->outgoing_buffer->used > 0) {
            ssize_t sent = send(sock, active_tunnel->outgoing_buffer->buffer, active_tunnel->outgoing_buffer->used, MSG_DONTWAIT);
            if (sent > 0) {
                frame_buffer_consume(active_tunnel->outgoing_buffer, sent);
                if (debug) {
                    printf("[DEBUG] Sent %zd bytes from buffer to local socket\n", sent);
                    fflush(stdout);
                }
            } else if (sent == -1 && errno != EAGAIN && errno != EWOULDBLOCK) {
                if (debug) {
                    perror("[DEBUG] send to local socket failed");
                    fflush(stdout);
                }
                pthread_mutex_unlock(&tunnel_mutex);
                break;
            }
        }

        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;
    }
    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]);
    }

    if (active_tunnel->outgoing_buffer) {
        // wsscp: Append to outgoing buffer
        pthread_mutex_lock(&tunnel_mutex);
        if (!frame_buffer_append(active_tunnel->outgoing_buffer, data, data_len)) {
            if (debug) {
                printf("[DEBUG] Failed to append to outgoing buffer, dropping %zu bytes\n", data_len);
                fflush(stdout);
            }
        }
        pthread_mutex_unlock(&tunnel_mutex);
    } else {
        // wsssh: Send directly to local socket
        send(active_tunnel->local_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);
        if (active_tunnel->outgoing_buffer) {
            frame_buffer_free(active_tunnel->outgoing_buffer);
        }
        free(active_tunnel);
        active_tunnel = NULL;
        if (debug) {
            printf("[DEBUG] Tunnel %s closed\n", request_id);
        }
    }
    pthread_mutex_unlock(&tunnel_mutex);
}

int reconnect_websocket(tunnel_t *tunnel, const char *wssshd_host, int wssshd_port, const char *client_id, const char *request_id, 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;

    if (debug) {
        printf("[DEBUG] Attempting to reconnect WebSocket connection...\n");
        fflush(stdout);
    }

    // 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;
    }

    // Create SSL context and connection
    ssl_ctx = create_ssl_context();
    if (!ssl_ctx) {
        close(sock);
        return -1;
    }

    ssl = create_ssl_connection(ssl_ctx, sock, debug);
    if (!ssl) {
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    // Perform WebSocket handshake
    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;
    }

    // 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;
    }

    // Read acknowledgment
    bytes_read = SSL_read(ssl, buffer, sizeof(buffer));
    if (bytes_read <= 0) {
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    // Check if it's a close frame
    if (bytes_read >= 2 && (buffer[0] & 0x8F) == 0x88) {
        if (debug) {
            printf("[DEBUG] Server sent close frame during reconnection\n");
            fflush(stdout);
        }
        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 during reconnection\n");
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

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

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

    // Success - update the tunnel's SSL connection
    if (tunnel->ssl) {
        SSL_free(tunnel->ssl);
    }
    tunnel->ssl = ssl;

    // Clean up old SSL context
    SSL_CTX_free(ssl_ctx);

    if (debug) {
        printf("[DEBUG] WebSocket reconnection successful\n");
        fflush(stdout);
    }

    return 0;
}

int setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id, int local_port, int debug, int use_buffer) {
    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;
    }

    // Create SSL context and connection
    ssl_ctx = create_ssl_context();
    if (!ssl_ctx) {
        close(sock);
        return -1;
    }

    ssl = create_ssl_connection(ssl_ctx, sock, debug);
    if (!ssl) {
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    // 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;
    }

    if (use_buffer) {
        active_tunnel->outgoing_buffer = frame_buffer_init();
        if (!active_tunnel->outgoing_buffer) {
            perror("Failed to initialize outgoing buffer");
            free(active_tunnel);
            SSL_free(ssl);
            SSL_CTX_free(ssl_ctx);
            close(sock);
            return -1;
        }
    } else {
        active_tunnel->outgoing_buffer = NULL;
    }

    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");
        if (use_buffer) frame_buffer_free(active_tunnel->outgoing_buffer);
        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);
        if (use_buffer) frame_buffer_free(active_tunnel->outgoing_buffer);
        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);
        if (use_buffer) frame_buffer_free(active_tunnel->outgoing_buffer);
        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);
    }

    // Clean up SSL context
    SSL_CTX_free(ssl_ctx);

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