#!/usr/bin/env python3
# ssh-via-proxy <host> <port> — stdio bridge for `ssh -o ProxyCommand`.
# Carries the SSH session inside an HTTP CONNECT tunnel through the egress proxy.
import os
import select
import socket
import sys
import urllib.parse
import base64


def main():
    target_host, target_port = sys.argv[1], int(sys.argv[2])
    u = urllib.parse.urlparse(os.environ.get("HTTPS_PROXY") or os.environ.get("https_proxy"))
    auth = base64.b64encode(
        f"{urllib.parse.unquote(u.username)}:{urllib.parse.unquote(u.password)}".encode()
    ).decode()

    s = socket.create_connection((u.hostname, u.port), timeout=20)
    s.sendall(
        f"CONNECT {target_host}:{target_port} HTTP/1.1\r\n"
        f"Host: {target_host}:{target_port}\r\n"
        f"Proxy-Authorization: Basic {auth}\r\n\r\n".encode()
    )
    head = b""
    while b"\r\n\r\n" not in head:
        chunk = s.recv(4096)
        if not chunk:
            sys.stderr.write("proxy closed during CONNECT\n")
            sys.exit(1)
        head += chunk
    status = head.split(b"\r\n", 1)[0].decode(errors="replace")
    if " 200 " not in f" {status} ":
        sys.stderr.write(f"CONNECT failed: {status}\n")
        sys.exit(1)
    rest = head.split(b"\r\n\r\n", 1)[1]

    stdin = sys.stdin.buffer
    stdout = sys.stdout.buffer
    s.setblocking(False)
    os.set_blocking(stdin.fileno(), False)
    if rest:
        stdout.write(rest)
        stdout.flush()
    while True:
        r, _, _ = select.select([s, stdin], [], [])
        if s in r:
            data = s.recv(65536)
            if not data:
                break
            stdout.write(data)
            stdout.flush()
        if stdin in r:
            data = os.read(stdin.fileno(), 65536)
            if not data:
                break
            s.sendall(data)


if __name__ == "__main__":
    main()
