#!/usr/bin/env python3
"""rafb-relay: a small WebSocket-to-TCP relay for browser SSH, SFTP and FTP clients.

A browser can't open TCP connections, so a web-based SSH client sends its bytes over a WebSocket to this relay,
which forwards them to the target server and back. SSH and SFTP are encrypted end to end between the browser and
the target: the relay only ever sees ciphertext. Plain FTP is not encrypted, so use FTP only through a relay you
trust (ideally your own).

Python 3.9+, standard library only, one file. Run it behind a TLS reverse proxy (Caddy, nginx) so browsers can
reach it as wss://. See README.md for the protocol, configuration and deployment.
"""

import argparse
import asyncio
import base64
import fnmatch
import hashlib
import hmac
import ipaddress
import json
import logging
import os
import signal
import socket
import struct
import sys
import time
from typing import Dict, List, Optional, Tuple
from urllib.parse import parse_qs, urlsplit

VERSION = "1.0.0"
PROTOCOL = 1
WS_GUID = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

log = logging.getLogger("rafb-relay")

DEFAULT_CONFIG = {
    "listen": "127.0.0.1:8022",
    # Browser origins allowed to open connections. "*" allows any site (only sensible with token/ticket auth).
    "allowed_origins": ["https://rafb.me"],
    # "none" (anyone who can reach the relay), "token" (shared secrets) or "ticket" (short-lived signed tickets
    # issued by a website that shares ticket_secret with the relay).
    "auth": {"mode": "token", "tokens": [], "ticket_secret": ""},
    "targets": {
        "ports": [22],
        # Allow private, loopback and link-local addresses (your LAN). Keep false on a public relay.
        "allow_private": False,
        # If non-empty, only these hosts may be reached: hostnames (wildcards allowed) or CIDR ranges.
        "allow_hosts": [],
        "deny_hosts": [],
    },
    "ftp": {
        # FTP opens a second connection per transfer to a port the server picks. When enabled, a user with an FTP
        # control connection (to a control port) may also connect to ports >= data_port_min on that same address.
        "enabled": False,
        "control_ports": [21],
        "data_port_min": 1024,
    },
    "limits": {
        "max_connections_per_user": 4,
        "max_connections_total": 200,
        "connect_timeout_seconds": 10,
        "idle_timeout_seconds": 1800,
        "max_session_seconds": 43200,
        "max_frame_bytes": 1048576,
    },
    # Read the client address from X-Forwarded-For (set by your reverse proxy). Only enable behind a proxy.
    "trust_proxy": True,
}


class RelayError(Exception):
    """Reported to the browser as {"type": "error", "code": ..., "message": ...} and close code `close_code`."""

    def __init__(self, code: str, message: str, close_code: int):
        super().__init__(message)
        self.code = code
        self.message = message
        self.close_code = close_code


def deep_merge(base: dict, override: dict) -> dict:
    out = dict(base)
    for key, value in override.items():
        out[key] = deep_merge(out[key], value) if isinstance(value, dict) and isinstance(out.get(key), dict) else value
    return out


def load_config(path: Optional[str]) -> dict:
    config = DEFAULT_CONFIG
    if path:
        with open(path, encoding="utf-8") as f:
            config = deep_merge(DEFAULT_CONFIG, json.load(f))
    # Secrets can come from the environment instead of the file.
    if os.environ.get("RELAY_TICKET_SECRET"):
        config = deep_merge(config, {"auth": {"ticket_secret": os.environ["RELAY_TICKET_SECRET"]}})
    if os.environ.get("RELAY_TOKENS"):
        config = deep_merge(config, {"auth": {"tokens": os.environ["RELAY_TOKENS"].split(",")}})
    mode = config["auth"]["mode"]
    if mode not in ("none", "token", "ticket"):
        raise ValueError("auth.mode must be none, token or ticket")
    if mode == "ticket" and len(config["auth"]["ticket_secret"]) < 16:
        raise ValueError("auth.ticket_secret must be at least 16 characters")
    if mode == "token" and not [t for t in config["auth"]["tokens"] if len(t) >= 16]:
        raise ValueError("auth.tokens needs at least one token of 16+ characters")
    return config


# ---------------------------------------------------------------------------------------------------------------
# Authentication
# ---------------------------------------------------------------------------------------------------------------

