"""Connect this computer to the Steam Frame: find it, create keys, add a `Host frame`
alias to ~/.ssh/config and get a key onto the headset. It first asks Valve's
SteamOS devkit service (port 32000) to pair, which needs only a tap on the
headset; if that service isn't there or says no, it copies the key over SSH,
asking for the Developer Mode password once. The Linux and Windows twin of
scripts/connect.sh (which the Mac app uses); same config block, so either can
re-run over the other. Idempotent.

Usage: python3 ui/frame_connect.py [HOST_OR_IP[:PORT]]
Env:   FRAME_USER (default steamos), FRAME_ALIAS (default frame)
"""
import base64
import json
import os
import platform
import re
import shutil
import socket
import subprocess
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

FRAME_USER = os.environ.get("FRAME_USER", "steamos")
USER_FROM_ENV = "FRAME_USER" in os.environ
FRAME_ALIAS = os.environ.get("FRAME_ALIAS", "frame")
SSH_DIR = Path.home() / ".ssh"
KEY = SSH_DIR / "id_ed25519_frame"
# The devkit service only accepts ssh-rsa keys (write_key in Valve's
# steamos-devkit-service), so pairing uses a second key next to the ed25519 one.
DEVKIT_KEY = SSH_DIR / "id_rsa_frame_devkit"
CONFIG = SSH_DIR / "config"
NAME_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*")
# Both go into ~/.ssh/config, so nothing that could add a line or a directive.
for _name, _value in (("FRAME_ALIAS", FRAME_ALIAS), ("FRAME_USER", FRAME_USER)):
    if not NAME_RE.fullmatch(_value):
        sys.exit(f"{_name} must be a plain name, not {_value!r}")
BEGIN = f"# >>> steam-frame ({FRAME_ALIAS}) >>>"
END = f"# <<< steam-frame ({FRAME_ALIAS}) <<<"

# Appends the key from stdin unless it's already there. base64 keeps it intact
# through Windows' command-line quoting.
ADD_KEY = """umask 077
mkdir -p ~/.ssh
k=$(cat)
grep -qxF "$k" ~/.ssh/authorized_keys 2>/dev/null || printf '%s\\n' "$k" >> ~/.ssh/authorized_keys
"""
ADD_KEY_CMD = 'sh -c "$(echo %s | base64 -d)"' % base64.b64encode(ADD_KEY.encode()).decode()


def say(msg):
    print(msg, flush=True)


# --- Valve's SteamOS devkit pairing (steamos-devkit-service on the headset). HTTP on
# port 32000: GET /properties.json names the user to log in as; POST /register with
# "ssh-rsa <key> <comment> <magic>" shows an approve prompt in the headset (the
# comment is what it displays, 30 s to answer), then installs the key and turns sshd on.
# The prompt only appears while Steam is on Settings > Developer > Pair new host;
# otherwise /register answers 403 "please put the Steam client in pairing mode".

DEVKIT_PORT = 32000
DEVKIT_SERVICE = "_steamos-devkit._tcp"
MAGIC_PHRASE = "900b919520e4cf601998a71eec318fec"  # fixed token Valve's client appends
REGISTER_TIMEOUT = 60
PAIRING_MODE_WAIT = 120  # seconds to keep asking while the user opens "Pair new host"
# A LAN host: never go through an HTTP(S)_PROXY from the environment.
_opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))


def key_comment(node):
    """"frame-control@<short host name>" as one word: the headset splits the body on spaces."""
    name = re.sub(r"[^A-Za-z0-9._-]+", "-", (node or "").split(".")[0]).strip("-.") or "computer"
    return f"frame-control@{name}"


def register_body(pub, comment):
    fields = pub.split()
    if len(fields) < 2 or fields[0] != "ssh-rsa":
        raise ValueError("the devkit service only takes ssh-rsa keys")
    return f"ssh-rsa {fields[1]} {comment} {MAGIC_PHRASE}\n"


def parse_login(raw):
    """The `login` from /properties.json if it's a plain user name, else None. "root"
    means several users are configured and Valve's client switches between them; we
    can't, so it counts as no answer."""
    props = json.loads(raw)
    if not isinstance(props, dict):
        raise ValueError("properties.json isn't a JSON object")
    login = props.get("login")
    if isinstance(login, str) and NAME_RE.fullmatch(login) and login != "root":
        return login
    return None


