/*
 * 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 <getopt.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <pthread.h>
#include <sys/select.h>

#include "wssshlib.h"
#include "websocket.h"
#include "wssh_ssl.h"
#include "tunnel.h"






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, "  --interval SEC       Connection retry interval in seconds (default: 30)\n");
    fprintf(stderr, "  --debug              Enable debug output\n");
    fprintf(stderr, "  --help               Show this help\n");
    fprintf(stderr, "  -p PORT              wssshd server port (default: 9898)\n");
    fprintf(stderr, "  --port PORT          Same as -p\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], "-p") == 0 && i + 1 < argc) {
            config->wssshd_port = atoi(argv[i + 1]);
            i++; // Skip the argument
        } else if (strcmp(argv[i], "--port") == 0 && i + 1 < argc) {
            config->wssshd_port = atoi(argv[i + 1]);
            i++; // Skip the argument
        } else if (strcmp(argv[i], "--interval") == 0 && i + 1 < argc) {
            config->interval = atoi(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, const char *config_domain) {
    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) {
        // No domain, use config
        if (!config_domain) {
            fprintf(stderr, "Error: Invalid hostname format. Expected client.domain format or domain in config\n");
            free(host_copy);
            return 0;
        }
        *client_id = strdup(host_copy);
        *wssshd_host = strdup(config_domain);
    } else {
        *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 debug) {
    *host = NULL;

    for (int i = 0; i < argc; i++) {
        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);
            }
            break;
        }
    }

    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;

    for (int i = 0; i < argc; i++) {
        // 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 main(int argc, char *argv[]) {
    // Read config
    char *config_domain = read_config_value("domain");

    wsssh_config_t config = {
        .local_port = NULL,
        .wssshd_port = 9898,
        .debug = 0,
        .interval = 30
    };

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

    // Handle --help and --free before parsing
    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0) {
            print_usage(argv[0]);
            free(config_domain);
            return 0;
        } else if (strcmp(argv[i], "--free") == 0) {
            print_palestinian_flag();
            free(config_domain);
            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);
        free(config_domain);
        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
    char *ssh_host = NULL;
    if (!parse_ssh_args(remaining_argc, remaining_argv, &ssh_host, config.debug)) {
        pthread_mutex_destroy(&tunnel_mutex);
        free(config_domain);
        return 1;
    }

    // Parse the SSH host to extract client_id and wssshd_host
    char *client_id = NULL;
    char *wssshd_host = NULL;
    int wssshd_port = config.wssshd_port;  // Use the wssshd port from config

    if (!parse_hostname(ssh_host, &client_id, &wssshd_host, config_domain)) {
        pthread_mutex_destroy(&tunnel_mutex);
        free(config_domain);
        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);
    }

    // Attempt initial tunnel setup with retry logic using --interval
    int listen_sock = -1;
    int setup_attempts = 0;
    int max_setup_attempts = 3;

    while (setup_attempts < max_setup_attempts && listen_sock < 0) {
        if (config.debug && setup_attempts > 0) {
            printf("[DEBUG] Initial tunnel setup attempt %d/%d\n", setup_attempts + 1, max_setup_attempts);
            fflush(stdout);
        }

        listen_sock = setup_tunnel(wssshd_host, wssshd_port, client_id, local_port, config.debug, 0);

        if (listen_sock < 0) {
            setup_attempts++;
            if (setup_attempts < max_setup_attempts) {
                if (config.debug) {
                    printf("[DEBUG] Initial tunnel setup failed, waiting %d seconds...\n", config.interval);
                    fflush(stdout);
                }
                sleep(config.interval);
            }
        }
    }

    if (listen_sock < 0) {
        fprintf(stderr, "Error: Failed to establish tunnel after %d attempts\n", max_setup_attempts);
        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 lost, attempting reconnection...\n");
                    fflush(stdout);
                }

                // Attempt reconnection - use 1 second intervals for WebSocket reconnections
                int reconnect_attempts = 0;
                int max_reconnect_attempts = 3;
                int reconnected = 0;

                while (reconnect_attempts < max_reconnect_attempts && !reconnected) {
                    if (config.debug) {
                        printf("[DEBUG] WebSocket reconnection attempt %d/%d\n", reconnect_attempts + 1, max_reconnect_attempts);
                        fflush(stdout);
                    }

                    if (reconnect_websocket(active_tunnel, wssshd_host, wssshd_port, client_id, active_tunnel->request_id, config.debug) == 0) {
                        reconnected = 1;
                        if (config.debug) {
                            printf("[DEBUG] WebSocket reconnection successful, continuing tunnel\n");
                            fflush(stdout);
                        }
                        // Update ssl_fd for select
                        ssl_fd = SSL_get_fd(active_tunnel->ssl);
                    } else {
                        reconnect_attempts++;
                        if (reconnect_attempts < max_reconnect_attempts) {
                            if (config.debug) {
                                printf("[DEBUG] WebSocket reconnection failed, waiting 1 second...\n");
                                fflush(stdout);
                            }
                            sleep(1);  // Use 1 second for WebSocket reconnections
                        }
                    }
                }

                if (!reconnected) {
                    if (config.debug) {
                        printf("[DEBUG] All reconnection attempts failed, exiting\n");
                        fflush(stdout);
                    }
                    break;
                }

                // Skip processing this iteration since we just reconnected
                continue;
            }

            // 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 == 8) { // The port string
            free(new_ssh_args[i]);
        } else if (strstr(new_ssh_args[i], "@localhost")) {
            // The user@localhost string (if allocated)
            free(new_ssh_args[i]);
        }
    }
    free(new_ssh_args);
    free(config_domain);
    pthread_mutex_destroy(&tunnel_mutex);

    return 0;
}