browserrig 0.1.0__py3-none-any.whl

browserrig/browser/chrome.py

@@ -0,0 +1,242 @@
+import asyncio
+import logging
+from contextlib import AsyncExitStack
+from typing import Any, Self
+from uuid import uuid4
+
+import httpx2
+from cdpify import CdpClient, EventSubscription, ReceivedEvent
+from cdpify.domains.browser import DownloadProgressEvent, DownloadWillBeginEvent
+from cdpify.domains.page import (
+    FrameNavigatedEvent,
+    FrameStartedLoadingEvent,
+    FrameStartedNavigatingEvent,
+    FrameStoppedLoadingEvent,
+    JavascriptDialogClosedEvent,
+    JavascriptDialogOpeningEvent,
+    LifecycleEventEvent,
+    NavigatedWithinDocumentEvent,
+)
+from cdpify.domains.runtime import (
+    ExecutionContextDestroyedEvent,
+    ExecutionContextsClearedEvent,
+)
+from cdpify.domains.target import (
+    TargetDestroyedEvent,
+    TargetInfoChangedEvent,
+)
+
+from browserrig.core import BrowserClosed, BrowserError
+from browserrig.events import (
+    ConnectedEvent,
+    DisconnectedEvent,
+)
+from browserrig.files import Download, FileSystem
+from browserrig.input import Input, InputTarget, Keyboard, Mouse
+
+from .base import BrowserEvents, EventSink, _discard
+from .config import BrowserOptions, CdpEndpoint, Launch
+from .downloads import Downloads
+from .tab import Tab, TabListener
+from .tabs import Tabs
+
+logger = logging.getLogger(__name__)
+
+# Page and runtime events arrive on tab sessions; target and download events on the root.
+_EVENTS = (
+    FrameNavigatedEvent,
+    NavigatedWithinDocumentEvent,
+    FrameStartedLoadingEvent,
+    FrameStartedNavigatingEvent,
+    LifecycleEventEvent,
+    FrameStoppedLoadingEvent,
+    JavascriptDialogOpeningEvent,
+    JavascriptDialogClosedEvent,
+    ExecutionContextDestroyedEvent,
+    ExecutionContextsClearedEvent,
+    TargetInfoChangedEvent,
+    TargetDestroyedEvent,
+    DownloadWillBeginEvent,
+    DownloadProgressEvent,
+)
+
+
+class ChromeBrowser:
+    def __init__(
+        self,
+        options: BrowserOptions | None = None,
+        *,
+        emit: EventSink = _discard,
+        on_tab_change: TabListener | None = None,
+    ) -> None:
+        self._options = options if options is not None else Launch()
+        files = self._options.files
+        self.files = (
+            FileSystem(files, download_directory=self._options.download_directory)
+            if files is not None
+            else None
+        )
+        self.id = uuid4().hex
+        self.events = BrowserEvents(self.id, emit)
+        self._stack: AsyncExitStack | None = None
+        self._downloads = Downloads(
+            files=self.files,
+            directory=self._options.download_directory,
+            emit=self.events.emit,
+        )
+        self._tasks: list[asyncio.Task[None]] = []
+        self._failure: Exception | None = None
+        self._connected = False
+        self.action_lock = asyncio.Lock()
+        # One mouse and keyboard for the whole browser, aimed at the active tab.
+        self.input = Input(self._focused_input)
+        self.tabs = Tabs(
+            self._options,
+            self.input,
+            check=self._check,
+            emit=self.events.emit,
+            on_change=on_tab_change,
+        )
+        self._close_task: asyncio.Task[None] | None = None
+        self.client: CdpClient
+
+    async def __aenter__(self) -> Self:
+        if self._stack is not None or self._close_task is not None:
+            raise BrowserError("A ChromeBrowser instance can be opened only once")
+        self._stack = AsyncExitStack()
+        try:
+            options = self._options
+            if self.files is not None:
+                await self.files.initialize()
+            async with asyncio.timeout(options.timeout):
+                endpoint = await self._resolve_endpoint()
+                self.client = await self._stack.enter_async_context(
+                    CdpClient(
+                        endpoint.url,
+                        additional_headers=endpoint.headers,
+                        default_timeout=options.timeout,
+                    )
+                )
+                context = await self.tabs.start(self.client, self._stack)
+                if context is not None:
+                    await self._downloads.enable(self.client, context)
+                self._start_listeners()
+                self._stack.push_async_callback(self._stop_listeners)
+                await self.client.target.set_discover_targets(discover=True)
+                await self.tabs.open_initial()
+                self._connected = True
+                await self.events.emit(
+                    ConnectedEvent(owned_process=options.owns_process)
+                )
+            return self
+        except (Exception, asyncio.CancelledError):
+            await self.aclose()
+            raise
+
+    async def __aexit__(self, *_: object) -> None:
+        await self.aclose()
+
+    @property
+    def active_tab(self) -> Tab:
+        return self.tabs.active
+
+    @property
+    def mouse(self) -> Mouse:
+        return self.input.mouse
+
+    @property
+    def keyboard(self) -> Keyboard:
+        return self.input.keyboard
+
+    async def downloads(self) -> tuple[Download, ...]:
+        return self._downloads.all()
+
+    async def aclose(self) -> None:
+        if self._close_task is None:
+            self._close_task = asyncio.create_task(
+                self._close(), name="browserrig:close"
+            )
+        await asyncio.shield(self._close_task)
+
+    async def _resolve_endpoint(self) -> CdpEndpoint:
+        """Resolve an HTTP discovery endpoint to the browser's WebSocket."""
+        assert self._stack is not None
+        options = self._options
+        endpoint = await options.cdp_endpoint(self._stack)
+        url = endpoint.url
+        if url.startswith(("http://", "https://")):
+            async with httpx2.AsyncClient(
+                timeout=options.timeout, headers=endpoint.headers, trust_env=False
+            ) as http:
+                response = await http.get(url.rstrip("/") + "/json/version")
+                response.raise_for_status()
+                url = response.json()["webSocketDebuggerUrl"]
+        if "/devtools/page/" in url:
+            raise BrowserError(
+                "Discovery returned a page endpoint; a browser endpoint is needed"
+            )
+        return CdpEndpoint(url=url, headers=endpoint.headers)
+
+    def _check(self) -> None:
+        if self._stack is None or not self.client.is_connected:
+            raise BrowserClosed("Browser connection is closed")
+        if self._failure is not None:
+            raise BrowserError(
+                "Browser event listener failed; reconnect before continuing"
+            ) from self._failure
+
+    def _focused_input(self) -> InputTarget | None:
+        tab = self.tabs.focused()
+        return None if tab is None else tab.input
+
+    def _start_listeners(self) -> None:
+        for event_type in _EVENTS:
+            subscription = self.client.listen_all(event_type)
+            self._tasks.append(
+                asyncio.create_task(
+                    self._consume(subscription),
+                    name=f"browserrig:{event_type.__name__}",
+                )
+            )
+
+    async def _consume(
+        self, subscription: EventSubscription[ReceivedEvent[Any]]
+    ) -> None:
+        try:
+            async with subscription:
+                async for received in subscription:
+                    try:
+                        await self._handle(received.value, received.session_id)
+                    except Exception:  # pylint: disable=broad-exception-caught
+                        # A target can disappear while its event is being handled.
+                        logger.exception(
+                            "Browser event %s failed", type(received.value).__name__
+                        )
+        except Exception as error:  # pylint: disable=broad-exception-caught
+            # Surface listener failures on the next browser operation.
+            self._failure = error
+
+    async def _handle(self, event: Any, session_id: str | None) -> None:
+        if session_id is None and isinstance(
+            event, (DownloadWillBeginEvent, DownloadProgressEvent)
+        ):
+            await self._downloads.handle(event)
+        else:
+            await self.tabs.handle(event, session_id)
+
+    async def _stop_listeners(self) -> None:
+        for task in self._tasks:
+            task.cancel()
+        await asyncio.gather(*self._tasks, return_exceptions=True)
+        self._tasks.clear()
+
+    async def _close(self) -> None:
+        if self._stack is not None:
+            stack, self._stack = self._stack, None
+            try:
+                await stack.aclose()
+            finally:
+                await self.tabs.discard()
+                if self._connected:
+                    self._connected = False
+                    await self.events.emit(DisconnectedEvent())