def devkit_error(status, raw):
    """A readable reason from a failed /register: its {"error": ...} JSON, or the text."""
    text = raw.decode("utf-8", "replace").strip()
    try:
        err = json.loads(text).get("error")
    except (ValueError, AttributeError):
        err = None
    return str(err or text or f"HTTP {status}")[:300]


def devkit_url(host, port, path):
    return f"http://[{host}]:{port}{path}" if ":" in host else f"http://{host}:{port}{path}"


def why(e):
    return str(getattr(e, "reason", None) or e)


def fetch_login(host, port=DEVKIT_PORT, timeout=5):
    """GET /properties.json. Raises OSError (HTTP errors included) or ValueError."""
    with _opener.open(devkit_url(host, port, "/properties.json"), timeout=timeout) as r:
        return parse_login(r.read())


def register(host, body, port=DEVKIT_PORT, timeout=REGISTER_TIMEOUT):
    """POST /register, which waits while someone answers the prompt. -> (ok, message)"""
    req = urllib.request.Request(devkit_url(host, port, "/register"), data=body.encode("ascii"),
                                 headers={"Content-Type": "text/plain"}, method="POST")
    try:
        with _opener.open(req, timeout=timeout) as r:
            return True, r.read().decode("utf-8", "replace").strip()
    except urllib.error.HTTPError as e:
        with e:
            return False, devkit_error(e.code, e.read())
    except OSError as e:
        return False, f"no answer ({why(e)})"


def devkit_pair(host, pub, comment, port=DEVKIT_PORT, on_login=None):
    """The password-free route. -> None once paired, else the reason, which means: fall
    back to copying the key with the password. on_login(user) runs before the prompt
    with the login properties.json names, so the fallback uses that user too."""
    try:
        login = fetch_login(host, port)
    except (OSError, ValueError) as e:
        return f"devkit service not reachable on port {port}: {why(e)}"
    if login and on_login:
        on_login(login)
    say("    In the headset: Steam Settings > Developer > Pair new host, then approve the request")
    body = register_body(pub, comment)
    ok, msg = register(host, body, port)
    # The headset refuses at once unless Steam is on its "Pair new host" screen
    # (verified on a Frame, 2026-09-26), so keep asking while the user opens it.
    deadline = time.monotonic() + PAIRING_MODE_WAIT
    while not ok and "pairing mode" in msg and time.monotonic() < deadline:
        time.sleep(3)
        ok, msg = register(host, body, port)
    return None if ok else f"devkit pairing failed: {msg}"


def split_port(arg):
    """"host:2222" -> ("host", 2222); anything else (IPv6 too) keeps port 22."""
    host, sep, port = arg.rpartition(":")
    if sep and port.isdigit() and ":" not in host:
        return host, int(port)
    return arg, 22


def port_open(host, port=22):
    try:
        with socket.create_connection((host, port), timeout=3):
            return True
    except OSError:
        return False


def reachable(host, port):
    """sshd, or the devkit service, which turns sshd on once a pairing is approved."""
    try:
        socket.getaddrinfo(host, port, type=socket.SOCK_STREAM)
    except OSError:
        return False
    return port_open(host, port) or port_open(host, DEVKIT_PORT)


# --- mDNS. There's no stdlib client, so this borrows dns-sd (macOS; Bonjour for
# Windows) or avahi-browse (Linux) when present, with short timeouts.

def run_for(args, seconds):
    """What a command printed within `seconds`; dns-sd never exits by itself."""
    try:
        out = subprocess.run(args, capture_output=True, timeout=seconds).stdout
    except subprocess.TimeoutExpired as e:
        out = e.stdout
    except OSError:
        out = b""
    return (out or b"").decode("utf-8", "replace")


def parse_dns_sd_browse(text):
    """Instance names from `dns-sd -B _steamos-devkit._tcp`, deduplicated, in order."""
    pat = re.compile(r"\sAdd\s+\d+\s+\d+\s+\S+\s+" + re.escape(DEVKIT_SERVICE) + r"\.\s+(.+?)\s*$")
    names = []
    for line in text.splitlines():
        m = pat.search(line)
        if m and m.group(1) not in names:
            names.append(m.group(1))
    return names


