/*
 * WSSSH Library - Tunnel management
 *
 * 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/>.
 */

#ifndef TUNNEL_H
#define TUNNEL_H

#include <openssl/ssl.h>
#include <pthread.h>
#include "data_messages.h"

// Frame buffer for wsscp
typedef struct {
    char *buffer;
    size_t size;
    size_t used;
} frame_buffer_t;

// Tunnel structure
typedef struct {
    int local_sock;  // Local TCP connection socket (listening socket for wsssh, connection socket for wsscp)
    int sock;  // SSH client connection socket (for wsssh) or SSH server connection socket (for wssshc)
    char request_id[37];  // UUID string
    int active;
    int broken;  // Flag to indicate if the tunnel is broken (not normal closure)
    SSL *ssl;  // WebSocket SSL connection
    frame_buffer_t *outgoing_buffer;  // Buffer for data to send to local socket (wsscp only)
    frame_buffer_t *incoming_buffer;  // Buffer for incoming data before connection is established
    int server_version_sent;  // Flag to indicate if server version was sent early
    pthread_t forward_thread;  // Thread ID for the forwarding thread

    // Keep-alive statistics and timing
    time_t last_keepalive_sent;     // Last time we sent a keep-alive
    time_t last_keepalive_received; // Last time we received a keep-alive
    unsigned long long total_bytes_sent;     // Total bytes sent through tunnel
    unsigned long long total_bytes_received; // Total bytes received through tunnel
    unsigned long long bytes_last_period;    // Bytes transferred in last 30-second period
    time_t last_stats_reset;         // When we last reset the period stats
    time_t start_time;               // When the tunnel was created
    char service[32];                // Service name (e.g., "ssh", "scp")
    char client_id[64];              // Client identifier
    int bin;  // Binary mode flag - if true, transmit data as binary instead of hex
    wsssh_encoding_t encoding;  // Data encoding mode

    // Reliable transmission
    retransmission_buffer_t *retransmission_buffer;  // Buffer for reliable message retransmission
} tunnel_t;

// Service configuration structure
typedef struct {
    char *name;
    int tunnel_port;
    char *tunnel_host;
    char *command;
    char *proto;  // "tcp" or "udp", default "tcp"
    int active;   // Whether this service is active, default true
} service_config_t;

// Thread arguments
typedef struct {
    SSL *ssl;
    tunnel_t *tunnel;
    int debug;
} thread_args_t;

// Global variables
extern tunnel_t *active_tunnel;  // For backward compatibility
extern tunnel_t **active_tunnels;
extern int active_tunnels_count;
extern int active_tunnels_capacity;
extern pthread_mutex_t tunnel_mutex;
extern pthread_mutex_t ssl_mutex;  // For synchronizing SSL operations
extern volatile int global_shutdown;  // Global flag to signal threads to stop

// Function declarations
frame_buffer_t *frame_buffer_init(void);
void frame_buffer_free(frame_buffer_t *fb);
int frame_buffer_resize(frame_buffer_t *fb, size_t new_size);
int frame_buffer_append(frame_buffer_t *fb, const char *data, size_t len);
int frame_buffer_consume(frame_buffer_t *fb, size_t len);

// Helper functions for managing multiple tunnels
tunnel_t *find_tunnel_by_request_id(const char *request_id);
int add_tunnel(tunnel_t *tunnel);
void remove_tunnel(const char *request_id);

void *forward_tcp_to_ws(void *arg);
void *forward_ws_to_local(void *arg);
void *forward_ws_to_ssh_server(void *arg);
void *tunnel_thread(void *arg);
void handle_tunnel_request(SSL *ssl, const char *request_id, int debug, const char *ssh_host, int ssh_port);
void handle_tunnel_request_with_enc(SSL *ssl, const char *request_id, int debug, const char *ssh_host, int ssh_port, wsssh_encoding_t encoding);
void handle_tunnel_request_with_service(SSL *ssl, const char *request_id, service_config_t *service, int debug);
void handle_tunnel_request_with_service_and_enc(SSL *ssl, const char *request_id, service_config_t *service, int debug, wsssh_encoding_t encoding);
void handle_tunnel_data(SSL *ssl, const char *request_id, const char *data_hex, int debug);
void handle_tunnel_close(SSL *ssl, const char *request_id, int debug);
void send_tunnel_close(SSL *ssl, const char *request_id, int debug);
void handle_tunnel_keepalive(SSL *ssl, const char *request_id, unsigned long long total_bytes, double rate_bps, int debug);
void handle_tunnel_keepalive_ack(SSL *ssl, const char *request_id, int debug);
void handle_tunnel_ack(SSL *ssl, const char *request_id, uint32_t frame_id, int debug);
void handle_tunnel_ko(SSL *ssl, const char *request_id, uint32_t frame_id, int debug);
void send_tunnel_keepalive(SSL *ssl, tunnel_t *tunnel, int debug);
void check_keepalive_timeouts(int debug);
void cleanup_tunnel(int debug);
int reconnect_websocket(tunnel_t *tunnel, const char *wssshd_host, int wssshd_port, const char *client_id, const char *request_id, int debug);

// CPU affinity functions
void init_cpu_affinity(void);
void set_thread_cpu_affinity(pthread_t thread);

// Structure for tunnel setup result
typedef struct {
    int listen_sock;
    SSL *ssl;
    SSL_CTX *ssl_ctx;
    char request_id[37];
} tunnel_setup_result_t;

// Function declarations
tunnel_setup_result_t setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id, int local_port, int debug, int use_buffer, const char *tunnel_host, wsssh_encoding_t encoding, int send_tunnel_request_immediately);

#endif // TUNNEL_H