cdp.py

# Licensed to the Software Freedom Conservancy (SFC) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The SFC licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.


import contextvars
import importlib
import itertools
import json
import logging
import os
import pathlib
from collections import defaultdict
from collections.abc import AsyncGenerator, AsyncIterator, Generator
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass
from typing import Any, TypeVar

import trio
from trio_websocket import ConnectionClosed as WsConnectionClosed
from trio_websocket import connect_websocket_url

logger = logging.getLogger("trio_cdp")
T = TypeVar("T")
MAX_WS_MESSAGE_SIZE = 2**24


def _resolve_max_message_size(explicit=None):
    """Return the WebSocket max message size to use.

    Priority: explicit argument > ``SE_CDP_MAX_WS_MESSAGE_SIZE`` env var > ``MAX_WS_MESSAGE_SIZE``.
    """
    if explicit is not None:
        return explicit
    return int(os.environ.get("SE_CDP_MAX_WS_MESSAGE_SIZE", MAX_WS_MESSAGE_SIZE))


devtools = None
version = None


def import_devtools(ver):
    """Attempt to load the current latest available devtools into the module cache for use later."""
    global devtools
    global version
    version = ver
    base = "selenium.webdriver.common.devtools.v"
    try:
        devtools = importlib.import_module(f"{base}{ver}")
        return devtools
    except ModuleNotFoundError:
        # Attempt to parse and load the 'most recent' devtools module. This is likely
        # because cdp has been updated but selenium python has not been released yet.
        devtools_path = pathlib.Path(__file__).parents[1].joinpath("devtools")
        versions = tuple(f.name for f in devtools_path.iterdir() if f.is_dir())
        available_versions = tuple(x for x in versions if x == "latest" or (x.startswith("v") and x[1:].isdigit()))
        numeric_versions = tuple(x[1:] for x in available_versions if x.startswith("v"))
        if not numeric_versions:
            raise
        latest = max(numeric_versions, key=int)
        selenium_logger = logging.getLogger(__name__)
        selenium_logger.debug("Falling back to loading `devtools`: v%s", latest)
        devtools = importlib.import_module(f"{base}{latest}")
        return devtools


_connection_context: contextvars.ContextVar = contextvars.ContextVar("connection_context")
_session_context: contextvars.ContextVar = contextvars.ContextVar("session_context")


def get_connection_context(fn_name):
    """Look up the current connection.

    If there is no current connection, raise a ``RuntimeError`` with a
    helpful message.
    """
    try:
        return _connection_context.get()
    except LookupError:
        raise RuntimeError(f"{fn_name}() must be called in a connection context.")


def get_session_context(fn_name):
    """Look up the current session.

    If there is no current session, raise a ``RuntimeError`` with a
    helpful message.
    """
    try:
        return _session_context.get()
    except LookupError:
        raise RuntimeError(f"{fn_name}() must be called in a session context.")


@contextmanager
def connection_context(connection):
    """Context manager installs ``connection`` as the session context for the current Trio task."""
    token = _connection_context.set(connection)
    try:
        yield
    finally:
        _connection_context.reset(token)


@contextmanager
def session_context(session):
    """Context manager installs ``session`` as the session context for the current Trio task."""
    token = _session_context.set(session)
    try:
        yield
    finally:
        _session_context.reset(token)


def set_global_connection(connection):
    """Install ``connection`` in the root context so that it will become the default connection for all tasks.

    This is generally not recommended, except it may be necessary in
    certain use cases such as running inside Jupyter notebook.
    """
    global _connection_context
    _connection_context = contextvars.ContextVar("_connection_context", default=connection)


def set_global_session(session):
    """Install ``session`` in the root context so that it will become the default session for all tasks.

    This is generally not recommended, except it may be necessary in
    certain use cases such as running inside Jupyter notebook.
    """
    global _session_context
    _session_context = contextvars.ContextVar("_session_context", default=session)


