browser-cli/browser_cli/serve/control.py

"""Built-in control commands handled directly by ``browser-cli serve``."""
from __future__ import annotations

import re
from pathlib import Path

from browser_cli.serve.logging import log_request

class ServeControlMixin:
  addr: tuple
  command: str
  auth_keys_path: Path | None
  auth_label: str | None

  async def send_error(self, msg: str, msg_id=None) -> None: ...
  async def send_ok(self, payload, command: str | None = None) -> None: ...

  async def handle_control_command(self, msg: dict) -> bool:
    if self.command == "browser-cli.targets":
      from browser_cli.client import active_browser_targets, send_command
      targets = []
      for target in active_browser_targets(include_remotes=False):
        item = {"profile": target.profile, "displayName": target.display_name}
        try:
          clients = send_command("clients.list", profile=target.profile, suppress_pq_warning=True)
          if clients:
            # Carry the full client info so a remote `clients` command can render
            # from this single roundtrip instead of issuing another clients.list.
            info = clients[0]
            for src, dst in (("name", "browserName"), ("version", "version"), ("extensionVersion", "extensionVersion")):
              value = info.get(src)
              if value:
                item[dst] = value
        except Exception:
          pass
        targets.append(item)
      await self.send_ok(targets, self.command)
      log_request(self.addr, self.command, None, "OK", identity=self.auth_label)
      return True

    if self.command == "browser-cli.auth.keys":
      if self.auth_keys_path is None:
        await self.send_error("no authorized keys file configured on this server")
        log_request(self.addr, self.command, None, "ERROR", "no authorized keys file", identity=self.auth_label)
        return True
      from browser_cli.auth import load_authorized_keys_with_policies
      entries = [
        {"pubkey": pk, "name": name, "allow": cats}
        for pk, name, cats in load_authorized_keys_with_policies(self.auth_keys_path)
      ]
      await self.send_ok(entries, self.command)
      log_request(self.addr, self.command, None, "OK", identity=self.auth_label)
      return True

    if self.command == "browser-cli.auth.trust":
      return await self._handle_trust(msg)

    if self.command == "browser-cli.auth.policy":
      return await self._handle_policy(msg)
    return False

  async def _handle_trust(self, msg: dict) -> bool:
    if self.auth_keys_path is None:
      await self.send_error("no authorized keys file configured on this server")
      log_request(self.addr, self.command, None, "ERROR", "no authorized keys file")
      return True
    from browser_cli.auth import add_authorized_key
    args = msg.get("args") or {}
    pubkey = str(args.get("pubkey") or "")
    name = str(args.get("name") or "")
    categories = args.get("allow")
    if not re.fullmatch(r"[0-9a-f]{64}", pubkey):
      await self.send_error("invalid pubkey: expected 64 lowercase hex characters")
      log_request(self.addr, self.command, None, "ERROR", "invalid pubkey", identity=self.auth_label)
      return True
    if not await self._validate_categories(categories):
      return True
    added = add_authorized_key(self.auth_keys_path, pubkey, name, categories)
    await self.send_ok({"added": added}, self.command)
    log_request(self.addr, self.command, None, "OK" if added else "ALREADY_TRUSTED", identity=self.auth_label)
    return True

  async def _handle_policy(self, msg: dict) -> bool:
    if self.auth_keys_path is None:
      await self.send_error("no authorized keys file configured on this server")
      log_request(self.addr, self.command, None, "ERROR", "no authorized keys file")
      return True
    from browser_cli.auth import set_authorized_key_policy
    args = msg.get("args") or {}
    identifier = str(args.get("identifier") or "")
    categories = args.get("allow")
    if not identifier.strip():
      await self.send_error("missing key identifier")
      log_request(self.addr, self.command, None, "ERROR", "missing identifier", identity=self.auth_label)
      return True
    if not await self._validate_categories(categories):
      return True
    try:
      updated = set_authorized_key_policy(self.auth_keys_path, identifier, categories)
    except ValueError as exc:
      await self.send_error(str(exc))
      log_request(self.addr, self.command, None, "ERROR", "ambiguous key", identity=self.auth_label)
      return True
    if updated is None:
      await self.send_error(f"trusted key not found: {identifier}")
      log_request(self.addr, self.command, None, "ERROR", "key not found", identity=self.auth_label)
      return True
    pubkey, name = updated
    await self.send_ok({"updated": True, "pubkey": pubkey, "name": name, "allow": categories}, self.command)
    log_request(self.addr, self.command, None, "OK", identity=self.auth_label)
    return True

  async def _validate_categories(self, categories) -> bool:
    if categories is not None and not isinstance(categories, list):
      await self.send_error("invalid allow: expected a list of category strings")
      log_request(self.addr, self.command, None, "ERROR", "invalid allow", identity=self.auth_label)
      return False
    if categories is not None:
      from browser_cli.serve.security import policy_from_categories
      try:
        policy_from_categories(categories)  # validate before persisting
      except ValueError as exc:
        await self.send_error(str(exc))
        log_request(self.addr, self.command, None, "ERROR", "invalid allow category", identity=self.auth_label)
        return False
    return True