def parse_dns_sd_resolve(text):
    """The target host from `dns-sd -L` ("... can be reached at frame.local.:32000")."""
    m = re.search(r"can be reached at (\S+?)\.?:\d+", text)
    return m.group(1) if m else None


def parse_avahi(text):
    """Host names, then IPv4 addresses, from `avahi-browse -rpt` resolved ("=") lines."""
    names, addrs = [], []
    for line in text.splitlines():
        f = line.split(";")
        if len(f) >= 9 and f[0] == "=" and f[2] == "IPv4":
            names.append(f[6])
            addrs.append(f[7])
    return list(dict.fromkeys(names + addrs))


def discover_devkit():
    if shutil.which("dns-sd"):
        hosts = []
        for name in parse_dns_sd_browse(run_for(["dns-sd", "-B", DEVKIT_SERVICE, "local."], 3))[:4]:
            host = parse_dns_sd_resolve(run_for(["dns-sd", "-L", name, DEVKIT_SERVICE, "local."], 2))
            if host and host not in hosts:
                hosts.append(host)
        return hosts
    if shutil.which("avahi-browse"):
        return parse_avahi(run_for(["avahi-browse", "-rpt", DEVKIT_SERVICE], 5))
    return []


HOST_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9.:%-]*")


def pick_host(arg):
    if arg and not HOST_RE.fullmatch(arg):
        say(f"  - {arg!r} isn't a host name or IP address")
        return None
    for cand in [arg] if arg else [f"{FRAME_ALIAS}.local", FRAME_ALIAS]:
        host, port = split_port(cand)
        if reachable(host, port):
            return host, port
        say(f"  - {cand}: not resolvable, or ports {port} and {DEVKIT_PORT} closed")
    if arg:
        return None
    say(f"  - asking mDNS for {DEVKIT_SERVICE}")
    for host in discover_devkit():
        if HOST_RE.fullmatch(host) and reachable(host, 22):
            return host, 22
        say(f"  - {host}: advertised, but not reachable")
    return None


def make_ssh_dir():
    # Windows: no mode. Python 3.12.4+ turns 0o700 into an owner-only ACL, which locks
    # the user out if the folder's owner is Administrators; the profile's ACL suffices.
    if os.name == "nt":
        SSH_DIR.mkdir(exist_ok=True)
    else:
        SSH_DIR.mkdir(mode=0o700, exist_ok=True)


def make_key(path, kind, comment):
    if path.exists():
        say(f"    exists: {path}")
        return
    bits = ["-b", "3072"] if kind == "rsa" else []
    subprocess.run(["ssh-keygen", "-q", "-t", kind, *bits, "-N", "", "-C", comment, "-f", str(path)], check=True)
    say(f"    created {path}")


def config_block(host, port=22, user=FRAME_USER):
    return [BEGIN, f"Host {FRAME_ALIAS}", f"  HostName {host}", *([f"  Port {port}"] if port != 22 else []),
            f"  User {user}",
            "  IdentityFile ~/.ssh/id_ed25519_frame", "  IdentityFile ~/.ssh/id_rsa_frame_devkit",
            "  IdentitiesOnly yes", "  ServerAliveInterval 30", "Host *", END]


def write_config(host, port=22, user=FRAME_USER):
    """Replace our managed block and put it first: ssh uses the first value it sees per
    option. The trailing "Host *" returns the rest of the file to global scope."""
    make_ssh_dir()
    old = CONFIG.read_text(encoding="utf-8") if CONFIG.exists() else ""
    kept, skip = [], False
    for line in old.splitlines():
        if line == BEGIN:
            skip = True
        elif line == END:
            skip = False
        elif not skip:
            kept.append(line)
    block = config_block(host, port, user)
    tmp = CONFIG.with_name("config.frame-control.tmp")
    tmp.write_text("\n".join(block + kept) + "\n", encoding="utf-8")
    if os.name != "nt":
        tmp.chmod(0o600)
    # On Windows a running ssh.exe (Frame Control's own, say) keeps the config open
    # and locked, so the swap can fail for a moment; keep trying for a while.
    for attempt in range(60):
        try:
            os.replace(tmp, CONFIG)
            return
        except PermissionError:
            if attempt == 0:
                say("    ~/.ssh/config is in use by another ssh; waiting for it...")
            time.sleep(0.5)
    tmp.unlink(missing_ok=True)
    raise SystemExit("~/.ssh/config stayed locked by another program. Quit Frame Control "
                     "and any ssh windows, then run the setup again.")