class BrowserError(Exception):
    """This exception is raised when the browser's response to a command indicates that an error occurred."""

    def __init__(self, obj):
        self.code = obj.get("code")
        self.message = obj.get("message")
        self.detail = obj.get("data")

    def __str__(self):
        return f"BrowserError<code={self.code} message={self.message}> {self.detail}"


class CdpConnectionClosed(WsConnectionClosed):
    """Raised when a public method is called on a closed CDP connection."""

    def __init__(self, reason):
        """Constructor.

        Args:
            reason: wsproto.frame_protocol.CloseReason
        """
        self.reason = reason

    def __repr__(self):
        """Return representation."""
        return f"{self.__class__.__name__}<{self.reason}>"


class InternalError(Exception):
    """This exception is only raised when there is faulty logic in TrioCDP or the integration with PyCDP."""

    pass


@dataclass
class CmEventProxy:
    """A proxy object returned by :meth:`CdpBase.wait_for()``.

    After the context manager executes, this proxy object will have a
    value set that contains the returned event.
    """

    value: Any = None


class CdpBase:
    def __init__(self, ws, session_id, target_id):
        self.ws = ws
        self.session_id = session_id
        self.target_id = target_id
        self.channels = defaultdict(set)
        self.id_iter = itertools.count()
        self.inflight_cmd = {}
        self.inflight_result = {}

    async def execute(self, cmd: Generator[dict, T, Any]) -> T:
        """Execute a command on the server and wait for the result.

        Args:
            cmd: any CDP command

        Returns:
            a CDP result
        """
        cmd_id = next(self.id_iter)
        cmd_event = trio.Event()
        self.inflight_cmd[cmd_id] = cmd, cmd_event
        request = next(cmd)
        request["id"] = cmd_id
        if self.session_id:
            request["sessionId"] = self.session_id
        request_str = json.dumps(request)
        if logger.isEnabledFor(logging.DEBUG):
            logger.debug(f"Sending CDP message: {cmd_id} {cmd_event}: {request_str}")
        try:
            await self.ws.send_message(request_str)
        except WsConnectionClosed as wcc:
            raise CdpConnectionClosed(wcc.reason) from None
        await cmd_event.wait()
        response = self.inflight_result.pop(cmd_id)
        if logger.isEnabledFor(logging.DEBUG):
            logger.debug(f"Received CDP message: {response}")
        if isinstance(response, Exception):
            if logger.isEnabledFor(logging.DEBUG):
                logger.debug(f"Exception raised by {cmd_event} message: {type(response).__name__}")
            raise response
        return response

    def listen(self, *event_types, buffer_size=10):
        """Listen for events.

        Returns:
            An async iterator that iterates over events matching the indicated types.
        """
        sender, receiver = trio.open_memory_channel(buffer_size)
        for event_type in event_types:
            self.channels[event_type].add(sender)
        return receiver

    @asynccontextmanager
    async def wait_for(self, event_type: type[T], buffer_size=10) -> AsyncGenerator[CmEventProxy, None]:
        """Wait for an event of the given type and return it.

        This is an async context manager, so you should open it inside
        an async with block. The block will not exit until the indicated
        event is received.
        """
        sender: trio.MemorySendChannel
        receiver: trio.MemoryReceiveChannel
        sender, receiver = trio.open_memory_channel(buffer_size)
        self.channels[event_type].add(sender)
        proxy = CmEventProxy()
        yield proxy
        async with receiver:
            event = await receiver.receive()
        proxy.value = event

    def _handle_data(self, data):
        """Handle incoming WebSocket data.

        Args:
            data: a JSON dictionary
        """
        if "id" in data:
            self._handle_cmd_response(data)
        else:
            self._handle_event(data)

    def _handle_cmd_response(self, data: dict):
        """Handle a response to a command.

        This will set an event flag that will return control to the
        task that called the command.

        Args:
            data: response as a JSON dictionary
        """
        cmd_id = data["id"]
        try:
            cmd, event = self.inflight_cmd.pop(cmd_id)
        except KeyError:
            logger.warning("Got a message with a command ID that does not exist: %s", data)
            return
        if "error" in data:
            # If the server reported an error, convert it to an exception and do
            # not process the response any further.
            self.inflight_result[cmd_id] = BrowserError(data["error"])
        else:
            # Otherwise, continue the generator to parse the JSON result
            # into a CDP object.
            try:
                _ = cmd.send(data["result"])
                raise InternalError("The command's generator function did not exit when expected!")
            except StopIteration as exit:
                return_ = exit.value
            self.inflight_result[cmd_id] = return_
        event.set()

    def _handle_event(self, data: dict):
        """Handle an event.

        Args:
            data: event as a JSON dictionary
        """
        global devtools
        if devtools is None:
            raise RuntimeError("CDP devtools module not loaded. Call import_devtools() first.")
        event = devtools.util.parse_json_event(data)
        logger.debug("Received event: %s", event)
        to_remove = set()
        for sender in self.channels[type(event)]:
            try:
                sender.send_nowait(event)
            except trio.WouldBlock:
                logger.error('Unable to send event "%r" due to full channel %s', event, sender)
            except trio.BrokenResourceError:
                to_remove.add(sender)
        if to_remove:
            self.channels[type(event)] -= to_remove


