#!/usr/bin/env python3
"""
WebSocket SSH Client (wssshc)
Connects to wssshd server and registers as a client.

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/>.
"""

import argparse
import asyncio
import ssl
import websockets
import json
import socket
import configparser
import os
import signal

debug = False

# Active tunnels: request_id -> {'reader': tcp_reader, 'writer': tcp_writer}
active_tunnels = {}

async def forward_tcp_to_ws(tcp_reader, websocket, request_id):
    try:
        while True:
            data = await tcp_reader.read(1024)
            if not data:
                break
            if debug: print(f"[DEBUG] Sending tunnel_response, data len: {len(data)}, data: {data[:50].decode('utf-8', errors='replace') if data else 'empty'}, hex: {data.hex()[:100]}...")
            await websocket.send(json.dumps({
                "type": "tunnel_response",
                "request_id": request_id,
                "data": data.hex()
            }))
    except Exception as e:
        if debug: print(f"[DEBUG] TCP to WS error: {e}")
    finally:
        # Cleanup tunnel
        if request_id in active_tunnels:
            del active_tunnels[request_id]
        if debug: print(f"[DEBUG] Tunnel {request_id} closed")

async def connect_to_server(server_ip, port, client_id, password, interval, shutdown_event=None):
    uri = f"wss://{server_ip}:{port}"
    ssl_context = ssl.create_default_context()
    ssl_context.check_hostname = False
    ssl_context.verify_mode = ssl.CERT_NONE

    while True:
        # Check for shutdown signal
        if shutdown_event and shutdown_event.is_set():
            if debug: print("[DEBUG] Shutdown signal received, exiting")
            break

        try:
            async with websockets.connect(uri, ssl=ssl_context) as websocket:
                # Register
                await websocket.send(json.dumps({"type": "register", "id": client_id, "password": password}))
                print(f"Connected and registered as {client_id}")

                # Wait for registration acknowledgment
                response = await websocket.recv()
                data = json.loads(response)
                if data.get('type') == 'registered':
                    print(f"Successfully registered as {client_id}")
                elif data.get('type') == 'registration_error':
                    print(f"Registration failed: {data.get('error', 'Unknown error')}")
                    return  # Exit and retry

                async for message in websocket:
                    if debug: print(f"[DEBUG] WebSocket message: {message[:100]}...")
                    data = json.loads(message)
                    if data.get('type') == 'tunnel_request':
                        request_id = data['request_id']
                        if debug: print(f"[DEBUG] Tunnel request received: {request_id}")
                        try:
                            # Connect to local SSH server
                            tcp_reader, tcp_writer = await asyncio.open_connection('localhost', 22)
                            active_tunnels[request_id] = {'reader': tcp_reader, 'writer': tcp_writer}

                            # Start forwarding from TCP to WS
                            asyncio.create_task(forward_tcp_to_ws(tcp_reader, websocket, request_id))

                            if debug: print(f"[DEBUG] Connected to local SSH server for tunnel {request_id}")
                        except Exception as e:
                            if debug: print(f"[DEBUG] Tunnel setup error: {e}")
                    elif data.get('type') == 'tunnel_data':
                        request_id = data['request_id']
                        if request_id in active_tunnels:
                            # Forward data to local SSH
                            tunnel_data = bytes.fromhex(data['data'])
                            if debug: print(f"[DEBUG] Received tunnel_data, hex len: {len(data['data'])}, hex: {data['data'][:100]}..., decoded: {tunnel_data[:50].decode('utf-8', errors='replace') if tunnel_data else 'empty'}")
                            active_tunnels[request_id]['writer'].write(tunnel_data)
                            await active_tunnels[request_id]['writer'].drain()
                    elif data.get('type') == 'tunnel_close':
                        request_id = data['request_id']
                        if request_id in active_tunnels:
                            # Close tunnel
                            active_tunnels[request_id]['writer'].close()
                            await active_tunnels[request_id]['writer'].wait_closed()
                            del active_tunnels[request_id]
                            if debug: print(f"[DEBUG] Tunnel {request_id} closed")
        except Exception as e:
            print(f"Connection failed: {e}, retrying in {interval} seconds")
            await asyncio.sleep(interval)

def load_config_file(config_path):
    """Load configuration from a specific file path"""
    if not os.path.exists(config_path):
        return {}

    config = configparser.ConfigParser()
    try:
        config.read(config_path)
        if 'wssshc' in config:
            section = config['wssshc']
            defaults = {}
            for key in ['password', 'id']:
                if key in section:
                    defaults[key] = section[key]
            if 'server-ip' in section:
                defaults['server_ip'] = section['server-ip']
            elif 'domain' in section:
                defaults['server_ip'] = section['domain']
            if 'port' in section:
                defaults['port'] = int(section['port'])
            if 'interval' in section:
                defaults['interval'] = int(section['interval'])
            return defaults
    except Exception as e:
        if debug: print(f"[DEBUG] Error reading config file {config_path}: {e}")

    return {}

def main():
    # Configuration hierarchy:
    # 1. System config (/etc/wssshc.conf)
    # 2. User config (~/.config/wsssh/wssshc.conf)
    # 3. Command line arguments (highest priority)

    # Load system config first (lowest priority)
    system_config = '/etc/wssshc.conf'
    defaults = load_config_file(system_config)

    # Load user config (medium priority, overrides system)
    user_config = os.path.expanduser('~/.config/wsssh/wssshc.conf')
    user_defaults = load_config_file(user_config)
    defaults.update(user_defaults)

    # Set default values if not set by config files
    defaults.setdefault('interval', 30)
    defaults.setdefault('port', 9898)

    # Parse command line arguments (highest priority)
    parser = argparse.ArgumentParser(description='WebSocket SSH Client (wssshc)')
    parser.set_defaults(**defaults)
    parser.add_argument('--config', help='Configuration file path (overrides default hierarchy)')
    parser.add_argument('--server-ip', help='Server IP address')
    parser.add_argument('--port', type=int, help='Server port')
    parser.add_argument('--id', help='Client ID')
    parser.add_argument('--password', help='Registration password')
    parser.add_argument('--interval', type=int, help='Reconnect interval in seconds (default: 30)')
    parser.add_argument('--debug', action='store_true', help='Enable debug output')

    # Parse just the config argument first to check for custom config file
    temp_parser = argparse.ArgumentParser(add_help=False)
    temp_parser.add_argument('--config')
    temp_args, remaining = temp_parser.parse_known_args()

    # If --config is specified, load that file instead of the hierarchy
    if temp_args.config:
        custom_defaults = load_config_file(temp_args.config)
        defaults.update(custom_defaults)
        parser.set_defaults(**defaults)

    args = parser.parse_args()

    args = parser.parse_args()

    # Check required arguments
    if not args.server_ip:
        parser.error('--server-ip is required')
    if not args.id:
        parser.error('--id is required')
    if not args.password:
        parser.error('--password is required')

    global debug
    debug = args.debug

    # Set up signal handling for clean exit
    shutdown_event = asyncio.Event()

    def signal_handler(signum, frame):
        if debug: print(f"[DEBUG] Received signal {signum}, initiating shutdown")
        shutdown_event.set()

    # Register signal handler for SIGINT (Ctrl+C)
    signal.signal(signal.SIGINT, signal_handler)

    asyncio.run(connect_to_server(args.server_ip, args.port, args.id, args.password, args.interval, shutdown_event))

if __name__ == '__main__':
    main()