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

debug = False

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':
                        print(f"Tunnel request received: {data['request_id']}")
                        # Client would handle tunnel setup here
                    elif data.get('type') == 'tunnel_data':
                        print(f"Tunnel data received: {len(data.get('data', ''))} bytes")
                        # Client would forward data here
                    elif data.get('type') == 'tunnel_close':
                        print(f"Tunnel closed: {data['request_id']}")
                        # Client would close tunnel here
        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()