#!/usr/bin/env python3
"""
WebSocket SCP (wsscp)
SCP 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 SCP 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_scp(server_ip, server_port, client_id, local_port, scp_args):
    """Connect to wssshd and run SCP"""
    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 SCP with modified arguments
                scp_cmd = ['scp'] + scp_args + ['-P', str(local_port)]
                if debug: print(f"[DEBUG] Launching: {' '.join(scp_cmd)}")

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

                # Wait for SCP 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 SCP (wsscp)', 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 'scp' to the remaining arguments
    modified_args = ['scp'] + remaining

    # Find host arguments and -P port option in modified args (SCP can have multiple host:paths)
    hosts = []
    scp_port = None
    for i, arg in enumerate(modified_args):
        if arg == '-P' and i + 1 < len(modified_args):
            try:
                scp_port = int(modified_args[i + 1])
            except ValueError:
                pass
        elif not arg.startswith('-') and ':' in arg:
            # Extract host from host:path format
            host_part = arg.split(':', 1)[0]
            if '@' in host_part:
                host = host_part.split('@', 1)[1]
            else:
                host = host_part
            hosts.append(host)

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

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

    # Use the first host for parsing
    host = hosts[0]

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

    if debug:
        print(f"[DEBUG] Hosts: {hosts}")
        print(f"[DEBUG] Client ID: {client_id}")
        print(f"[DEBUG] WSSSHD Host: {wssshd_host}")
        print(f"[DEBUG] WSSSHD Port: {scp_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 = []
    for arg in modified_args:
        if ':' in arg:
            host_part, path_part = arg.split(':', 1)
            if '@' in host_part:
                user, host_in_arg = host_part.split('@', 1)
                if host_in_arg == host:  # Only modify the first host
                    final_args.append(f"{user}@localhost:{path_part}")
                else:
                    final_args.append(arg)
            else:
                if host_part == host:  # Only modify the first host
                    final_args.append(f"localhost:{path_part}")
                else:
                    final_args.append(arg)
        else:
            final_args.append(arg)

    # Add port argument for local tunnel
    final_args.extend(['-P', str(local_port)])

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

    # Run the async SCP wrapper
    exit_code = asyncio.run(run_scp(
        wssshd_host,  # Use parsed wssshd_host as server_ip
        scp_port,     # Use -P port as wssshd_port
        client_id,
        local_port,
        final_args
    ))

    sys.exit(exit_code)

if __name__ == '__main__':
    main()