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

debug = False

async def handle_tunnel(ws, local_port, request_id):
    """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}")

        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):
    """Connect to wssshd and run SSH"""
    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())

    try:
        async with websockets.connect(uri, ssl=ssl_context) as websocket:
            # 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
                tunnel_task = asyncio.create_task(handle_tunnel(websocket, local_port, request_id))

                # Launch SSH with modified arguments
                ssh_cmd = ['ssh'] + ssh_args + ['-p', str(local_port), 'localhost']
                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
                )

                # Wait for SSH to complete
                await process.wait()

                # Close tunnel
                await websocket.send(json.dumps({
                    "type": "tunnel_close",
                    "request_id": request_id
                }))

                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

    return 0

def parse_hostname(hostname):
    """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 = 'localhost'  # Default fallback

    return client_id, wssshd_host

def main():
    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('--debug', action='store_true', help='Enable debug output')

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

    global debug
    debug = args.debug

    # Automatically prepend 'ssh' to the remaining arguments
    modified_args = ['ssh'] + 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 == 1:  # Skip 'ssh' at index 0
            # First non-option argument after 'ssh' might be host
            host = arg
            break

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

    if not ssh_port:
        print("Error: Could not determine wssshd port from -p option")
        sys.exit(1)

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

    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 port argument for local tunnel
    final_args.extend(['-p', str(local_port)])

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

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

    sys.exit(exit_code)

if __name__ == '__main__':
    main()