#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>

#define BUFFER_SIZE 4096
#define DEFAULT_PORT 22

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

void print_usage(const char *program_name) {
    fprintf(stderr, "Usage: %s [options] [scp_options...] source destination\n", program_name);
    fprintf(stderr, "WebSocket SCP Wrapper - SCP 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[], wsscp_config_t *config) {
    static struct option long_options[] = {
        {"local-port", required_argument, 0, 'l'},
        {"debug", no_argument, 0, 'd'},
        {"help", no_argument, 0, 'h'},
        {0, 0, 0, 0}
    };

    int opt;
    while ((opt = getopt_long(argc, argv, "l:dh", long_options, NULL)) != -1) {
        switch (opt) {
            case 'l':
                config->local_port = strdup(optarg);
                break;
            case 'd':
                config->debug = 1;
                break;
            case 'h':
            default:
                print_usage(argv[0]);
                return 0;
        }
    }

    return 1;
}

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

    char *host_part = strndup(hostname, colon_pos - hostname);
    char *at_pos = strchr(host_part, '@');

    if (at_pos) {
        *at_pos = '\0';
        // Skip user part for client_id extraction
    }

    char *dot_pos = strchr(host_part, '.');
    if (!dot_pos) {
        fprintf(stderr, "Error: Invalid hostname format. Expected client.domain format\n");
        free(host_part);
        return 0;
    }

    *dot_pos = '\0';
    *client_id = strdup(host_part);
    *wssshd_host = strdup(dot_pos + 1);
    *wssshd_port = DEFAULT_PORT;  // SCP uses -P option, not hostname suffix

    free(host_part);
    return 1;
}

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

int websocket_handshake(int sock, 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(sock, request, strlen(request), 0) < 0) {
        perror("WebSocket handshake send failed");
        return 0;
    }

    // Read response
    bytes_read = recv(sock, response, sizeof(response) - 1, 0);
    if (bytes_read <= 0) {
        perror("WebSocket handshake recv failed");
        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(int sock, const char *type, const char *client_id) {
    char message[1024];

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

    if (send(sock, message, strlen(message), 0) < 0) {
        perror("Send failed");
        return 0;
    }

    return 1;
}

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;

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

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

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

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

    // 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 (SSL_connect(ssl) <= 0) {
        ERR_print_errors_fp(stderr);
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return 0;
    }

    // Perform WebSocket handshake
    if (!websocket_handshake(sock, wssshd_host, wssshd_port, "/")) {
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return 0;
    }

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

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

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

    buffer[bytes_read] = '\0';

    if (debug) {
        printf("[DEBUG] Received: %s\n", buffer);
    }

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

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

    // Keep the connection alive for data transfer
    // In a real implementation, we'd need to handle bidirectional data transfer
    // For now, just keep the connection open

    // Cleanup (this would normally be handled differently)
    SSL_free(ssl);
    SSL_CTX_free(ssl_ctx);
    close(sock);

    return 1;
}

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

    if (!parse_args(argc, argv, &config)) {
        return 1;
    }

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

    // Need at least two arguments (source and destination)
    if (optind + 1 >= argc) {
        fprintf(stderr, "Error: Source and destination required\n");
        print_usage(argv[0]);
        return 1;
    }

    // Find SCP port from -P option
    int scp_port = DEFAULT_PORT;
    for (int i = optind; i < argc; i++) {
        if (strcmp(argv[i], "-P") == 0 && i + 1 < argc) {
            scp_port = atoi(argv[i + 1]);
            break;
        }
    }

    // Find host in arguments
    char *client_id = NULL;
    char *wssshd_host = NULL;
    int wssshd_port = scp_port;  // Use the -P port as wssshd port

    for (int i = optind; i < argc; i++) {
        if (strchr(argv[i], ':') && !strchr(argv[i], '=')) {
            // This looks like host:path
            char *colon_pos = strchr(argv[i], ':');
            if (colon_pos) {
                char *host_part = strndup(argv[i], colon_pos - argv[i]);
                if (parse_hostname(argv[i], &client_id, &wssshd_host, &wssshd_port)) {
                    break;
                }
                free(host_part);
            }
        }
    }

    if (!client_id || !wssshd_host) {
        fprintf(stderr, "Error: Could not determine target host\n");
        return 1;
    }

    if (config.debug) {
        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);
        return 1;
    }

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

    // Setup tunnel
    if (!setup_tunnel(wssshd_host, wssshd_port, client_id, local_port, config.debug)) {
        fprintf(stderr, "Error: Failed to establish tunnel\n");
        free(client_id);
        free(wssshd_host);
        return 1;
    }

    // Prepare SCP arguments
    char port_str[16];
    snprintf(port_str, sizeof(port_str), "%d", local_port);

    // Build SCP command
    char **scp_args = malloc((argc - optind + 4) * sizeof(char *));
    if (!scp_args) {
        perror("Memory allocation failed");
        free(client_id);
        free(wssshd_host);
        return 1;
    }

    int arg_count = 0;
    scp_args[arg_count++] = "scp";
    scp_args[arg_count++] = "-P";
    scp_args[arg_count++] = port_str;

    // Process remaining arguments, replacing host:path with localhost:path
    for (int i = optind; i < argc; i++) {
        if (strchr(argv[i], ':') && !strchr(argv[i], '=')) {
            // This is host:path, replace with localhost:path
            char *colon_pos = strchr(argv[i], ':');
            if (colon_pos) {
                char *path_part = colon_pos + 1;
                char *new_arg = malloc(strlen(path_part) + 10); // "localhost:" + path
                if (new_arg) {
                    sprintf(new_arg, "localhost:%s", path_part);
                    scp_args[arg_count++] = new_arg;
                }
            }
        } else {
            scp_args[arg_count++] = argv[i];
        }
    }
    scp_args[arg_count] = NULL;

    if (config.debug) {
        printf("[DEBUG] SCP command:");
        for (int i = 0; scp_args[i]; i++) {
            printf(" %s", scp_args[i]);
        }
        printf("\n");
    }

    // Execute SCP
    execvp("scp", scp_args);

    // If we reach here, execvp failed
    perror("execvp failed");

    // Free allocated memory
    for (int i = 3; i < arg_count; i++) {
        if (strstr(scp_args[i], "localhost:")) {
            free(scp_args[i]);
        }
    }
    free(scp_args);
    free(client_id);
    free(wssshd_host);
    free(config.local_port);

    return 1;
}