browser-cli/browser_cli/cli.py

#!/usr/bin/env -S uv run
"""
browser-cli — Control your running browser from the terminal.
"""
import click
import sys
import os
import json
import stat
import shutil
import re
from importlib.metadata import PackageNotFoundError, version as package_version
from pathlib import Path
from rich.console import Console

from browser_cli.commands.navigate import nav_group
from browser_cli.commands.tabs import tabs_group
from browser_cli.commands.groups import group_group
from browser_cli.commands.windows import windows_group
from browser_cli.commands.dom import dom_group
from browser_cli.commands.extract import extract_group
from browser_cli.commands.session import session_group
from browser_cli.commands.search import search_group
from browser_cli.commands.page import page_group
from browser_cli.commands.storage import storage_group
from browser_cli.commands.cookies import cookies_group
from browser_cli.commands.serve import cmd_serve
from browser_cli.client import (
    send_command,
    BrowserNotConnected,
    REGISTRY_PATH,
    active_browser_targets,
    display_browser_name,
)
from browser_cli.platform import install_base_dir, is_windows

console = Console()

NATIVE_HOST_NAME = "com.browsercli.host"
EXTENSION_ID = "bfpmkhngkjnfhabmfckgeohlilokodkg"

NATIVE_HOST_DIRS = {
    "chrome": {
        "linux":  [Path.home() / ".config/google-chrome/NativeMessagingHosts"],
        "darwin": [Path.home() / "Library/Application Support/Google/Chrome/NativeMessagingHosts"],
    },
    "chromium": {
        "linux":  [Path.home() / ".config/chromium/NativeMessagingHosts"],
        "darwin": [Path.home() / "Library/Application Support/Chromium/NativeMessagingHosts"],
    },
    "brave": {
        "linux":  [Path.home() / ".config/BraveSoftware/Brave-Browser/NativeMessagingHosts"],
        "darwin": [Path.home() / "Library/Application Support/BraveSoftware/Brave-Browser/NativeMessagingHosts"],
    },
    "edge": {
        "linux":  [Path.home() / ".config/microsoft-edge/NativeMessagingHosts"],
        "darwin": [Path.home() / "Library/Application Support/Microsoft Edge/NativeMessagingHosts"],
    },
    "vivaldi": {
        "linux":  [Path.home() / ".config/vivaldi/NativeMessagingHosts"],
        "darwin": [Path.home() / "Library/Application Support/Vivaldi/NativeMessagingHosts"],
    },
}

WINDOWS_NATIVE_HOST_REGISTRY_KEYS = {
    "chrome": [r"Software\Google\Chrome\NativeMessagingHosts"],
    "chromium": [r"Software\Chromium\NativeMessagingHosts"],
    "brave": [r"Software\BraveSoftware\Brave-Browser\NativeMessagingHosts"],
    "edge": [r"Software\Microsoft\Edge\NativeMessagingHosts"],
    "vivaldi": [r"Software\Vivaldi\NativeMessagingHosts"],
}


def _rename_target_profile(target_browser: str | None) -> str | None:
    if target_browser:
        return target_browser

    active = active_browser_targets()
    if len(active) == 1:
        return active[0].profile
    return None


def _ensure_unique_browser_alias(alias: str, target_browser: str | None) -> None:
    target_profile = _rename_target_profile(target_browser)

    profiles: dict[str, str] = {}
    if REGISTRY_PATH.exists():
        try:
            profiles = json.loads(REGISTRY_PATH.read_text())
        except Exception:
            profiles = {}

    if alias in profiles and alias != target_profile:
        raise click.ClickException(f"Browser alias '{alias}' already exists")


def _native_host_wrapper_path() -> Path:
    base_dir = install_base_dir()
    if is_windows():
        return base_dir / "libexec" / "native-host.cmd"
    return base_dir / "libexec" / "native-host"


def _native_host_script_path() -> Path:
    return _native_host_wrapper_path().with_name("native_host.py")


def _windows_registry_views():
    import winreg

    return [0, getattr(winreg, "KEY_WOW64_32KEY", 0), getattr(winreg, "KEY_WOW64_64KEY", 0)]


