#!/usr/bin/env python3
"""
WebSocket SSH (wsssh)
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/>.
"""

import argparse
import asyncio
import ssl
import websockets
import json
import socket
import subprocess
import sys
import uuid
import configparser
import os
import signal

debug = False

async def handle_tunnel(ws, local_port, request_id, ready_event=None):
    """Handle tunnel data forwarding"""
    try:
        # Open local TCP listener
        server = await asyncio.start_server(
            lambda r, w: handle_local_connection(r, w, ws, request_id),
            'localhost', local_port
        )

        if debug: print(f"[DEBUG] Listening on localhost:{local_port}")

        # Signal that server is ready
        if ready_event:
            ready_event.set()

        async with server:
            await server.serve_forever()
    except Exception as e:
        if debug: print(f"[DEBUG] Tunnel handler error: {e}")

async def handle_local_connection(reader, writer, ws, request_id):
    """Handle connections from local SSH client"""
    try:
        # Forward data between local connection and WebSocket
        async def forward_to_ws():
            try:
                while True:
                    data = await reader.read(1024)
                    if not data:
                        break
                    await ws.send(json.dumps({
                        "type": "tunnel_data",
                        "request_id": request_id,
                        "data": data.hex()  # Use hex encoding for binary data
                    }))
            except Exception as e:
                if debug: print(f"[DEBUG] Local to WS error: {e}")

        async def forward_from_ws():
            try:
                async for message in ws:
                    data = json.loads(message)
                    if data.get('type') == 'tunnel_data' and data.get('request_id') == request_id:
                        # Decode and forward data
                        tunnel_data = bytes.fromhex(data['data'])
                        writer.write(tunnel_data)
                        await writer.drain()
                    elif data.get('type') == 'tunnel_close' and data.get('request_id') == request_id:
                        break
            except Exception as e:
                if debug: print(f"[DEBUG] WS to local error: {e}")

        await asyncio.gather(forward_to_ws(), forward_from_ws())
    except Exception as e:
        if debug: print(f"[DEBUG] Local connection error: {e}")
    finally:
        writer.close()
        await writer.wait_closed()

