browser-cli/browser_cli/remote/known_hosts.py

"""SSH-style known-hosts pinning for browser-cli remote servers."""
from __future__ import annotations

import json
import os
import sys
from pathlib import Path

from browser_cli.constants import CONFIG_DIR
from browser_cli.endpoints import _normalize_endpoint
from browser_cli.errors import BrowserNotConnected

KNOWN_HOSTS_PATH = CONFIG_DIR / "known_hosts.json"

def fingerprint(pubkey_hex: str) -> str:
  """Return a compact SHA256 fingerprint for display."""
  import base64
  import hashlib

  digest = hashlib.sha256(bytes.fromhex(pubkey_hex)).digest()
  return "SHA256:" + base64.b64encode(digest).decode("ascii").rstrip("=")

def load_known_hosts(path: Path | None = None) -> dict[str, str]:
  path = path or KNOWN_HOSTS_PATH
  if not path.exists():
    return {}
  try:
    data = json.loads(path.read_text(encoding="utf-8"))
  except Exception:
    return {}
  if not isinstance(data, dict):
    return {}
  return {_normalize_endpoint(str(endpoint)): str(pubkey) for endpoint, pubkey in data.items() if isinstance(pubkey, str)}

def save_known_host(endpoint: str, pubkey_hex: str, path: Path | None = None) -> None:
  path = path or KNOWN_HOSTS_PATH
  known = load_known_hosts(path)
  known[_normalize_endpoint(endpoint)] = pubkey_hex
  path.parent.mkdir(parents=True, exist_ok=True)
  fd = os.open(str(path), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
  with os.fdopen(fd, "w", encoding="utf-8") as file:
    file.write(json.dumps(known, indent=2, sort_keys=True) + "\n")

def remove_known_host(endpoint: str, path: Path | None = None) -> bool:
  path = path or KNOWN_HOSTS_PATH
  known = load_known_hosts(path)
  normalized = _normalize_endpoint(endpoint)
  if normalized not in known:
    return False
  del known[normalized]
  path.parent.mkdir(parents=True, exist_ok=True)
  fd = os.open(str(path), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
  with os.fdopen(fd, "w", encoding="utf-8") as file:
    file.write(json.dumps(known, indent=2, sort_keys=True) + "\n")
  return True

def _is_loopback_endpoint(endpoint: str) -> bool:
  host, sep, _port = endpoint.rpartition(":")
  check = host if sep else endpoint
  return check in {"127.0.0.1", "localhost", "::1"}

def verify_known_host(endpoint: str, challenge: dict | None) -> None:
  """Verify and pin the server identity from a challenge frame.

  First contact auto-adds the host when the process is interactive, mirroring
  SSH's trust-on-first-use flow. Non-interactive clients must pin explicitly via
  `browser-cli remote trust-host ENDPOINT`.
  """
  if not isinstance(challenge, dict):
    return
  pubkey = challenge.get("server_pubkey")
  if not isinstance(pubkey, str) or not pubkey:
    return

  from browser_cli.auth.server_identity import verify_challenge_signature
  if not verify_challenge_signature(challenge):
    raise BrowserNotConnected("Remote server identity signature is invalid")

  normalized = _normalize_endpoint(endpoint)
  known = load_known_hosts()
  expected = known.get(normalized)
  if expected is None and _is_loopback_endpoint(endpoint):
    return
  if expected is None:
    if not sys.stdin.isatty():
      raise BrowserNotConnected(
        f"Unknown remote server identity for {normalized} ({fingerprint(pubkey)}).\n"
        f"Run: browser-cli remote trust-host {normalized}"
      )
    sys.stderr.write(
      f"The authenticity of remote '{normalized}' can't be established.\n"
      f"Server key fingerprint is {fingerprint(pubkey)}.\n"
      "Trust this server and add it to known hosts? [y/N] "
    )
    answer = sys.stdin.readline().strip().lower()
    if answer not in {"y", "yes"}:
      raise BrowserNotConnected("Remote server identity was not trusted")
    save_known_host(normalized, pubkey)
    sys.stderr.write(f"Added {normalized} to browser-cli known hosts.\n")
    return

  if expected != pubkey:
    raise BrowserNotConnected(
      f"REMOTE SERVER IDENTITY CHANGED for {normalized}!\n"
      f"Known: {fingerprint(expected)}\n"
      f"Seen:  {fingerprint(pubkey)}\n"
      f"If this is expected, run: browser-cli remote untrust-host {normalized} && browser-cli remote trust-host {normalized}"
    )