def _register_windows_native_host(browser: str, manifest_path: Path) -> list[str]:
    import winreg

    installed = []
    for key_path in WINDOWS_NATIVE_HOST_REGISTRY_KEYS[browser]:
        full_key = f"{key_path}\\{NATIVE_HOST_NAME}"
        for view in _windows_registry_views():
            try:
                access = winreg.KEY_WRITE | view
                key = winreg.CreateKeyEx(winreg.HKEY_CURRENT_USER, full_key, 0, access)
                with key:
                    winreg.SetValueEx(key, "", 0, winreg.REG_SZ, str(manifest_path))
                installed.append(f"HKCU\\{full_key}")
            except OSError as e:
                console.print(f"[yellow]Could not write registry key {full_key}: {e}[/yellow]")
    return installed


def _project_version() -> str:
    pyproject_path = Path(__file__).resolve().parent.parent / "pyproject.toml"
    try:
        content = pyproject_path.read_text(encoding="utf-8")
        match = re.search(r'^version\s*=\s*"([^"]+)"', content, re.MULTILINE)
        if match:
            return match.group(1)
    except OSError:
        pass

    try:
        return package_version("browser-cli")
    except PackageNotFoundError:
        return "unknown"


def _print_version(ctx, param, value):
    if not value or ctx.resilient_parsing:
        return
    click.echo(_project_version())
    ctx.exit()

@click.group()
@click.option(
    "-V", "--version",
    is_flag=True,
    is_eager=True,
    expose_value=False,
    callback=_print_version,
    help="Show the browser-cli version and exit.",
)
@click.option(
    "--browser", default=None, metavar="ALIAS",
    help="Browser profile alias to target (required when multiple browsers are active).",
)
@click.option(
    "--remote", default=None, metavar="HOST:PORT",
    help="Connect to a remote browser exposed via 'browser-cli serve'.",
)
@click.option(
    "--token", default=None, metavar="TOKEN",
    help="Auth token for the remote browser-cli serve instance.",
)
@click.pass_context
def main(ctx, browser, remote, token):
    """Control your running browser from the terminal via a Chrome extension."""
    ctx.ensure_object(dict)
    ctx.obj["browser"] = browser
    ctx.obj["browser_explicit"] = browser is not None
    if browser:
        os.environ["BROWSER_CLI_PROFILE"] = browser
    ctx.obj["remote"] = remote
    ctx.obj["token"] = token
    if remote:
        os.environ["BROWSER_CLI_REMOTE"] = remote
    if token:
        os.environ["BROWSER_CLI_TOKEN"] = token


# ── Sub-command groups ─────────────────────────────────────────────────────────
main.add_command(nav_group)
main.add_command(tabs_group)
main.add_command(group_group)
main.add_command(windows_group)
main.add_command(dom_group)
main.add_command(extract_group)
main.add_command(session_group)
main.add_command(search_group)
main.add_command(page_group)
main.add_command(storage_group)
main.add_command(cookies_group)
main.add_command(cmd_serve)


# ── clients ────────────────────────────────────────────────────────────────────

@click.group("clients", invoke_without_command=True)
@click.pass_context
def clients_group(ctx):
    """Inspect and manage connected browser clients."""
    if ctx.invoked_subcommand is not None:
        return

    all_clients = []

    remote = (ctx.obj or {}).get("remote") or os.environ.get("BROWSER_CLI_REMOTE")
    if remote:
        try:
            result = send_command("clients.list", profile=(ctx.obj or {}).get("browser"))
            for c in (result or []):
                c["profile"] = c.get("profile") or (ctx.obj or {}).get("browser") or "remote"
                all_clients.append(c)
        except (BrowserNotConnected, RuntimeError) as e:
            console.print(f"[red]Error:[/red] {e}")
            sys.exit(1)
    else:
        profiles: dict[str, str] = {}
        if REGISTRY_PATH.exists():
            try:
                profiles = json.loads(REGISTRY_PATH.read_text())
            except Exception:
                pass

        for profile_name, sock_path in profiles.items():
            display_profile = display_browser_name(profile_name, sock_path)
            try:
                result = send_command("clients.list", profile=profile_name)
                for c in (result or []):
                    c["profile"] = display_profile
                    all_clients.append(c)
            except (BrowserNotConnected, RuntimeError):
                # Socket registered but browser no longer connected
                all_clients.append({
                    "profile": display_profile,
                    "name": "—",
                    "version": "—",
                    "extensionVersion": "disconnected",
                })

    if not all_clients:
        console.print("[yellow]No browser clients found. Start a browser with the extension enabled first.[/yellow]")
        sys.exit(1)

    from rich.table import Table
    table = Table(show_header=True, header_style="bold cyan")
    table.add_column("Profile")
    table.add_column("Browser")
    table.add_column("Version")
    table.add_column("Extension Version")
    for c in all_clients:
        table.add_row(
            c.get("profile", ""),
            c.get("name", ""),
            c.get("version", ""),
            c.get("extensionVersion", ""),
        )
    console.print(table)