async def run_ssh(server_ip, server_port, client_id, local_port, ssh_args, interval=30, shutdown_event=None):
    """Connect to wssshd and run SSH with reconnection support"""
    uri = f"wss://{server_ip}:{server_port}"
    ssl_context = ssl.create_default_context()
    ssl_context.check_hostname = False
    ssl_context.verify_mode = ssl.CERT_NONE

    request_id = str(uuid.uuid4())

    # Initial connection with retry logic
    websocket = None
    max_initial_attempts = 3
    initial_attempt = 0

    while initial_attempt < max_initial_attempts and websocket is None:
        try:
            if debug and initial_attempt > 0:
                print(f"[DEBUG] Initial connection attempt {initial_attempt + 1}/{max_initial_attempts}")
            websocket = await websockets.connect(uri, ssl=ssl_context)
            if debug: print("[DEBUG] Initial WebSocket connection established")
        except Exception as e:
            initial_attempt += 1
            if initial_attempt < max_initial_attempts:
                if debug: print(f"[DEBUG] Initial connection failed, waiting {interval} seconds...")
                await asyncio.sleep(interval)
            else:
                print(f"Connection failed after {max_initial_attempts} attempts: {e}")
                return 1

    try:
        # Set up ping/pong handlers for explicit handling
        async def ping_handler(payload):
            if debug: print(f"[DEBUG] Received ping: {payload}")
            # Pong is sent automatically by websockets library

        async def pong_handler(payload):
            if debug: print(f"[DEBUG] Received pong: {payload}")

        websocket.ping_handler = ping_handler
        websocket.pong_handler = pong_handler

        # Request tunnel
        await websocket.send(json.dumps({
            "type": "tunnel_request",
            "client_id": client_id,
            "request_id": request_id
        }))

        # Wait for acknowledgment
        response = await websocket.recv()
        data = json.loads(response)
        if data.get('type') == 'tunnel_ack':
            if debug: print(f"[DEBUG] Tunnel request acknowledged: {request_id}")

            # Start tunnel handler (opens listening port)
            ready_event = asyncio.Event()
            tunnel_task = asyncio.create_task(handle_tunnel(websocket, local_port, request_id, ready_event))

            # Wait for tunnel to be ready (listening port opened)
            await ready_event.wait()

            # Launch SSH with modified arguments
            ssh_cmd = ['ssh'] + ssh_args
            if debug: print(f"[DEBUG] Launching: {' '.join(ssh_cmd)}")

            # Run SSH process
            process = await asyncio.create_subprocess_exec(
                *ssh_cmd,
                stdin=sys.stdin,
                stdout=sys.stdout,
                stderr=sys.stderr
            )

            # Monitor both SSH process and WebSocket connection
            tunnel_active = True
            while tunnel_active:
                # Check if SSH process is still running
                if process.returncode is not None:
                    if debug: print(f"[DEBUG] SSH process finished with code: {process.returncode}")
                    tunnel_active = False
                    break

                # Check for shutdown signal
                if shutdown_event and shutdown_event.is_set():
                    if debug: print("[DEBUG] Shutdown signal received, terminating tunnel")
                    tunnel_active = False
                    break

                # Check WebSocket connection and attempt reconnection if needed
                try:
                    # Try to send a small message to test connection
                    await asyncio.wait_for(websocket.ping(), timeout=5.0)
                except Exception as e:
                    if debug: print(f"[DEBUG] WebSocket connection lost, attempting reconnection...")

                    # Attempt WebSocket reconnection
                    reconnect_attempts = 0
                    max_reconnect_attempts = 3
                    reconnected = False

                    while reconnect_attempts < max_reconnect_attempts and not reconnected:
                        # Check for shutdown signal during reconnection
                        if shutdown_event and shutdown_event.is_set():
                            if debug: print("[DEBUG] Shutdown signal received during reconnection")
                            tunnel_active = False
                            break

                        try:
                            if debug: print(f"[DEBUG] WebSocket reconnection attempt {reconnect_attempts + 1}/{max_reconnect_attempts}")

                            # Create new WebSocket connection
                            new_websocket = await websockets.connect(uri, ssl=ssl_context)
                            new_websocket.ping_handler = ping_handler
                            new_websocket.pong_handler = pong_handler

                            # Re-request tunnel
                            await new_websocket.send(json.dumps({
                                "type": "tunnel_request",
                                "client_id": client_id,
                                "request_id": request_id
                            }))

                            # Wait for acknowledgment
                            response = await new_websocket.recv()
                            data = json.loads(response)
                            if data.get('type') == 'tunnel_ack':
                                # Update tunnel task with new WebSocket
                                tunnel_task.cancel()
                                tunnel_task = asyncio.create_task(handle_tunnel(new_websocket, local_port, request_id, None))
                                websocket = new_websocket
                                reconnected = True
                                if debug: print("[DEBUG] WebSocket reconnection successful")
                            else:
                                await new_websocket.close()
                        except Exception as reconnect_error:
                            reconnect_attempts += 1
                            if reconnect_attempts < max_reconnect_attempts:
                                if debug: print(f"[DEBUG] WebSocket reconnection failed, waiting 1 second...")
                                await asyncio.sleep(1)  # Fast reconnection for WebSocket
                            else:
                                if debug: print(f"[DEBUG] All reconnection attempts failed: {reconnect_error}")
                                tunnel_active = False
                                break

                if tunnel_active:
                    await asyncio.sleep(2)  # Check every 2 seconds

            # Wait for SSH to complete if still running
            if process.returncode is None:
                await process.wait()

            # Close tunnel
            try:
                await websocket.send(json.dumps({
                    "type": "tunnel_close",
                    "request_id": request_id
                }))
            except Exception:
                pass  # WebSocket might already be closed

            if tunnel_task and not tunnel_task.done():
                tunnel_task.cancel()

        elif data.get('type') == 'tunnel_error':
            print(f"Error: {data.get('error', 'Unknown error')}")
            return 1

    except Exception as e:
        print(f"Connection failed: {e}")
        return 1
    finally:
        if websocket:
            try:
                await websocket.close()
            except Exception:
                pass

    return 0