def b64url(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def b64url_decode(text: str) -> bytes:
    return base64.urlsafe_b64decode(text + "=" * (-len(text) % 4))


def make_ticket(secret: str, payload: dict) -> str:
    """Ticket = "v1." + b64url(json payload) + "." + b64url(HMAC-SHA256(secret, "v1." + b64url(json payload)))."""
    body = "v1." + b64url(json.dumps(payload, separators=(",", ":")).encode("utf-8"))
    return body + "." + b64url(hmac.new(secret.encode("utf-8"), body.encode("ascii"), hashlib.sha256).digest())


def verify_ticket(secret: str, ticket: str, now: Optional[float] = None) -> dict:
    """Returns the payload of a valid, unexpired ticket. Payload: {"sub": user id, "exp": unix seconds}."""
    try:
        version, payload_b64, sig_b64 = ticket.split(".")
        if version != "v1":
            raise ValueError
        expected = hmac.new(secret.encode("utf-8"), f"v1.{payload_b64}".encode("ascii"), hashlib.sha256).digest()
        if not hmac.compare_digest(expected, b64url_decode(sig_b64)):
            raise ValueError
        payload = json.loads(b64url_decode(payload_b64))
        sub, exp = str(payload["sub"]), float(payload["exp"])
    except (ValueError, KeyError, TypeError):
        raise RelayError("unauthorized", "Invalid relay ticket.", 4401)
    if exp < (now if now is not None else time.time()) - 30:  # small allowance for clock skew
        raise RelayError("unauthorized", "Relay ticket expired; get a new one.", 4401)
    if not sub or len(sub) > 128:
        raise RelayError("unauthorized", "Invalid relay ticket.", 4401)
    return payload


def authenticate(config: dict, credential: str, client_ip: str) -> str:
    """Returns the user id used for per-user limits and logs."""
    auth = config["auth"]
    if auth["mode"] == "none":
        return "ip:" + client_ip
    if not credential:
        raise RelayError("unauthorized", "This relay needs an auth token or ticket.", 4401)
    if auth["mode"] == "ticket":
        return verify_ticket(auth["ticket_secret"], credential)["sub"]
    for token in auth["tokens"]:
        if len(token) >= 16 and hmac.compare_digest(token.encode("utf-8"), credential.encode("utf-8")):
            return "token:" + hashlib.sha256(token.encode("utf-8")).hexdigest()[:8]
    raise RelayError("unauthorized", "Wrong relay token.", 4401)


# ---------------------------------------------------------------------------------------------------------------
# Target policy
# ---------------------------------------------------------------------------------------------------------------

def is_public(ip: ipaddress._BaseAddress) -> bool:
    if isinstance(ip, ipaddress.IPv6Address) and ip.ipv4_mapped:
        ip = ip.ipv4_mapped
    return ip.is_global and not ip.is_multicast


def host_matches(patterns: List[str], host: str, ip: Optional[ipaddress._BaseAddress]) -> bool:
    for pattern in patterns:
        try:
            if ip is not None and ip in ipaddress.ip_network(pattern, strict=False):
                return True
            continue
        except ValueError:
            pass
        if fnmatch.fnmatch(host.lower().rstrip("."), pattern.lower()):
            return True
    return False


class Relay:
    def __init__(self, config: dict):
        self.config = config
        self.connections: Dict[str, int] = {}  # user -> open connections
        self.ftp_controls: Dict[Tuple[str, str], int] = {}  # (user, ip) -> open FTP control connections
        self.total = 0

    # ---- policy ----

    async def resolve_target(self, user: str, host: str, port: int) -> Tuple[str, bool]:
        """Checks the request and returns (ip to connect to, is_ftp_control)."""
        targets, ftp = self.config["targets"], self.config["ftp"]
        if not host or len(host) > 253 or not (1 <= port <= 65535):
            raise RelayError("bad_request", "Give a host and a port between 1 and 65535.", 4400)
        if host_matches(targets["deny_hosts"], host, None):
            raise RelayError("forbidden", f"{host} is blocked on this relay.", 4403)

        is_control = ftp["enabled"] and port in ftp["control_ports"]
        port_allowed = port in targets["ports"] or is_control
        maybe_ftp_data = ftp["enabled"] and not port_allowed and port >= ftp["data_port_min"]
        if not port_allowed and not maybe_ftp_data:
            allowed = sorted(set(targets["ports"]) | (set(ftp["control_ports"]) if ftp["enabled"] else set()))
            raise RelayError("forbidden", f"Port {port} isn't allowed on this relay (allowed: {allowed}).", 4403)

        try:
            infos = await asyncio.get_running_loop().getaddrinfo(host, port, type=socket.SOCK_STREAM)
        except socket.gaierror:
            raise RelayError("connect_failed", f"Couldn't resolve {host}.", 4502)

        for info in infos:
            ip = ipaddress.ip_address(info[4][0].split("%")[0])
            if host_matches(targets["deny_hosts"], host, ip):
                continue
            explicitly_allowed = bool(targets["allow_hosts"]) and host_matches(targets["allow_hosts"], host, ip)
            if targets["allow_hosts"] and not explicitly_allowed:
                continue
            if not is_public(ip) and not targets["allow_private"] and not explicitly_allowed:
                continue
            if maybe_ftp_data and not self.ftp_controls.get((user, str(ip))):
                continue
            return str(ip), is_control

        if maybe_ftp_data:
            raise RelayError("forbidden", f"Port {port} is only allowed for FTP data, after an FTP login to that server.", 4403)
        if targets["allow_hosts"]:
            raise RelayError("forbidden", f"{host} isn't on this relay's list of allowed hosts.", 4403)
        raise RelayError("forbidden", f"{host} resolves to a private or reserved address, which this relay doesn't allow.", 4403)

    def acquire(self, user: str) -> None:
        limits = self.config["limits"]
        if self.total >= limits["max_connections_total"]:
            raise RelayError("busy", "The relay is at capacity; try again shortly.", 4429)
        if self.connections.get(user, 0) >= limits["max_connections_per_user"]:
            raise RelayError("too_many", f"You already have {limits['max_connections_per_user']} connections open through this relay.", 4429)
        self.connections[user] = self.connections.get(user, 0) + 1
        self.total += 1

    def release(self, user: str) -> None:
        self.total -= 1
        self.connections[user] -= 1
        if not self.connections[user]:
            del self.connections[user]

    # ---- HTTP ----

    async def handle(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        try:
            head = await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), 10)
        except (asyncio.TimeoutError, asyncio.IncompleteReadError, asyncio.LimitOverrunError, ConnectionError):
            writer.close()
            return
        try:
            request_line, *header_lines = head.decode("latin-1").split("\r\n")
            method, target, _ = request_line.split(" ", 2)
            headers = {}
            for line in header_lines:
                if ":" in line:
                    name, value = line.split(":", 1)
                    headers[name.strip().lower()] = value.strip()
        except ValueError:
            await self.respond(writer, 400, "Bad request")
            return

        url = urlsplit(target)
        origin = headers.get("origin", "")
        client_ip = (writer.get_extra_info("peername") or ("?",))[0]
        if self.config["trust_proxy"] and headers.get("x-forwarded-for"):
            client_ip = headers["x-forwarded-for"].split(",")[0].strip()

        if url.path == "/healthz":
            await self.respond(writer, 200, "ok")
        elif url.path == "/" and method == "GET":
            await self.respond(writer, 200, json.dumps(self.info()), "application/json", cors=True)
        elif url.path == "/v1/connect" and headers.get("upgrade", "").lower() == "websocket":
            if not self.origin_allowed(origin):
                log.info("rejected origin=%s ip=%s", origin or "-", client_ip)
                await self.respond(writer, 403, "Origin not allowed")
                return
            await self.websocket(reader, writer, headers, parse_qs(url.query), client_ip, origin)
        else:
            await self.respond(writer, 404, "Not found")

    def info(self) -> dict:
        t, ftp = self.config["targets"], self.config["ftp"]
        return {
            "name": "rafb-relay", "version": VERSION, "protocol": PROTOCOL, "auth": self.config["auth"]["mode"],
            "ports": t["ports"], "allowPrivate": t["allow_private"], "restrictedHosts": bool(t["allow_hosts"]),
            "ftp": ftp["enabled"], "ftpControlPorts": ftp["control_ports"] if ftp["enabled"] else [],
        }

    def origin_allowed(self, origin: str) -> bool:
        allowed = self.config["allowed_origins"]
        return "*" in allowed or origin in allowed

    async def respond(self, writer, status: int, body: str, content_type: str = "text/plain", cors: bool = False):
        reasons = {200: "OK", 400: "Bad Request", 403: "Forbidden", 404: "Not Found"}
        data = body.encode("utf-8")
        head = (f"HTTP/1.1 {status} {reasons.get(status, '')}\r\nContent-Type: {content_type}; charset=utf-8\r\n"
                f"Content-Length: {len(data)}\r\nConnection: close\r\nCache-Control: no-store\r\n")
        if cors:
            head += "Access-Control-Allow-Origin: *\r\n"
        writer.write(head.encode("latin-1") + b"\r\n" + data)
        try:
            await writer.drain()
        finally:
            writer.close()

    # ---- WebSocket <-> TCP ----

    async def websocket(self, reader, writer, headers, query, client_ip, origin) -> None:
        key = headers.get("sec-websocket-key", "")
        try:
            key_ok = len(base64.b64decode(key, validate=True)) == 16
        except ValueError:
            key_ok = False
        if headers.get("sec-websocket-version") != "13" or not key_ok:
            await self.respond(writer, 400, "Bad WebSocket handshake")
            return
        accept = base64.b64encode(hashlib.sha1(key.encode("ascii") + WS_GUID).digest()).decode("ascii")
        writer.write(("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n"
                      f"Sec-WebSocket-Accept: {accept}\r\n\r\n").encode("ascii"))
        ws = WebSocket(reader, writer, self.config["limits"]["max_frame_bytes"])

        host = (query.get("host") or [""])[0].strip().strip("[]")
        port_text = (query.get("port") or [""])[0]
        user, acquired, control_key = None, False, None
        try:
            if not port_text.isdigit():
                raise RelayError("bad_request", "Give a host and a port between 1 and 65535.", 4400)
            port = int(port_text)
            user = authenticate(self.config, (query.get("auth") or [""])[0], client_ip)
            ip, is_control = await self.resolve_target(user, host, port)
            self.acquire(user)
            acquired = True
            try:
                tcp_reader, tcp_writer = await asyncio.wait_for(
                    asyncio.open_connection(ip, port), self.config["limits"]["connect_timeout_seconds"])
            except (OSError, asyncio.TimeoutError) as e:
                raise RelayError("connect_failed", f"Couldn't connect to {host}:{port} ({type(e).__name__}).", 4502)
            if is_control:
                control_key = (user, ip)
                self.ftp_controls[control_key] = self.ftp_controls.get(control_key, 0) + 1
            await ws.send_text(json.dumps({"type": "connected", "address": ip, "port": port}))
            log.info("open user=%s client=%s origin=%s target=%s:%d ip=%s", user, client_ip, origin or "-", host, port, ip)
            started = time.monotonic()
            up, down = await self.pump(ws, tcp_reader, tcp_writer)
            log.info("close user=%s target=%s:%d up=%d down=%d secs=%d", user, host, port, up, down, time.monotonic() - started)
        except RelayError as e:
            log.info("refused user=%s client=%s target=%s:%s reason=%s", user or "-", client_ip, host, port_text, e.code)
            await ws.send_text(json.dumps({"type": "error", "code": e.code, "message": e.message}))
            await ws.close(e.close_code, e.code)
        except (ConnectionError, asyncio.IncompleteReadError):
            pass
        finally:
            if control_key:
                self.ftp_controls[control_key] -= 1
                if not self.ftp_controls[control_key]:
                    del self.ftp_controls[control_key]
            if acquired:
                self.release(user)
            writer.close()

    async def pump(self, ws: "WebSocket", tcp_reader, tcp_writer) -> Tuple[int, int]:
        limits = self.config["limits"]
        counts = [0, 0]
        last_activity = [time.monotonic()]

        async def ws_to_tcp():
            while True:
                data = await ws.receive()
                if data is None:
                    return
                last_activity[0] = time.monotonic()
                counts[0] += len(data)
                tcp_writer.write(data)
                await tcp_writer.drain()

        async def tcp_to_ws():
            while True:
                data = await tcp_reader.read(65536)
                if not data:
                    return
                last_activity[0] = time.monotonic()
                counts[1] += len(data)
                await ws.send_binary(data)

        async def watchdog():
            started = time.monotonic()
            while True:
                await asyncio.sleep(5)
                now = time.monotonic()
                if now - last_activity[0] > limits["idle_timeout_seconds"]:
                    return "idle"
                if now - started > limits["max_session_seconds"]:
                    return "session time limit"

        tasks = [asyncio.ensure_future(t()) for t in (ws_to_tcp, tcp_to_ws, watchdog)]
        reason = None
        try:
            done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
            for task in done:
                if task is tasks[2]:
                    reason = task.result()
                else:
                    task.exception()  # a dropped connection just ends the session
        finally:
            for task in tasks:
                task.cancel()
            tcp_writer.close()
        await ws.close(4408 if reason else 1000, reason or "closed")
        return counts[0], counts[1]


