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

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):
    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:
        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 main():
    parser = argparse.ArgumentParser(description='WebSocket SSH Client (wssshc)')
    parser.add_argument('--server-ip', required=True, help='Server IP address')
    parser.add_argument('--port', type=int, required=True, help='Server port')
    parser.add_argument('--id', required=True, help='Client ID')
    parser.add_argument('--password', required=True, help='Registration password')
    parser.add_argument('--interval', type=int, default=30, help='Reconnect interval in seconds (default: 30)')
    parser.add_argument('--debug', action='store_true', help='Enable debug output')

    args = parser.parse_args()
    global debug
    debug = args.debug

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

if __name__ == '__main__':
    main()