#!/usr/bin/env python3
"""
SSH-over-WebSocket tunnel client. No dependencies beyond stdlib.

Usage:
  python3 tunnel.py <host> -d <destination> [-p PORT] [--dns DNS_SERVER]

Destinations:
  parspack   185.208.174.152:22  (default)
  mequ       188.213.68.133:2022
  alipi      10.66.66.2:22

Examples:
  python3 tunnel.py nginx-0651fe8398-manpache.apps.ir-central1.arvancaas.ir -d parspack -p 1212
  python3 tunnel.py nginx-0651fe8398-manpache.apps.ir-central1.arvancaas.ir -d mequ -p 1213
  python3 tunnel.py nginx-0651fe8398-manpache.apps.ir-central1.arvancaas.ir -d alipi -p 1214 --dns 8.8.8.8
  ssh -p 1212 user@localhost
"""

import asyncio
import base64
import hashlib
import os
import socket
import ssl
import struct
import sys


DESTINATIONS = {
    "parspack": "185.208.174.152:22",
    "mequ":     "188.213.68.133:2022",
    "alipi":    "10.66.66.2:22",
}


def ws_make_frame(opcode, data, masked=True):
    """Build a WebSocket frame (client frames must be masked)."""
    frame = bytes([0x80 | opcode])
    if len(data) < 126:
        frame += bytes([0x80 | len(data)] if masked else [len(data)])
    elif len(data) < 65536:
        frame += struct.pack("!BH", (0x80 | 126) if masked else 126, len(data))
    else:
        frame += struct.pack("!BQ", (0x80 | 127) if masked else 127, len(data))
    if masked:
        mask = os.urandom(4)
        frame += mask
        data = bytes(b ^ mask[i % 4] for i, b in enumerate(data))
    return frame + data


async def ws_read_frame(reader):
    """Read one WebSocket frame, return (opcode, payload)."""
    b0, b1 = struct.unpack("!BB", await reader.readexactly(2))
    opcode = b0 & 0x0F
    masked = b1 & 0x80
    length = b1 & 0x7F
    if length == 126:
        length = struct.unpack("!H", await reader.readexactly(2))[0]
    elif length == 127:
        length = struct.unpack("!Q", await reader.readexactly(8))[0]
    mask = await reader.readexactly(4) if masked else None
    data = await reader.readexactly(length)
    if mask:
        data = bytes(b ^ mask[i % 4] for i, b in enumerate(data))
    return opcode, data


async def ws_handshake(reader, writer, host, dest):
    """Perform WebSocket handshake, return True on success."""
    key = base64.b64encode(os.urandom(16)).decode()
    req = (f"GET /tunnel/{dest} HTTP/1.1\r\n"
           f"Host: {host}\r\n"
           f"Upgrade: websocket\r\n"
           f"Connection: Upgrade\r\n"
           f"Sec-WebSocket-Key: {key}\r\n"
           f"Sec-WebSocket-Version: 13\r\n"
           f"\r\n")
    writer.write(req.encode())
    await writer.drain()

    status_line = (await reader.readline()).decode().strip()
    if "101" not in status_line:
        # Read remaining response
        body = b""
        while True:
            line = await reader.readline()
            if not line or line == b"\r\n":
                break
        rest = await asyncio.wait_for(reader.read(1024), timeout=1.0) if "404" in status_line else b""
        print(f"Handshake failed: {status_line}")
        if rest:
            print(f"  {rest.decode(errors='replace')}")
        return False

    # consume headers
    while True:
        line = (await reader.readline()).decode().strip()
        if not line:
            break

    return True


def resolve_with_dns(hostname, dns_server):
    """Resolve hostname using a specific DNS server via raw UDP query."""
    import random
    txn_id = random.randint(0, 65535)
    flags = 0x0100
    header = struct.pack("!HHHHHH", txn_id, flags, 1, 0, 0, 0)
    query = b""
    for label in hostname.split("."):
        query += bytes([len(label)]) + label.encode()
    query += b"\x00"
    query += struct.pack("!HH", 1, 1)

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(5)
    try:
        sock.sendto(header + query, (dns_server, 53))
        resp, _ = sock.recvfrom(1024)
    finally:
        sock.close()

    pos = 12
    while resp[pos] != 0:
        pos += resp[pos] + 1
    pos += 5

    ans_count = struct.unpack("!H", resp[6:8])[0]
    if ans_count == 0:
        raise RuntimeError(f"DNS {dns_server} returned no answers for {hostname}")

    if resp[pos] & 0xC0 == 0xC0:
        pos += 2
    else:
        while resp[pos] != 0:
            pos += resp[pos] + 1
        pos += 1

    rtype, rclass, ttl, rdlen = struct.unpack("!HHIH", resp[pos:pos+10])
    pos += 10
    if rtype == 1 and rdlen == 4:
        return socket.inet_ntoa(resp[pos:pos+4])
    raise RuntimeError(f"DNS {dns_server}: unexpected record type {rtype}")