def parse_hostname(hostname, config_domain=None):
    """Parse hostname to extract CLIENT_ID and WSSSHD_HOST"""
    # Split by dots to get client_id and wssshd_host
    parts = hostname.split('.')
    if len(parts) >= 2:
        client_id = parts[0]
        wssshd_host = '.'.join(parts[1:])
    else:
        # No domain, assume whole hostname is client_id
        client_id = hostname
        wssshd_host = config_domain if config_domain else 'localhost'  # Default fallback

    return client_id, wssshd_host

def main():
    # Read config file
    config = configparser.ConfigParser()
    config_file = os.path.expanduser("~/.config/wsssh/wsssh.conf")
    config_port = None
    config_domain = None
    if os.path.exists(config_file):
        config.read(config_file)
        if 'default' in config:
            if 'port' in config['default']:
                try:
                    config_port = int(config['default']['port'])
                except ValueError:
                    pass
            if 'domain' in config['default']:
                config_domain = config['default']['domain']

    parser = argparse.ArgumentParser(description='WebSocket SSH (wsssh)', add_help=False)
    parser.add_argument('--local-port', type=int, default=0, help='Local port for tunnel (0 = auto)')
    parser.add_argument('--interval', type=int, default=30, help='Connection retry interval in seconds (default: 30)')
    parser.add_argument('--debug', action='store_true', help='Enable debug output')

    # Parse our arguments first
    args, remaining = parser.parse_known_args()

    global debug
    debug = args.debug

    # Don't prepend 'ssh' - we'll handle it in execvp
    modified_args = remaining

    # Find the host argument and -p port option in modified args
    host = None
    ssh_port = None
    for i, arg in enumerate(modified_args):
        if arg == '-p' and i + 1 < len(modified_args):
            try:
                ssh_port = int(modified_args[i + 1])
            except ValueError:
                pass
        elif not arg.startswith('-') and i > 0 and modified_args[i-1] in ['-h', '--host']:
            host = arg
            break
        elif not arg.startswith('-') and '@' in arg:
            # Handle user@host format
            host = arg.split('@', 1)[1]
            break
        elif not arg.startswith('-') and i == 0:
            # First non-option argument might be host
            host = arg
            break

    if not host:
        print("Error: Could not determine target host")
        sys.exit(1)

    if not ssh_port:
        if config_port:
            ssh_port = config_port
        else:
            print("Error: Could not determine wssshd port from -p option or config file")
            sys.exit(1)

    # Parse hostname to extract client_id and wssshd_host
    client_id, wssshd_host = parse_hostname(host, config_domain)

    if debug:
        print(f"[DEBUG] Host: {host}")
        print(f"[DEBUG] Client ID: {client_id}")
        print(f"[DEBUG] WSSSHD Host: {wssshd_host}")
        print(f"[DEBUG] WSSSHD Port: {ssh_port}")

    # Find available local port
    if args.local_port == 0:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(('localhost', 0))
            local_port = s.getsockname()[1]
    else:
        local_port = args.local_port

    if debug: print(f"[DEBUG] Using local port: {local_port}")

    # Modify modified_args to use localhost:local_port
    final_args = []
    skip_next = False
    for i, arg in enumerate(modified_args):
        if skip_next:
            skip_next = False
            continue
        if arg in ['-h', '--host', '-p', '--port']:
            skip_next = True
            continue
        elif '@' in arg and arg.split('@', 1)[1] == host:
            # Replace user@host with user@localhost
            user = arg.split('@', 1)[0]
            final_args.append(f"{user}@localhost")
            continue
        elif arg == host:
            final_args.append('localhost')
            continue
        else:
            final_args.append(arg)

    # Add StrictHostKeyChecking=no and port argument for local tunnel
    final_args.extend(['-o', 'StrictHostKeyChecking=no', '-o', 'UserKnownHostsFile=/dev/null', '-o', 'LogLevel ERROR', '-p', str(local_port)])

    if debug: print(f"[DEBUG] Final SSH args: {final_args}")

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

    # Run the async SSH wrapper
    exit_code = asyncio.run(run_ssh(
        wssshd_host,  # Use parsed wssshd_host as server_ip
        ssh_port,     # Use -p port as wssshd_port
        client_id,
        local_port,
        final_args,
        args.interval,  # Pass the interval parameter
        shutdown_event  # Pass the shutdown event
    ))

    sys.exit(exit_code)

if __name__ == '__main__':
    main()