browserrig/browser/config.py

@@ -0,0 +1,119 @@
+from contextlib import AsyncExitStack
+from dataclasses import dataclass
+from pathlib import Path
+from typing import ClassVar, Self
+from urllib.parse import urlsplit
+
+from pydantic import Field, InstanceOf, SecretStr, field_validator, model_validator
+
+from browserrig.core import Options, TargetId
+from browserrig.files import FileSystemOptions
+
+from .process import ChromeProcess
+from .setup import SessionSetup
+
+
+class Viewport(Options):
+    width: int = Field(default=1280, ge=320, le=7680)
+    height: int = Field(default=800, ge=240, le=4320)
+    scale: float = Field(default=1, gt=0, le=4)
+
+
+@dataclass(frozen=True, slots=True)
+class CdpEndpoint:
+    """Where to reach a browser over CDP, and the headers that request needs."""
+
+    url: str
+    headers: dict[str, str]
+
+
+class Launch(Options):
+    """Start a Chrome process that browserrig owns and stops again."""
+
+    owns_process: ClassVar[bool] = True
+    # A launched browser always starts in a fresh tab of its own context.
+    target_id: ClassVar[None] = None
+    executable: Path | None = None
+    headless: bool = True
+    url: str = "about:blank"
+    viewport: Viewport = Viewport()
+    session_setup: tuple[InstanceOf[SessionSetup], ...] = ()
+    timeout: float = Field(default=20, gt=0, le=120)
+    arguments: tuple[str, ...] = ()
+    download_directory: str | None = None
+    # The agent's workspace for downloads and local files.
+    files: FileSystemOptions = FileSystemOptions()
+
+    @field_validator("arguments")
+    @classmethod
+    def _owned_arguments(cls, arguments: tuple[str, ...]) -> tuple[str, ...]:
+        reserved = ("--remote-debugging", "--user-data-dir", "--headless")
+        if any(
+            arg.startswith(reserved) or not arg.startswith("--") for arg in arguments
+        ):
+            raise ValueError(
+                "Arguments must be flags and cannot override process ownership"
+            )
+        return arguments
+
+    def process(self) -> ChromeProcess:
+        """The Chrome process these options describe, not started yet."""
+        flags = [f"--window-size={self.viewport.width},{self.viewport.height}"]
+        if self.headless:
+            flags.append("--headless=new")
+        return ChromeProcess(
+            executable=self.executable,
+            flags=(*flags, *self.arguments),
+            timeout=self.timeout,
+        )
+
+    async def cdp_endpoint(self, stack: AsyncExitStack) -> CdpEndpoint:
+        """Start Chrome and return its endpoint; closing `stack` stops the process."""
+        process = await stack.enter_async_context(self.process())
+        return CdpEndpoint(url=process.endpoint, headers={})
+
+
+class Connect(Options):
+    """Attach to a browser that runs elsewhere; browserrig never stops it."""
+
+    owns_process: ClassVar[bool] = False
+    endpoint: SecretStr
+    target_id: TargetId | None = None
+    url: str = "about:blank"
+    timeout: float = Field(default=20, gt=0, le=120)
+    headers: dict[str, SecretStr] = Field(default_factory=dict, repr=False)
+    viewport: Viewport | None = None
+    session_setup: tuple[InstanceOf[SessionSetup], ...] = ()
+    download_directory: str | None = None
+    # A browser running elsewhere gets a local workspace only when one is configured.
+    files: FileSystemOptions | None = None
+
+    @model_validator(mode="after")
+    def _target_ownership(self) -> Self:
+        if self.target_id and (self.url != "about:blank" or self.download_directory):
+            raise ValueError(
+                "An attached target keeps its URL and browser-wide download policy"
+            )
+        return self
+
+    @field_validator("endpoint")
+    @classmethod
+    def _browser_endpoint(cls, endpoint: SecretStr) -> SecretStr:
+        parsed = urlsplit(endpoint.get_secret_value())
+        if parsed.scheme not in {"http", "https", "ws", "wss"} or not parsed.hostname:
+            raise ValueError(
+                "Expected an HTTP(S) discovery endpoint or browser WebSocket"
+            )
+        if parsed.scheme in {"ws", "wss"} and "/devtools/page/" in parsed.path:
+            raise ValueError(
+                "A browser WebSocket is required; a page WebSocket cannot manage tabs"
+            )
+        return endpoint
+
+    async def cdp_endpoint(self, _stack: AsyncExitStack) -> CdpEndpoint:
+        """Return the configured endpoint and its request headers."""
+        headers = {key: value.get_secret_value() for key, value in self.headers.items()}
+        return CdpEndpoint(url=self.endpoint.get_secret_value(), headers=headers)
+
+
+type BrowserOptions = Launch | Connect