class CdpSession(CdpBase):
    """Contains the state for a CDP session.

    Generally you should not instantiate this object yourself; you should call
    :meth:`CdpConnection.open_session`.
    """

    def __init__(self, ws, session_id, target_id):
        """Constructor.

        Args:
            ws: trio_websocket.WebSocketConnection
            session_id: devtools.target.SessionID
            target_id: devtools.target.TargetID
        """
        super().__init__(ws, session_id, target_id)

        self._dom_enable_count = 0
        self._dom_enable_lock = trio.Lock()
        self._page_enable_count = 0
        self._page_enable_lock = trio.Lock()

    @asynccontextmanager
    async def dom_enable(self):
        """Context manager that executes ``dom.enable()`` when it enters and then calls ``dom.disable()``.

        This keeps track of concurrent callers and only disables DOM
        events when all callers have exited.
        """
        global devtools
        async with self._dom_enable_lock:
            self._dom_enable_count += 1
            if self._dom_enable_count == 1:
                await self.execute(devtools.dom.enable())

        yield

        async with self._dom_enable_lock:
            self._dom_enable_count -= 1
            if self._dom_enable_count == 0:
                await self.execute(devtools.dom.disable())

    @asynccontextmanager
    async def page_enable(self):
        """Context manager executes ``page.enable()`` when it enters and then calls ``page.disable()`` when it exits.

        This keeps track of concurrent callers and only disables page
        events when all callers have exited.
        """
        global devtools
        async with self._page_enable_lock:
            self._page_enable_count += 1
            if self._page_enable_count == 1:
                await self.execute(devtools.page.enable())

        yield

        async with self._page_enable_lock:
            self._page_enable_count -= 1
            if self._page_enable_count == 0:
                await self.execute(devtools.page.disable())