async def handle_local_client(local_reader, local_writer, host, dest, use_tls, resolved_ip=None):
    """Handle one SSH connection by tunneling it over WebSocket."""
    ctx = None
    if use_tls:
        ctx = ssl.create_default_context()
        if resolved_ip:
            ctx.check_hostname = False
            ctx.verify_mode = ssl.CERT_NONE

    connect_host = resolved_ip or host
    try:
        ws_reader, ws_writer = await asyncio.open_connection(
            connect_host, 443 if use_tls else 80, ssl=ctx,
            server_hostname=host if resolved_ip else None)
    except Exception as e:
        print(f"Failed to connect to {host}: {e}")
        local_writer.close()
        return

    if not await ws_handshake(ws_reader, ws_writer, host, dest):
        local_writer.close()
        ws_writer.close()
        return

    print(f"  tunnel established → {host} → {dest} ({DESTINATIONS[dest]})")

    async def local_to_ws():
        try:
            while True:
                data = await local_reader.read(16384)
                if not data:
                    break
                ws_writer.write(ws_make_frame(0x2, data))
                await ws_writer.drain()
        except Exception:
            pass
        finally:
            ws_writer.write(ws_make_frame(0x8, struct.pack("!H", 1000), masked=True))
            try:
                await ws_writer.drain()
            except Exception:
                pass
            ws_writer.close()

    async def ws_to_local():
        try:
            while True:
                op, data = await ws_read_frame(ws_reader)
                if op == 0x8:
                    break
                if op == 0x9:
                    ws_writer.write(ws_make_frame(0xA, data))
                    await ws_writer.drain()
                    continue
                if op in (0x1, 0x2):
                    local_writer.write(data)
                    await local_writer.drain()
        except Exception:
            pass
        finally:
            local_writer.close()

    await asyncio.gather(local_to_ws(), ws_to_local())
    print(f"  tunnel closed")


async def main():
    import argparse
    parser = argparse.ArgumentParser(
        description="SSH over WebSocket tunnel",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="Destinations:\n" + "\n".join(f"  {k:10s} {v}" for k, v in DESTINATIONS.items())
    )
    parser.add_argument("host", help="WebSocket server hostname")
    parser.add_argument("-d", "--destination", default="parspack",
                        choices=DESTINATIONS.keys(),
                        help="SSH destination (default: parspack)")
    parser.add_argument("-p", "--port", type=int, default=1212,
                        help="Local port to listen on (default: 1212)")
    parser.add_argument("--dns", default=None,
                        help="DNS server to resolve hostname (e.g. 8.8.8.8)")
    parser.add_argument("--no-tls", action="store_true",
                        help="Use ws:// instead of wss://")
    args = parser.parse_args()

    use_tls = not args.no_tls
    resolved_ip = None
    dest = args.destination

    if args.dns:
        try:
            resolved_ip = resolve_with_dns(args.host, args.dns)
            print(f"Resolved {args.host} → {resolved_ip} (via {args.dns})")
        except Exception as e:
            print(f"DNS resolution failed: {e}")
            sys.exit(1)

    async def on_connect(reader, writer):
        addr = writer.get_extra_info("peername")
        print(f"  new connection from {addr}")
        await handle_local_client(reader, writer, args.host, dest, use_tls, resolved_ip)

    server = await asyncio.start_server(on_connect, "127.0.0.1", args.port)
    proto = "wss" if use_tls else "ws"
    print(f"Tunnel listening on localhost:{args.port}")
    print(f"  → {proto}://{args.host}/tunnel/{dest} → {DESTINATIONS[dest]}")
    print(f"  Run: ssh -p {args.port} user@localhost")

    async with server:
        await server.serve_forever()


if __name__ == "__main__":
    asyncio.run(main())