browserrig/browser/downloads.py

@@ -0,0 +1,128 @@
+from time import monotonic
+
+from cdpify import CdpClient
+from cdpify.domains.browser import DownloadProgressEvent, DownloadWillBeginEvent
+
+from browserrig.events import (
+    DownloadCanceledEvent,
+    DownloadCompletedEvent,
+    DownloadFailedEvent,
+    DownloadProgressedEvent,
+    DownloadStartedEvent,
+)
+from browserrig.files import Download, DownloadState, FileSystem
+
+from .base import EventSink
+
+# Chrome reports every received chunk; viewers only need a smooth progress ring.
+_PROGRESS_INTERVAL_SECONDS = 0.2
+
+
+class Downloads:
+    """Track browser downloads and move completed files into the workspace."""
+
+    def __init__(
+        self, *, files: FileSystem | None, directory: str | None, emit: EventSink
+    ) -> None:
+        self._files = files
+        self._directory = directory
+        self._emit = emit
+        self._items: dict[str, Download] = {}
+        self._progress_at: dict[str, float] = {}
+
+    async def enable(self, client: CdpClient, context_id: str | None) -> None:
+        directory = (
+            str(self._files.pending_directory) if self._files else self._directory
+        )
+        if directory is not None:
+            await client.browser.set_download_behavior(
+                behavior="allowAndName",
+                download_path=directory,
+                browser_context_id=context_id,
+                events_enabled=True,
+            )
+
+    def all(self) -> tuple[Download, ...]:
+        return tuple(self._items.values())
+
+    async def handle(
+        self, event: DownloadWillBeginEvent | DownloadProgressEvent
+    ) -> None:
+        match event:
+            case DownloadWillBeginEvent():
+                self._items[event.guid] = Download(
+                    id=event.guid, url=event.url, filename=event.suggested_filename
+                )
+                await self._emit(
+                    DownloadStartedEvent(
+                        download_id=event.guid,
+                        filename=event.suggested_filename,
+                        url=event.url,
+                    )
+                )
+            case DownloadProgressEvent() if event.guid in self._items:
+                await self._progress(event)
+
+    async def _progress(self, event: DownloadProgressEvent) -> None:
+        guid = event.guid
+        previous = self._items[guid]
+        # Terminal CDP events can be repeated; never move an already completed file twice.
+        if previous.state != DownloadState.IN_PROGRESS:
+            return
+        total = int(event.total_bytes) if event.total_bytes > 0 else None
+        if event.state == DownloadState.IN_PROGRESS:
+            await self._report_progress(previous, int(event.received_bytes), total)
+            return
+        path, error = event.file_path, None
+        if event.state == DownloadState.COMPLETED and self._files is not None:
+            try:
+                path = await self._files.complete_download(
+                    guid, previous.filename, path
+                )
+            except (OSError, ValueError) as failure:
+                # A single missing file must not stop browser event processing.
+                path, error = None, f"Downloaded file could not be saved: {failure}"
+        download = Download(
+            id=guid,
+            url=previous.url,
+            filename=previous.filename,
+            state=DownloadState(event.state),
+            received_bytes=int(event.received_bytes),
+            total_bytes=total,
+            path=path,
+            error=error,
+        )
+        self._items[guid] = download
+        self._progress_at.pop(guid, None)
+        match download.state:
+            case DownloadState.COMPLETED if error is not None:
+                await self._emit(DownloadFailedEvent(download_id=guid, reason=error))
+            case DownloadState.COMPLETED:
+                await self._emit(
+                    DownloadCompletedEvent(
+                        download_id=guid,
+                        filename=download.filename,
+                        size=download.received_bytes,
+                    )
+                )
+            case DownloadState.CANCELED:
+                await self._emit(DownloadCanceledEvent(download_id=guid))
+
+    async def _report_progress(
+        self, download: Download, received: int, total: int | None
+    ) -> None:
+        self._items[download.id] = download.model_copy(
+            update={"received_bytes": received, "total_bytes": total}
+        )
+        now = monotonic()
+        if (
+            now - self._progress_at.get(download.id, float("-inf"))
+            < _PROGRESS_INTERVAL_SECONDS
+        ):
+            return
+        self._progress_at[download.id] = now
+        await self._emit(
+            DownloadProgressedEvent(
+                download_id=download.id, received_bytes=received, total_bytes=total
+            )
+        )

