#!/usr/bin/env python3
"""Local helper for SSH ProxyCommand using HTTPS CONNECT or WebSockets.

Example:
ssh -o ProxyCommand='./ssh_https_connect.py --proxy-host proxy.example.com --proxy-port 443 --proxy-user proxyuser --target-host %h --target-port %p --insecure' user@server
"""
from __future__ import annotations

import argparse
import base64
import getpass
import os
import selectors
import socket
import ssl
import sys
import urllib.parse
import uuid

BUFFER_SIZE = 65536


def parse_args():
    p = argparse.ArgumentParser(description="SSH over HTTPS CONNECT / WebSocket helper")
    p.add_argument("--proxy-host", required=True)
    p.add_argument("--proxy-port", type=int, default=443)
    p.add_argument("--proxy-user", default=os.getenv("PROXY_USER", ""))
    p.add_argument("--proxy-pass", default=os.getenv("PROXY_PASS", ""))
    p.add_argument("--target-host", required=True)
    p.add_argument("--target-port", type=int, default=22)
    p.add_argument("--insecure", action="store_true", help="Disable TLS certificate verification")
    p.add_argument("--server-name", help="SNI name; defaults to --proxy-host")
    p.add_argument("--host-header", help="Custom Host header; defaults to --proxy-host")
    p.add_argument("--mode", choices=["connect", "ws", "websocket", "auto"], default="auto",
                   help="Tunneling mode: connect, ws (websocket), or auto (defaults to auto)")
    return p.parse_args()


class WebSocketTunnel:
    def __init__(self, sock: ssl.SSLSocket, initial_data: bytes = b""):
        self.sock = sock
        self.read_buf = initial_data
        self.frame_payloads = []
        if self.read_buf:
            self._parse_frames()

    def setblocking(self, flag: bool):
        self.sock.setblocking(flag)

    def fileno(self) -> int:
        return self.sock.fileno()

    def recv(self, bufsize: int) -> bytes:
        if self.frame_payloads:
            return self.frame_payloads.pop(0)
        try:
            data = self.sock.recv(bufsize)
        except ssl.SSLWantReadError:
            raise
        if not data:
            return b""
        self.read_buf += data
        self._parse_frames()
        if self.frame_payloads:
            return self.frame_payloads.pop(0)
        raise ssl.SSLWantReadError()

    def _parse_frames(self):
        while len(self.read_buf) >= 2:
            fin_opcode = self.read_buf[0]
            opcode = fin_opcode & 0x0f
            mask_len = self.read_buf[1]
            has_mask = bool(mask_len & 0x80)
            length = mask_len & 0x7f

            idx = 2
            if length == 126:
                if len(self.read_buf) < 4:
                    break
                length = int.from_bytes(self.read_buf[2:4], "big")
                idx = 4
            elif length == 127:
                if len(self.read_buf) < 10:
                    break
                length = int.from_bytes(self.read_buf[2:10], "big")
                idx = 10

            mask_offset = idx
            if has_mask:
                idx += 4

            if len(self.read_buf) < idx + length:
                break

            payload = self.read_buf[idx : idx + length]
            if has_mask:
                mask_key = self.read_buf[mask_offset : mask_offset + 4]
                payload = bytes(b ^ mask_key[i % 4] for i, b in enumerate(payload))

            self.read_buf = self.read_buf[idx + length :]
            if opcode == 8:  # Close frame
                self.frame_payloads.append(b"")
                return
            if opcode in (1, 2):  # Text or Binary
                self.frame_payloads.append(payload)

    def sendall(self, data: bytes):
        mask = os.urandom(4)
        length = len(data)
        if length <= 125:
            header = bytes([0x82, 0x80 | length])
        elif length <= 65535:
            header = bytes([0x82, 0x80 | 126]) + length.to_bytes(2, "big")
        else:
            header = bytes([0x82, 0x80 | 127]) + length.to_bytes(8, "big")
        masked_data = bytes(b ^ mask[i % 4] for i, b in enumerate(data))
        self.sock.sendall(header + mask + masked_data)

    def shutdown(self, how):
        if how == socket.SHUT_WR:
            try:
                self.sock.sendall(b"\x88\x80\x00\x00\x00\x00")
            except OSError:
                pass
        try:
            self.sock.shutdown(how)
        except OSError:
            pass

    def close(self):
        self.sock.close()