class CdpConnection(CdpBase, trio.abc.AsyncResource):
    """Contains the connection state for a Chrome DevTools Protocol server.

    CDP can multiplex multiple "sessions" over a single connection. This
    class corresponds to the "root" session, i.e. the implicitly created
    session that has no session ID. This class is responsible for
    reading incoming WebSocket messages and forwarding them to the
    corresponding session, as well as handling messages targeted at the
    root session itself. You should generally call the
    :func:`open_cdp()` instead of instantiating this class directly.
    """

    def __init__(self, ws):
        """Constructor.

        Args:
            ws: trio_websocket.WebSocketConnection
        """
        super().__init__(ws, session_id=None, target_id=None)
        self.sessions = {}

    async def aclose(self):
        """Close the underlying WebSocket connection.

        This will cause the reader task to gracefully exit when it tries
        to read the next message from the WebSocket. All of the public
        APIs (``execute()``, ``listen()``, etc.) will raise
        ``CdpConnectionClosed`` after the CDP connection is closed. It
        is safe to call this multiple times.
        """
        await self.ws.aclose()

    @asynccontextmanager
    async def open_session(self, target_id) -> AsyncIterator[CdpSession]:
        """Context manager opens a session and enables the "simple" style of calling CDP APIs.

        For example, inside a session context, you can call ``await
        dom.get_document()`` and it will execute on the current session
        automatically.
        """
        session = await self.connect_session(target_id)
        with session_context(session):
            yield session

    async def connect_session(self, target_id) -> "CdpSession":
        """Returns a new :class:`CdpSession` connected to the specified target."""
        global devtools
        if devtools is None:
            raise RuntimeError("CDP devtools module not loaded. Call import_devtools() first.")
        session_id = await self.execute(devtools.target.attach_to_target(target_id, True))
        session = CdpSession(self.ws, session_id, target_id)
        self.sessions[session_id] = session
        return session

    async def _reader_task(self):
        """Runs in the background and handles incoming messages.

        Dispatches responses to commands and events to listeners.
        """
        global devtools
        if devtools is None:
            raise RuntimeError("CDP devtools module not loaded. Call import_devtools() first.")
        while True:
            try:
                message = await self.ws.get_message()
            except WsConnectionClosed:
                # If the WebSocket is closed, we don't want to throw an
                # exception from the reader task. Instead we will throw
                # exceptions from the public API methods, and we can quietly
                # exit the reader task here.
                break
            try:
                data = json.loads(message)
            except json.JSONDecodeError:
                raise BrowserError(
                    {
                        "code": -32700,
                        "message": "Client received invalid JSON",
                        "data": message,
                    }
                )
            logger.debug("Received message %r", data)
            if "sessionId" in data:
                session_id = devtools.target.SessionID(data["sessionId"])
                try:
                    session = self.sessions[session_id]
                except KeyError:
                    raise BrowserError(
                        {
                            "code": -32700,
                            "message": "Browser sent a message for an invalid session",
                            "data": f"{session_id!r}",
                        }
                    )
                session._handle_data(data)
            else:
                self._handle_data(data)

        for _, session in self.sessions.items():
            for _, senders in session.channels.items():
                for sender in senders:
                    sender.close()


@asynccontextmanager
async def open_cdp(url, max_message_size=None) -> AsyncIterator[CdpConnection]:
    """Async context manager opens a connection to the browser then closes the connection when the block exits.

    The context manager also sets the connection as the default
    connection for the current task, so that commands like ``await
    target.get_targets()`` will run on this connection automatically. If
    you want to use multiple connections concurrently, it is recommended
    to open each on in a separate task.

    Args:
        url: WebSocket URL of the browser's CDP endpoint.
        max_message_size: Maximum WebSocket message size in bytes. Defaults to the
            ``SE_CDP_MAX_WS_MESSAGE_SIZE`` environment variable, or 16 MiB if unset.
    """
    async with trio.open_nursery() as nursery:
        conn = await connect_cdp(nursery, url, max_message_size=max_message_size)
        try:
            with connection_context(conn):
                yield conn
        finally:
            await conn.aclose()


async def connect_cdp(nursery, url, max_message_size=None) -> CdpConnection:
    """Connect to the browser specified by ``url`` and spawn a background task in the specified nursery.

    The ``open_cdp()`` context manager is preferred in most situations.
    You should only use this function if you need to specify a custom
    nursery. This connection is not automatically closed! You can either
    use the connection object as a context manager (``async with
    conn:``) or else call ``await conn.aclose()`` on it when you are
    done with it. If ``set_context`` is True, then the returned
    connection will be installed as the default connection for the
    current task. This argument is for unusual use cases, such as
    running inside of a notebook.

    Args:
        nursery: Trio nursery to spawn the reader task in.
        url: WebSocket URL of the browser's CDP endpoint.
        max_message_size: Maximum WebSocket message size in bytes. Defaults to the
            ``SE_CDP_MAX_WS_MESSAGE_SIZE`` environment variable, or 16 MiB if unset.
    """
    ws = await connect_websocket_url(nursery, url, max_message_size=_resolve_max_message_size(max_message_size))
    cdp_conn = CdpConnection(ws)
    nursery.start_soon(cdp_conn._reader_task)
    return cdp_conn