main.add_command(clients_group)


@clients_group.command("rename")
@click.option(
    "--browser", "target_browser", default=None, metavar="ALIAS",
    help="Browser profile alias to rename. Overrides the global --browser option for this command.",
)
@click.argument("alias")
def cmd_clients_rename(target_browser, alias):
    """Set the profile alias used to identify this browser instance."""
    try:
        _ensure_unique_browser_alias(alias, target_browser)
        send_command("clients.rename_profile", {"alias": alias}, profile=target_browser)
    except BrowserNotConnected as e:
        console.print(f"[red]Error:[/red] {e}")
        sys.exit(1)
    console.print(f"[green]Profile renamed to '{alias}'[/green]")


# ── install ────────────────────────────────────────────────────────────────────

@main.command("install")
@click.argument("browser", type=click.Choice(["chrome", "chromium", "brave", "edge", "vivaldi"]), default="chrome")
def cmd_install(browser):
    """Register the native messaging host and print extension load instructions."""

    # Install wrapper outside PATH — the browser uses the absolute path from the
    # native messaging manifest, so only `browser-cli` needs to be on PATH.
    wrapper_path = _native_host_wrapper_path()
    native_host_script_path = _native_host_script_path()
    wrapper_path.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(Path(__file__).with_name("native_host.py"), native_host_script_path)
    if not is_windows():
        native_host_script_path.chmod(
            native_host_script_path.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
        )
        wrapper_content = f'#!/bin/sh\nexec "{sys.executable}" "{native_host_script_path}" "$@"\n'
        wrapper_path.write_text(wrapper_content)
        wrapper_path.chmod(wrapper_path.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)
    else:
        wrapper_content = f'@echo off\r\n"{sys.executable}" "{native_host_script_path}" %*\r\n'
        wrapper_path.write_text(wrapper_content, encoding="utf-8")

    # Load extension
    ext_urls = {
        "chrome": "chrome://extensions",
        "chromium": "chrome://extensions",
        "brave": "brave://extensions",
        "edge": "edge://extensions",
        "vivaldi": "vivaldi://extensions",
    }
    ext_url = ext_urls[browser]
    console.print("\n[bold]Step 1:[/bold] Load the extension in your browser")
    console.print(f"  1. Open [cyan]{ext_url}[/cyan]")
    console.print("  2. Enable [bold]Developer mode[/bold] (top-right toggle)")
    console.print(f"  3. Click [bold]Load unpacked[/bold] → select: [cyan]{Path(__file__).parent.parent / 'extension'}[/cyan]")
    console.print(f"  4. Extension ID will be [cyan]{EXTENSION_ID}[/cyan] (fixed by built-in key)\n")

    extension_id = EXTENSION_ID

    # Build native messaging manifest
    manifest = {
        "name": NATIVE_HOST_NAME,
        "description": "browser-cli native messaging host",
        "path": str(wrapper_path),
        "type": "stdio",
        "allowed_origins": [f"chrome-extension://{extension_id}/"],
    }

    installed = []
    if is_windows():
        manifest_dir = wrapper_path.parent
        manifest_dir.mkdir(parents=True, exist_ok=True)
        manifest_path = manifest_dir / f"{NATIVE_HOST_NAME}.json"
        manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")
        installed = _register_windows_native_host(browser, manifest_path)
    else:
        platform = "darwin" if sys.platform == "darwin" else "linux"
        dirs = NATIVE_HOST_DIRS[browser][platform]

        for d in dirs:
            try:
                d.mkdir(parents=True, exist_ok=True)
                manifest_path = d / f"{NATIVE_HOST_NAME}.json"
                manifest_path.write_text(json.dumps(manifest, indent=2))
                installed.append(manifest_path)
            except Exception as e:
                console.print(f"[yellow]Could not write to {d}: {e}[/yellow]")

    if not installed:
        console.print("[red]Failed to install native host manifest[/red]")
        sys.exit(1)

    for p in installed:
        if is_windows():
            console.print(f"[green]✓[/green] Registered native host: {p}")
        else:
            console.print(f"[green]✓[/green] Wrote native host manifest: {p}")
    console.print(f"[green]✓[/green] Installed native host script: {native_host_script_path}")
    console.print(f"[green]✓[/green] Installed native host wrapper: {wrapper_path}")
    if is_windows():
        console.print("\n[green]✓[/green] Wrote native host manifest:", manifest_path)

    console.print(f"\n[bold]Step 2:[/bold] Restart {browser.capitalize()} completely (quit app, then reopen)")
    console.print("\n[green bold]✓ Installation complete![/green bold]")
    console.print("  After restarting the browser, try: [cyan]browser-cli tabs list[/cyan]")