def open_connect_tunnel(args) -> ssl.SSLSocket:
    raw = socket.create_connection((args.proxy_host, args.proxy_port), timeout=20)
    if args.insecure:
        ctx = ssl._create_unverified_context()
    else:
        ctx = ssl.create_default_context()
    tls = ctx.wrap_socket(raw, server_hostname=args.server_name or args.proxy_host)

    host_header = args.host_header or args.proxy_host
    if not args.host_header and args.proxy_port not in (80, 443):
        host_header = f"{args.proxy_host}:{args.proxy_port}"

    headers = [
        f"CONNECT {args.target_host}:{args.target_port} HTTP/1.1",
        f"Host: {host_header}",
        "User-Agent: ssh-https-connect/1.0",
        "Proxy-Connection: Keep-Alive",
    ]
    password = args.proxy_pass
    if args.proxy_user and not password:
        password = getpass.getpass("Proxy password: ")
        args.proxy_pass = password
    if args.proxy_user:
        token = base64.b64encode(f"{args.proxy_user}:{password}".encode()).decode()
        headers.append(f"Proxy-Authorization: Basic {token}")
        headers.append(f"Authorization: Basic {token}")
        headers.append(f"X-Authorization: Basic {token}")
        headers.append(f"X-Auth: Basic {token}")
    request = "\r\n".join(headers) + "\r\n\r\n"
    tls.sendall(request.encode())

    response = b""
    while b"\r\n\r\n" not in response:
        chunk = tls.recv(4096)
        if not chunk:
            raise RuntimeError("Proxy closed connection before CONNECT response")
        response += chunk
        if len(response) > 65536:
            raise RuntimeError("CONNECT response too large")
    status_line = response.split(b"\r\n", 1)[0].decode("iso-8859-1", "replace")
    if not status_line.startswith("HTTP/") or " 200 " not in status_line:
        raise RuntimeError(f"CONNECT failed: {status_line}")
    return tls


def open_ws_tunnel(args) -> WebSocketTunnel:
    raw = socket.create_connection((args.proxy_host, args.proxy_port), timeout=20)
    if args.insecure:
        ctx = ssl._create_unverified_context()
    else:
        ctx = ssl.create_default_context()
    tls = ctx.wrap_socket(raw, server_hostname=args.server_name or args.proxy_host)

    ws_key = base64.b64encode(uuid.uuid4().bytes).decode()
    target_encoded = urllib.parse.quote(f"{args.target_host}:{args.target_port}")
    path = f"/tunnel?target={target_encoded}"

    host_header = args.host_header or args.proxy_host
    if not args.host_header and args.proxy_port not in (80, 443):
        host_header = f"{args.proxy_host}:{args.proxy_port}"

    headers = [
        f"GET {path} HTTP/1.1",
        f"Host: {host_header}",
        "Upgrade: websocket",
        "Connection: Upgrade",
        f"Sec-WebSocket-Key: {ws_key}",
        "Sec-WebSocket-Version: 13",
        f"X-Target: {args.target_host}:{args.target_port}",
        "User-Agent: ssh-https-connect/1.0",
    ]
    password = args.proxy_pass
    if args.proxy_user and not password:
        password = getpass.getpass("Proxy password: ")
        args.proxy_pass = password
    if args.proxy_user:
        token = base64.b64encode(f"{args.proxy_user}:{password}".encode()).decode()
        headers.append(f"Proxy-Authorization: Basic {token}")
        headers.append(f"Authorization: Basic {token}")
        headers.append(f"X-Authorization: Basic {token}")
        headers.append(f"X-Auth: Basic {token}")
    request = "\r\n".join(headers) + "\r\n\r\n"
    tls.sendall(request.encode())

    response = b""
    while b"\r\n\r\n" not in response:
        chunk = tls.recv(4096)
        if not chunk:
            raise RuntimeError("Proxy closed connection before WS handshake response")
        response += chunk
        if len(response) > 65536:
            raise RuntimeError("Handshake response too large")

    header_part, body_part = response.split(b"\r\n\r\n", 1)
    status_line = header_part.split(b"\r\n", 1)[0].decode("iso-8859-1", "replace")
    if not status_line.startswith("HTTP/") or " 101 " not in status_line:
        raise RuntimeError(f"WS handshake failed: {status_line}\nResponse: {header_part.decode('iso-8859-1')}")

    return WebSocketTunnel(tls, body_part)


def pump(sock: socket.socket | WebSocketTunnel) -> int:
    sel = selectors.DefaultSelector()
    stdin = sys.stdin.buffer
    stdout = sys.stdout.buffer
    sock.setblocking(False)
    os.set_blocking(stdin.fileno(), False)
    sel.register(sock, selectors.EVENT_READ, "sock")
    sel.register(stdin, selectors.EVENT_READ, "stdin")
    try:
        while True:
            for key, _ in sel.select():
                if key.data == "sock":
                    try:
                        data = sock.recv(BUFFER_SIZE)
                    except ssl.SSLWantReadError:
                        continue
                    if not data:
                        return 0
                    stdout.write(data)
                    stdout.flush()
                else:
                    try:
                        data = os.read(stdin.fileno(), BUFFER_SIZE)
                    except BlockingIOError:
                        continue
                    if not data:
                        try:
                            sock.shutdown(socket.SHUT_WR)
                        except OSError:
                            pass
                        sel.unregister(stdin)
                    else:
                        sock.sendall(data)
    finally:
        try:
            sel.close()
        except Exception:
            pass


def main() -> int:
    args = parse_args()
    try:
        if args.mode in ("ws", "websocket"):
            sock = open_ws_tunnel(args)
        elif args.mode == "connect":
            sock = open_connect_tunnel(args)
        else:
            # auto mode
            try:
                sock = open_connect_tunnel(args)
            except Exception as exc:
                err_str = str(exc)
                if any(x in err_str for x in ("404", "405", "502", "Bad Request", "closed connection before CONNECT")):
                    print(f"CONNECT failed ({exc}), falling back to WebSocket tunnel...", file=sys.stderr)
                    sock = open_ws_tunnel(args)
                else:
                    raise
        return pump(sock)
    except Exception as exc:
        print(f"ssh_https_connect: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())