def key_login_works():
    # accept-new: trust a first-seen host key (as the copy step does); a changed one still fails.
    return subprocess.run(["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=5",
                           "-o", "StrictHostKeyChecking=accept-new", FRAME_ALIAS, "true"],
                          capture_output=True).returncode == 0


def configured_user():
    """The User in our managed block, so a re-run keeps one the headset named earlier."""
    if not CONFIG.exists():
        return None
    inside = False
    for line in CONFIG.read_text(encoding="utf-8").splitlines():
        if line in (BEGIN, END):
            inside = line == BEGIN
        elif inside and line.startswith("  User "):
            name = line[7:].strip()
            return name if NAME_RE.fullmatch(name) else None
    return None


def pair_with_devkit(host, port, user):
    """Try devkit pairing and confirm key login. -> (user, None) or (user, reason to fall back)."""
    say("==> Pairing through the headset's SteamOS devkit service (no password)")
    chosen = [user]

    def use_login(login):
        if login == chosen[0]:
            return
        if USER_FROM_ENV:
            say(f"    the headset logs in as '{login}'; keeping FRAME_USER={user}")
        else:
            chosen[0] = login
            say(f"    the headset logs in as '{login}'")
            write_config(host, port, login)

    try:
        pub = DEVKIT_KEY.with_suffix(".pub").read_text(encoding="utf-8")
    except OSError as e:
        return user, f"can't read the pairing key: {e}"
    reason = devkit_pair(host, pub, key_comment(platform.node()), on_login=use_login)
    if reason:
        return chosen[0], reason
    # The approval is what turns sshd on, so it may take a moment to answer.
    for _ in range(10):
        if key_login_works():
            return chosen[0], None
        time.sleep(1)
    return chosen[0], "paired, but key login still fails"


def main(argv):
    if argv and argv[0] in ("-h", "--help"):
        sys.exit(__doc__)
    say("==> Looking for the Steam Frame")
    found = pick_host(argv[0] if argv else None)
    while not found:
        say("Could not reach the Frame over SSH.")
        say("Check: Developer Mode on and a user password set; same network; no client isolation.")
        try:
            typed = input("Type the Frame's IP address (Quick Settings shows it), or press Enter to quit: ").strip()
        except EOFError:
            typed = ""
        if not typed:
            return 1
        found = pick_host(typed)
    host, port = found
    say(f"    found: {host}" + (f" port {port}" if port != 22 else ""))

    say("==> SSH keys")
    make_ssh_dir()
    make_key(KEY, "ed25519", f"{platform.node() or 'computer'}->steam-frame")
    make_key(DEVKIT_KEY, "rsa", key_comment(platform.node()))

    user = FRAME_USER if USER_FROM_ENV else (configured_user() or FRAME_USER)
    say(f"==> ~/.ssh/config alias '{FRAME_ALIAS}' -> {host}")
    write_config(host, port, user)

    say("==> Checking key login")
    if key_login_works():
        say("    key login already works")
    else:
        user, reason = pair_with_devkit(host, port, user)
        if not reason:
            say("    paired; key login OK")
        else:
            say(f"    {reason}; falling back to the password")
            say("    copying the key: enter the Developer Mode password when asked")
            pub = KEY.with_suffix(".pub").read_text(encoding="utf-8").strip()
            r = subprocess.run(["ssh", "-o", "StrictHostKeyChecking=accept-new", "-o", "PubkeyAuthentication=no",
                                "-p", str(port), f"{user}@{host}", ADD_KEY_CMD], input=pub + "\n", text=True)
            if r.returncode != 0 or not key_login_works():
                say("Key login still isn't working. Check the password and run this again.")
                return 1
            say("    key login OK")
    say(f"\nDone. Frame Control can reach the Frame now. In a terminal: ssh {FRAME_ALIAS}")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))