wabot/src/wabot/hosts.py

"""Driver hosts: the processes that keep a browser alive, and how to reach them.

Persistence-capable browsers are always driven over the WebDriver wire
protocol (a URL). Two hosts provide that URL: ``ExternalServer`` (a
Selenium Grid or standalone driver the user runs) and ``ManagedService``
(a chromedriver/geckodriver wabot spawns detached so it outlives Python).
"""

from __future__ import annotations

import http.client
import json
import logging
import os
import shutil
import signal
import socket
import subprocess
import time
import urllib.request
from pathlib import Path

import platformdirs
from selenium import webdriver

LOGGER = logging.getLogger("wabot")

BROWSER_ALIASES = {"chrome": "chromium"}
DRIVER_BINARIES = {"chromium": "chromedriver", "firefox": "geckodriver"}


def normalize_browser(browser: str) -> str:
    browser = BROWSER_ALIASES.get(browser, browser)
    if browser not in DRIVER_BINARIES:
        raise ValueError(
            f"unsupported browser {browser!r}; use 'firefox' or 'chromium' (alias: 'chrome')"
        )
    return browser


def build_options(
    browser: str, *, headless: bool = False, user_agent: str | None = None
) -> webdriver.ChromeOptions | webdriver.FirefoxOptions:
    """Browser-appropriate Options. UA defaults to the browser's own."""
    browser = normalize_browser(browser)
    if browser == "chromium":
        opts = webdriver.ChromeOptions()
        binary = _chromium_binary()
        if binary:
            opts.binary_location = binary
        if headless:
            opts.add_argument("--headless=new")
        if user_agent:
            opts.add_argument(f"--user-agent={user_agent}")
        return opts
    opts = webdriver.FirefoxOptions()
    if headless:
        opts.add_argument("-headless")
    if user_agent:
        opts.set_preference("general.useragent.override", user_agent)
    return opts


def _chromium_binary() -> str | None:
    for name in ("chromium", "chromium-browser", "google-chrome", "chrome"):
        found = shutil.which(name)
        if found:
            return found
    return None


def service_alive(url: str, timeout: float = 2.0) -> bool:
    """True if a WebDriver server answers GET /status at ``url``.

    Liveness only: geckodriver reports ready=false while its single session
    is in use, so the ``ready`` flag is deliberately ignored.
    The ``timeout`` bounds each socket operation, not total wall-clock time.
    """
    try:
        with urllib.request.urlopen(f"{url.rstrip('/')}/status", timeout=timeout) as resp:
            json.loads(resp.read())
            return resp.status == 200
    except (OSError, ValueError, http.client.HTTPException):
        return False


class ExternalServer:
    """A WebDriver server somebody else runs (Selenium Grid, bare driver)."""

    def __init__(self, url: str):
        self.url = url.rstrip("/")

    def ensure_running(self) -> str:
        if not service_alive(self.url):
            raise ConnectionError(f"no WebDriver server answering at {self.url}/status")
        return self.url


def find_driver_binary(browser: str) -> str:
    """Locate chromedriver/geckodriver: env override, PATH, then Selenium Manager.

    The Selenium Manager fallback may download a driver over the network.
    """
    browser = normalize_browser(browser)
    name = DRIVER_BINARIES[browser]
    override = os.environ.get(f"WABOT_{name.upper()}")
    if override:
        return override
    found = shutil.which(name)
    if found:
        return found
    try:
        # Selenium Manager is bundled with selenium but explicitly beta
        # ("may change") — keep it as a guarded fallback, never the primary.
        from selenium.webdriver.common.selenium_manager import SeleniumManager

        args = ["--browser", "chrome" if browser == "chromium" else browser]
        chromium = _chromium_binary()
        if browser == "chromium" and chromium:
            args += ["--browser-path", chromium]
        return SeleniumManager().binary_paths(args)["driver_path"]
    # Selenium Manager is beta; any failure folds into a clean FileNotFoundError
    except Exception as ex:
        raise FileNotFoundError(
            f"could not find {name!r} on PATH (set WABOT_{name.upper()} to override); "
            f"Selenium Manager fallback also failed: {ex}"
        ) from ex


def _free_port() -> int:
    with socket.socket() as sock:
        sock.bind(("127.0.0.1", 0))
        return sock.getsockname()[1]


class ManagedService:
    """A driver service wabot spawns detached, so it outlives this process.

    ensure_running() is not idempotent: each call spawns a new service.
    """

    def __init__(
        self,
        browser: str,
        log_dir: Path | None = None,
        driver_binary: str | None = None,
        startup_timeout: float = 15.0,
    ):
        self.browser = normalize_browser(browser)
        self.log_dir = Path(log_dir) if log_dir else Path(platformdirs.user_data_dir("wabot"))
        self.driver_binary = driver_binary
        self.startup_timeout = startup_timeout
        self.pid: int | None = None
        self.port: int | None = None

    def ensure_running(self) -> str:
        binary = self.driver_binary or find_driver_binary(self.browser)
        self.port = _free_port()
        url = f"http://127.0.0.1:{self.port}"
        self.log_dir.mkdir(parents=True, exist_ok=True)
        log_path = self.log_dir / f"{DRIVER_BINARIES[self.browser]}-{self.port}.log"
        with open(log_path, "ab") as log:
            process = subprocess.Popen(
                [binary, f"--port={self.port}"],
                stdout=log,
                stderr=log,
                stdin=subprocess.DEVNULL,
                start_new_session=True,  # own session: survives this python process
            )
        self.pid = process.pid
        LOGGER.info("spawned %s (pid %s) on %s, log %s", binary, self.pid, url, log_path)
        deadline = time.monotonic() + self.startup_timeout
        while time.monotonic() < deadline:
            if service_alive(url):
                return url
            returncode = process.poll()
            if returncode is not None:
                raise RuntimeError(
                    f"{binary} exited early with code {returncode} before "
                    f"answering /status on {url} (see {log_path})"
                )
            time.sleep(0.2)
        # deadline expired but the driver is still running and unresponsive:
        # terminate and reap it so we never leak a process that would otherwise
        # outlive the interpreter.
        process.terminate()
        try:
            process.wait(timeout=5)
        except subprocess.TimeoutExpired:
            process.kill()
            process.wait()
        raise TimeoutError(
            f"{binary} did not answer /status on {url} within "
            f"{self.startup_timeout}s (see {log_path})"
        )


def stop_service(pid: int) -> bool:
    """SIGTERM a managed driver service. False if it was already gone.

    If the service is a child of THIS process (the same-process spawn→stop
    lifecycle), it is reaped to avoid a zombie. In the cross-process reattach
    case the driver is not our child and init reaps it. A driver that ignores
    SIGTERM is left running (no SIGKILL escalation).
    """
    try:
        os.kill(pid, signal.SIGTERM)
    except (ProcessLookupError, PermissionError):
        return False
    deadline = time.monotonic() + 5.0
    while time.monotonic() < deadline:
        try:
            reaped, _ = os.waitpid(pid, os.WNOHANG)
        except ChildProcessError:
            break  # not our child (cross-process) — init reaps it
        if reaped:
            break
        time.sleep(0.05)
    return True