# ── native-host (hidden, called by Chrome via native messaging) ────────────────

@main.command("native-host", hidden=True)
def cmd_native_host():
    """Native messaging host — called by Chrome, not for direct use."""
    from browser_cli.native_host import main as _main
    _main()


# ── completion ─────────────────────────────────────────────────────────────────

@main.command("completion")
@click.argument("shell", type=click.Choice(["zsh", "bash", "fish"]))
@click.option("--script", is_flag=True, help="Output the raw completion script instead of instructions")
def cmd_completion(shell, script):
    """Print shell completion setup instructions (or output the script with --script)."""
    if script:
        from click.shell_completion import BashComplete, ZshComplete, FishComplete
        cls = {"zsh": ZshComplete, "bash": BashComplete, "fish": FishComplete}[shell]
        comp = cls(main, {}, "browser-cli", "_BROWSER_CLI_COMPLETE")
        click.echo(comp.source())
        return

    exe = sys.executable.replace("/python", "/browser-cli").replace("/python3", "/browser-cli")
    if not Path(exe).exists():
        exe = "browser-cli"

    env_var = "_BROWSER_CLI_COMPLETE"

    if shell == "zsh":
        console.print("[bold]Quickest setup — generate the file once:[/bold]")
        console.print()
        console.print(f"  [cyan]uv run browser-cli completion zsh --script > ~/.zfunc/_browser-cli[/cyan]")
        console.print()
        console.print("  Then add these lines to [bold]~/.zshrc[/bold] (before any compinit call):")
        console.print("    [cyan]fpath=(~/.zfunc $fpath)[/cyan]")
        console.print("    [cyan]autoload -Uz compinit && compinit[/cyan]")
        console.print()
        console.print("  Reload: [cyan]exec zsh[/cyan]")
        console.print()
        console.print("[bold]Alternative — eval on every shell start (simpler but slower):[/bold]")
        console.print(f'  [cyan]eval "$({env_var}=zsh_source {exe})"[/cyan]')
    elif shell == "bash":
        console.print("[bold]Quickest setup — generate the file once:[/bold]")
        console.print()
        console.print(f"  [cyan]uv run browser-cli completion bash --script > ~/.bash_completion.d/browser-cli[/cyan]")
        console.print()
        console.print("  Reload: [cyan]source ~/.bashrc[/cyan]")
        console.print()
        console.print("[bold]Alternative — eval on every shell start:[/bold]")
        console.print(f'  [cyan]eval "$({env_var}=bash_source {exe})"[/cyan]')
    elif shell == "fish":
        console.print("[bold]Setup:[/bold]")
        console.print()
        console.print(f"  [cyan]uv run browser-cli completion fish --script > ~/.config/fish/completions/browser-cli.fish[/cyan]")


if __name__ == "__main__":
    main()