browser-cli/browser_cli/commands/__init__.py

"""Shared helpers for the Click CLI command modules.

Every CLI command is a thin presentation layer over the Python SDK: it builds a
:class:`~browser_cli.BrowserCLI` from the global ``--browser/--remote/--key``
options, calls the matching SDK namespace method, and renders the result. The
SDK is the single source of truth for command strings, argument shapes, and
multi-browser routing.
"""
import functools
import click
from rich.console import Console
from rich.table import Table

from browser_cli import BrowserCLI, BrowserCounts
from browser_cli.client import BrowserNotConnected
from browser_cli.constants import GENTLE_MODES

_console = Console()

# Reusable ``--tab`` option: select a tab by ID (default: the active tab).
tab_option = click.option("--tab", "tab_id", type=int, default=None, help="Tab ID (default: active tab)")


def gentle_mode_option(help_text: str):
  """Reusable ``--gentle-mode`` Click option (throttle mode for large operations)."""
  return click.option(
    "--gentle-mode",
    type=click.Choice(GENTLE_MODES),
    default="auto",
    show_default=True,
    help=help_text,
  )

def command_policy_options(fn):
  """Reusable raw-command safety flags for /command-like entry points."""
  fn = click.option(
    "--allow-all",
    is_flag=True,
    help="Allow every command (equivalent to --allow-read-page --allow-control --allow-dangerous --allow-keys)",
  )(fn)
  fn = click.option(
    "--allow-keys",
    is_flag=True,
    help="Allow key-management commands (list/trust authorized keys over --remote)",
  )(fn)
  fn = click.option(
    "--allow-dangerous",
    is_flag=True,
    help="Allow high-risk commands such as dom.eval, storage.*, screenshots",
  )(fn)
  fn = click.option(
    "--allow-control",
    is_flag=True,
    help="Allow browser-control commands such as nav.*, tabs.close, dom.click",
  )(fn)
  fn = click.option(
    "--allow-read-page",
    is_flag=True,
    help="Allow page-content read commands such as extract.* and dom.text",
  )(fn)
  return fn

def command_policy_from_options(*, allow_read_page: bool, allow_control: bool, allow_dangerous: bool, allow_keys: bool = False, allow_all: bool = False):
  """Build a CommandPolicy from shared raw-command safety flags."""
  from browser_cli.command_security import CommandPolicy

  if allow_all:
    return CommandPolicy.unrestricted()
  return CommandPolicy(
    allow_read_page=allow_read_page,
    allow_control=allow_control,
    allow_dangerous=allow_dangerous,
    allow_keys=allow_keys,
  )

def command_categories_from_options(*, allow_read_page: bool, allow_control: bool, allow_dangerous: bool, allow_keys: bool = False, allow_all: bool = False):
  """Convert the shared --allow-* flags into a category list, or None if none were set.

  None means "no explicit policy" — the key falls back to the server-wide default.
  """
  if allow_all:
    return ["all"]
  cats = []
  if allow_read_page:
    cats.append("read-page")
  if allow_control:
    cats.append("control")
  if allow_dangerous:
    cats.append("dangerous")
  if allow_keys:
    cats.append("keys")
  return cats or None

def print_counts(result, noun: str, *, single_suffix: str = "") -> None:
  """Render a count result.

  In multi-browser mode (*result* is a :class:`~browser_cli.BrowserCounts`) print a
  per-browser table with a Total row; otherwise print a single ``N noun(s)`` line.
  """
  if isinstance(result, BrowserCounts):
    table = Table(show_header=True, header_style="bold cyan")
    table.add_column("Browser", no_wrap=True)
    table.add_column(f"{noun.capitalize()}s", justify="right")
    rendered_groups: set[str] = set()
    for name, count in result.by_browser.items():
      group = result.browser_groups.get(name)
      if group:
        if group not in rendered_groups:
          group_total = sum(
            browser_count
            for browser_name, browser_count in result.by_browser.items()
            if result.browser_groups.get(browser_name) == group
          )
          table.add_row(f"[bold]{group}[/bold]", str(group_total))
          rendered_groups.add(group)
        display_name = name.removeprefix(f"{group}:")
        table.add_row(f"  {display_name}", str(count))
      else:
        table.add_row(name, str(count))
    table.add_row("Total", str(result.total))
    _console.print(table)
  else:
    _console.print(f"[bold]{result}[/bold] {noun}(s){single_suffix}")

def client_from_ctx() -> BrowserCLI:
  """Build a BrowserCLI from the root context's global options.

  Reads ``browser``/``remote``/``key`` set by the top-level ``main`` group.
  Falls back to an unconfigured client when a command group is invoked
  standalone (e.g. in unit tests).
  """
  obj = click.get_current_context().find_root().obj or {}
  return BrowserCLI(browser=obj.get("browser"), remote=obj.get("remote"), key=obj.get("key"))

def handle_errors(fn):
  """Decorate a CLI command so SDK exceptions become clean errors + exit(1).

  Apply as the innermost decorator (directly above ``def``) so Click's option
  decorators attach their params to the wrapper.
  """
  @functools.wraps(fn)
  def wrapper(*args, **kwargs):
    try:
      return fn(*args, **kwargs)
    except BrowserNotConnected as e:
      _console.print(f"[red]Error:[/red] {e}")
      raise SystemExit(1)
    except PermissionError as e:
      _console.print(f"[red]Blocked:[/red] {e}")
      raise SystemExit(1)
    except RuntimeError as e:
      _console.print(f"[red]Browser error:[/red] {e}")
      raise SystemExit(1)

  return wrapper