browserrig/browser/favicon.py

@@ -0,0 +1,119 @@
+import asyncio
+import base64
+import re
+from contextlib import suppress
+from urllib.parse import unquote_to_bytes, urlsplit
+
+from cdpify import CdpSession
+from cdpify.domains.network import LoadNetworkResourceOptions
+from cdpify.exceptions import CDPError
+
+# Icons travel inside live events, so oversized ones are dropped instead of streamed.
+_MAX_BYTES = 64 * 1024
+_TIMEOUT_SECONDS = 5
+_CLOSE_SECONDS = 1
+_CHUNK_BYTES = 16 * 1024
+
+_SIGNATURES = (
+    (b"\x89PNG\r\n\x1a\n", "image/png"),
+    (b"\x00\x00\x01\x00", "image/x-icon"),
+    (b"GIF87a", "image/gif"),
+    (b"GIF89a", "image/gif"),
+    (b"\xff\xd8\xff", "image/jpeg"),
+    (b"BM", "image/bmp"),
+)
+_SVG_PREAMBLE = re.compile(rb"^(?:\s|<\?xml[^>]*\?>|<!--.*?-->)*", re.S)
+
+
+async def load_favicon(cdp: CdpSession, frame_id: str, url: str) -> str:
+    """Return the page's icon as a base64 `data:` URI, or "" when it is missing or unusable.
+
+    Web icons load through the page's own network stack without cookies, so icons behind
+    a login do not load. Inline `data:` icons are decoded locally. Either way the type
+    comes from the bytes. SVG stays safe only while viewers render it with `<img>`, never
+    through `innerHTML` or `<object>`.
+    """
+    scheme = urlsplit(url).scheme
+    try:
+        if scheme == "data":
+            data = _decode_data_uri(url)
+        elif scheme in {"http", "https"}:
+            async with asyncio.timeout(_TIMEOUT_SECONDS):
+                data = await _fetch(cdp, frame_id, url)
+        else:
+            return ""
+    except (CDPError, TimeoutError, ValueError):
+        return ""
+    if not data or (media_type := _media_type(data)) is None:
+        return ""
+    return f"data:{media_type};base64,{base64.b64encode(data).decode()}"
+
+
+def _decode_data_uri(url: str) -> bytes | None:
+    header, separator, payload = url.removeprefix("data:").partition(",")
+    if not separator:
+        return None
+    if header.lower().endswith(";base64"):
+        encoded = b"".join(unquote_to_bytes(payload).split())
+        data = base64.b64decode(encoded, validate=True)
+    else:
+        data = unquote_to_bytes(payload)
+    return data if len(data) <= _MAX_BYTES else None
+
+
+async def _fetch(cdp: CdpSession, frame_id: str, url: str) -> bytes | None:
+    result = await cdp.network.load_network_resource(
+        frame_id=frame_id,
+        url=url,
+        options=LoadNetworkResourceOptions(
+            disable_cache=False, include_credentials=False
+        ),
+    )
+    resource = result.resource
+    if not resource.success or resource.stream is None:
+        return None
+    try:
+        status = resource.http_status_code
+        if status is not None and not 200 <= status < 300:
+            return None
+        return await _read(cdp, resource.stream)
+    finally:
+        # Closing must neither hang after a timeout nor replace the result with its error.
+        with suppress(CDPError, TimeoutError):
+            async with asyncio.timeout(_CLOSE_SECONDS):
+                await cdp.io.close(handle=resource.stream)
+
+
+async def _read(cdp: CdpSession, handle: str) -> bytes | None:
+    data = bytearray()
+    while len(data) <= _MAX_BYTES:
+        chunk = await cdp.io.read(handle=handle, size=_CHUNK_BYTES)
+        data += (
+            base64.b64decode(chunk.data, validate=True)
+            if chunk.base64_encoded
+            else chunk.data.encode()
+        )
+        if chunk.eof:
+            return bytes(data) if len(data) <= _MAX_BYTES else None
+    return None
+
+
+def _media_type(data: bytes) -> str | None:
+    # The page controls both the icon URL and its headers, so only the bytes decide.
+    for signature, media_type in _SIGNATURES:
+        if data.startswith(signature):
+            return media_type
+    if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
+        return "image/webp"
+    if data[4:12] in (b"ftypavif", b"ftypavis"):
+        return "image/avif"
+    head = data[:2048].lstrip(b"\xef\xbb\xbf").lower()
+    # An XHTML page may start with an XML prolog and embed inline SVG.
+    if b"<html" in head:
+        return None
+    body = _SVG_PREAMBLE.sub(b"", head)
+    if body.startswith(b"<svg") or (
+        body.startswith(b"<!doctype svg") and b"<svg" in body
+    ):
+        return "image/svg+xml"
+    return None