class WebSocket:
    """Server side of RFC 6455, just what the relay needs: binary messages, ping/pong and close."""

    def __init__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, max_message: int):
        self.reader, self.writer, self.max_message = reader, writer, max_message
        self.closed = False

    async def receive(self) -> Optional[bytes]:
        """The next binary message, or None once the connection is closing."""
        message = bytearray()
        while True:
            b1, b2 = await self.reader.readexactly(2)
            fin, opcode, masked, length = b1 & 0x80, b1 & 0x0F, b2 & 0x80, b2 & 0x7F
            if not masked:
                await self.close(1002, "client frames must be masked")
                return None
            if length == 126:
                length = struct.unpack(">H", await self.reader.readexactly(2))[0]
            elif length == 127:
                length = struct.unpack(">Q", await self.reader.readexactly(8))[0]
            if length + len(message) > self.max_message:
                await self.close(1009, "message too big")
                return None
            mask = await self.reader.readexactly(4)
            payload = await self.reader.readexactly(length)
            if length:
                payload = (int.from_bytes(payload, "big") ^ int.from_bytes((mask * (length // 4 + 1))[:length], "big")).to_bytes(length, "big")

            if opcode == 0x8:  # close
                await self.close(1000, "")
                return None
            if opcode == 0x9:  # ping
                await self._send(0xA, payload)
                continue
            if opcode == 0xA:  # pong
                continue
            if opcode in (0x0, 0x1, 0x2):  # continuation, text (ignored), binary
                message += payload
                if fin:
                    if opcode == 0x1:
                        message = bytearray()
                        continue
                    return bytes(message)
                continue
            await self.close(1002, "unknown opcode")
            return None

    async def send_binary(self, data: bytes) -> None:
        await self._send(0x2, data)

    async def send_text(self, text: str) -> None:
        await self._send(0x1, text.encode("utf-8"))

    async def close(self, code: int, reason: str) -> None:
        if self.closed:
            return
        try:
            await self._send(0x8, struct.pack(">H", code) + reason.encode("utf-8")[:120])
        except (ConnectionError, RuntimeError):
            pass
        self.closed = True

    async def _send(self, opcode: int, payload: bytes) -> None:
        if self.closed:
            return
        n = len(payload)
        if n < 126:
            head = struct.pack(">BB", 0x80 | opcode, n)
        elif n < 65536:
            head = struct.pack(">BBH", 0x80 | opcode, 126, n)
        else:
            head = struct.pack(">BBQ", 0x80 | opcode, 127, n)
        self.writer.write(head + payload)
        await self.writer.drain()


async def serve(config: dict) -> None:
    host, _, port = config["listen"].rpartition(":")
    relay = Relay(config)
    server = await asyncio.start_server(relay.handle, host.strip("[]") or None, int(port), limit=16384)
    log.info("rafb-relay %s listening on %s (auth=%s, ports=%s, private=%s, ftp=%s, origins=%s)", VERSION,
             config["listen"], config["auth"]["mode"], config["targets"]["ports"], config["targets"]["allow_private"],
             config["ftp"]["enabled"], ",".join(config["allowed_origins"]))
    stop = asyncio.Event()
    for sig in (signal.SIGINT, signal.SIGTERM):
        try:
            asyncio.get_running_loop().add_signal_handler(sig, stop.set)
        except (NotImplementedError, RuntimeError):  # Windows
            pass
    async with server:
        await stop.wait()


def main(argv: Optional[List[str]] = None) -> int:
    parser = argparse.ArgumentParser(description="WebSocket-to-TCP relay for browser SSH/SFTP/FTP clients.")
    parser.add_argument("--config", "-c", help="JSON config file (see config.example.json)")
    parser.add_argument("--listen", help="host:port to listen on, overriding the config")
    parser.add_argument("--make-token", action="store_true", help="print a random token or secret and exit")
    args = parser.parse_args(argv)
    if args.make_token:
        print(b64url(os.urandom(32)))
        return 0
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s", stream=sys.stdout)
    config = load_config(args.config)
    if args.listen:
        config["listen"] = args.listen
    try:
        asyncio.run(serve(config))
    except KeyboardInterrupt:
        pass
    return 0


if __name__ == "__main__":
    sys.exit(main())