"""BiDi session/API driver
This directly talks the https://w3c.github.io/webdriver-bidi/ protocol from async Python,
without any node/JS in between. The only dependencies are aiohttp, firefox, and/or
chromedriver+chromium (or headless_shell).
"""
import asyncio
import collections
import contextlib
import json
import logging
import os
import socket
import sys
import tempfile
from collections.abc import AsyncIterator, Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Any, AsyncIterable, Iterable, Self
import aiohttp
log_proto = logging.getLogger("bidi.proto")
log_command = logging.getLogger("bidi.command")
JsonObject = dict[str, Any]
class WebdriverError(RuntimeError):
pass
@dataclass
class ConsoleMessage:
level: str
type: str
timestamp: int
args: list[object]
text: str
def __init__(self, message_params: JsonObject) -> None:
self.level = message_params["level"]
self.type = message_params["type"]
self.timestamp = message_params["timestamp"]
self.args = message_params.get("args", [])
self.text = message_params["text"]
def __str__(self) -> str:
return f"> {self.level}: {self.text}"
class ConsoleLog(Iterable[ConsoleMessage], AsyncIterable[ConsoleMessage]):
def __init__(self) -> None:
self.ready = asyncio.Event()
self.log = collections.deque[ConsoleMessage]()
self.eof = False
def append(self, message: ConsoleMessage) -> None:
self.log.append(message)
self.ready.set()
def close(self) -> None:
self.eof = True
self.ready.set()
def __iter__(self) -> Iterator[ConsoleMessage]:
"""Iterates over all items currently in the log. If .eof is false,
then more items might still be added."""
return iter(self.log)
async def __aiter__(self) -> AsyncIterator[ConsoleMessage]:
"""Iterates over all items in the log, waiting for the next item to
come, and only stopping once the browser session is closed. Consumes
the log entries."""
while not self.eof:
if len(self.log):
yield self.log.popleft()
else:
self.ready.clear()
await self.ready.wait()
@dataclass
class BidiSession:
ws_url: str
session_url: str
process: asyncio.subprocess.Process
def pick_ports(count: int) -> list[int]:
sockets: list[socket.socket] = []
ports: list[int] = []
try:
for _ in range(count):
sock = socket.socket()
sock.bind(('127.0.0.1', 0))
sockets.append(sock)
ports.append(sock.getsockname()[1])
finally:
for s in sockets:
s.close()
return ports
def jsquote(js: object) -> str:
return json.dumps(js)
def unpack_value(raw: Any) -> Any:
"""Convert a WebDriver type/value object into a plain Python object"""
if not isinstance(raw, dict):
return raw
try:
type_ = raw["type"]
except KeyError as e:
raise ValueError(f"No type in {raw!r}") from e
if type_ in ["undefined", "null"]:
return None
if type_ in ["string", "boolean", "number"]:
return raw["value"]
if type_ in ["function", "node", "window"]:
return raw
if type_ == "array":
return [unpack_value(v) for v in raw["value"]]
if type_ == "object":
obj = {}
for k, v in raw["value"]:
obj[k] = unpack_value(v)
return obj
raise ValueError(f"Unknown type in {raw!r}")
class WebdriverBidi:
http_session: aiohttp.ClientSession
def __init__(self, *, headless: bool = False) -> None:
self.headless = headless
self.last_id = 0
self.pending_commands: dict[int, asyncio.Future[JsonObject]] = {}
self.logs = ConsoleLog()
self.bidi_session: BidiSession | None = None
self.future_wait_page_load = None
self.top_context: str | None = None
self.context: str | None
self.homedir_temp = tempfile.TemporaryDirectory(prefix="cockpit-test-browser-home-")
self.homedir = Path(self.homedir_temp.name)
self.download_dir = self.homedir / 'Downloads'
self.download_dir.mkdir()
async def start_bidi_session(self) -> None:
raise NotImplementedError
async def close_bidi_session(self) -> None:
raise NotImplementedError
async def close(self) -> None:
assert self.bidi_session is not None
log_proto.debug("cleaning up webdriver")
await self.ws.close()
await self.close_bidi_session()
try:
self.bidi_session.process.terminate()
except ProcessLookupError:
pass
await self.bidi_session.process.wait()
self.bidi_session = None
await self.http_session.close()
await self.task_reader
def ws_done_callback(self, _task: asyncio.Task[None]) -> None:
for fut in self.pending_commands.values():
if not fut.done():
fut.set_exception(WebdriverError("websocket closed"))
self.logs.close()
async def start_session(self) -> None:
self.http_session = aiohttp.ClientSession(raise_for_status=True)
await self.start_bidi_session()
assert self.bidi_session
self.ws = await self.http_session.ws_connect(self.bidi_session.ws_url)
self.task_reader = asyncio.create_task(self.ws_reader(self.ws), name="bidi_reader")
self.task_reader.add_done_callback(self.ws_done_callback)
await self.bidi("session.subscribe", events=[
"log.entryAdded", "browsingContext.domContentLoaded",
])
for _ in range(10):
realms = (await self.bidi("script.getRealms"))["realms"]
if len(realms) > 0:
self.top_context = realms[0]["context"]
self.context = self.top_context
break
await asyncio.sleep(0.5)
else:
raise WebdriverError("timed out waiting for default realm")
async def __aenter__(self) -> Self:
await self.start_session()
return self
async def __aexit__(self, *_excinfo: Any) -> None:
if self.bidi_session is not None:
await self.close()
async def ws_reader(self, ws: aiohttp.client.ClientWebSocketResponse) -> None:
async for msg in ws:
if msg.type == aiohttp.WSMsgType.TEXT:
data = json.loads(msg.data)
log_proto.debug("ws TEXT → %r", data)
id_ = data.get("id")
if id_ in self.pending_commands:
log_proto.debug("ws_reader: resolving pending command %i", id_)
pending_command = self.pending_commands[id_]
if data["type"] == "success":
if data["result"].get("type") == "exception":
pending_command.set_exception(
WebdriverError(data["result"]["exceptionDetails"]["text"]))
else:
pending_command.set_result(data["result"])
elif data["type"] == "exception":
pending_command.set_exception(
WebdriverError(data["exceptionDetails"]["text"]))
elif data["type"] == "error":
pending_command.set_exception(
WebdriverError(f"{data['error']}: {data.get('message', '')}"))
else:
pending_command.set_exception(
WebdriverError(f"unknown response type: {data!r}"))
del self.pending_commands[id_]
continue
if data["type"] == "event":
if data["method"] == "log.entryAdded":
message = ConsoleMessage(data["params"])
self.logs.append(message)
print(str(message), file=sys.stderr)
continue
if data["method"] == "browsingContext.domContentLoaded":
if self.future_wait_page_load:
log_command.debug("page loaded: %r, resolving wait page load future", data["params"])
self.future_wait_page_load.set_result(data["params"]["url"])
else:
log_command.debug("page loaded: %r (not awaited)", data["params"])
continue
log_proto.warning("ws_reader: unhandled message %r", data)
elif msg.type == aiohttp.WSMsgType.ERROR:
log_proto.error("BiDi failure: %s", msg)
break
async def bidi(self, method: str, *, quiet: bool = False, timeout: int = 10, **params: Any) -> JsonObject:
"""Send a Webdriver BiDi command and return the JSON response
Most commands ought to be quick, so the default timeout is 10 seconds.
Set a custom one if you call a function which waits/polls for something
for a non-trivial duration.
"""
tag = self.last_id
self.last_id += 1
payload = json.dumps({"id": tag, "method": method, "params": params})
log_command.info("← %s(%r) [id %i]", method, 'quiet' if quiet else params, tag)
await self.ws.send_str(payload)
future = asyncio.get_event_loop().create_future()
self.pending_commands[tag] = future
res = await asyncio.wait_for(future, timeout=timeout)
if "result" in res:
log_proto.debug("[id %i] unpacking raw result %r", tag, res["result"])
value = unpack_value(res["result"])
res["result"] = value
if not quiet:
log_command.info("→ %r", res)
return res
async def locate(self, selector: str) -> str:
r = await self.bidi("browsingContext.locateNodes", context=self.context,
locator={"type": "css", "value": selector})
nodes = r["nodes"]
if len(nodes) == 0:
raise WebdriverError(f"no element found for {selector}")
if len(nodes) > 1:
raise WebdriverError(f"selector {selector} is ambiguous: {nodes}")
log_command.info("locate(%s) = %r", selector, nodes[0])
return nodes[0]
async def switch_to_frame(self, name: str) -> None:
self.switch_to_top()
frame = await self.locate(f"iframe[name='{name}']")
cw = await self.bidi("script.callFunction",
functionDeclaration="f => f.contentWindow",
arguments=[frame],
awaitPromise=False,
target={"context": self.top_context})
self.context = cw["result"]["value"]["context"]
log_command.info("← switch_to_frame(%s)", name)
def switch_to_top(self) -> None:
self.context = self.top_context
log_command.info("← switch_to_top")
def in_top_context(self) -> bool:
return self.context == self.top_context
@contextlib.contextmanager
def restore_context(self) -> Iterator[None]:
saved = self.context
try:
yield
finally:
self.context = saved
class ChromiumBidi(WebdriverBidi):
def __init__(self, *, headless: bool = False) -> None:
super().__init__(headless=headless)
self.cdp_ws: aiohttp.client.ClientWebSocketResponse | None = None
async def start_bidi_session(self) -> None:
assert self.bidi_session is None
candidate_binaries = ["/usr/bin/chromium-browser", "/usr/bin/chromium"]
if self.headless:
candidate_binaries.append("/usr/lib64/chromium-browser/headless_shell")
binaries = [path for path in candidate_binaries if os.path.exists(path)]
if not binaries:
raise WebdriverError(f"no Chromium binary found: tried {' '.join(candidate_binaries)}")
chrome_binary = binaries[0]
log_proto.debug("webdriver chromium binary: %s", chrome_binary)
session_args = {"capabilities": {
"alwaysMatch": {
"webSocketUrl": True,
"acceptInsecureCerts": True,
"goog:chromeOptions": {
"binary": chrome_binary,
"args": [
"--font-render-hinting=none",
"--disable-pushstate-throttle",
"--disable-popup-blocking",
] + (["--headless"] if self.headless else []),
},
}
}}
[webdriver_port] = pick_ports(1)
driver = await asyncio.create_subprocess_exec(
"chromedriver", "--port=" + str(webdriver_port),
env=dict(os.environ, HOME=str(self.homedir)))
wd_url = f"http://localhost:{webdriver_port}"
for retry in range(1, 10):
try:
async with self.http_session.post(f"{wd_url}/session",
data=json.dumps(session_args).encode()) as resp:
session_info = json.loads(await resp.text())["value"]
log_proto.debug("webdriver session request: %r %r", resp, session_info)
break
except (IOError, aiohttp.client.ClientResponseError) as e:
log_proto.debug("waiting for webdriver: %s", e)
await asyncio.sleep(0.1 * retry)
else:
raise WebdriverError("could not connect to chromedriver")
self.cdp_address = session_info["capabilities"]["goog:chromeOptions"]["debuggerAddress"]
self.last_cdp_id = 0
self.bidi_session = BidiSession(
session_url=f"{wd_url}/session/{session_info['sessionId']}",
ws_url=session_info["capabilities"]["webSocketUrl"],
process=driver)
log_proto.debug("Established chromium session %r, CDP address %s", self.bidi_session, self.cdp_address)
async def close_cdp_session(self) -> None:
if self.cdp_ws is not None:
await self.cdp_ws.close()
self.cdp_ws = None
async def close_bidi_session(self) -> None:
assert self.bidi_session is not None
await self.close_cdp_session()
await self.http_session.delete(self.bidi_session.session_url)
async def cdp(self, method: str, **params: Any) -> dict[str, Any]:
"""Send a Chrome DevTools command and return the JSON response
This is currently *not* safe for enabling events! These should be handled via BiDi,
this is only an escape hatch for CDP specific functionality such as Profiler.
"""
if self.cdp_ws is None:
self.cdp_ws = await self.http_session.ws_connect(
f"ws://{self.cdp_address}/devtools/page/{self.top_context}",
max_msg_size=16000000)
reply = None
payload = json.dumps({"id": self.last_cdp_id, "method": method, "params": params})
log_command.info("CDP ← %r", payload)
await self.cdp_ws.send_str(payload)
async for msg in self.cdp_ws:
if msg.type == aiohttp.WSMsgType.TEXT:
reply = json.loads(msg.data)
if reply.get("id") == self.last_cdp_id:
break
else:
log_proto.warning("CDP message: %r", reply)
else:
log_proto.debug("CDP non-text message: %r", msg)
log_command.info("CDP → %r", reply)
assert reply
self.last_cdp_id += 1
return reply
class FirefoxBidi(WebdriverBidi):
async def start_bidi_session(self) -> None:
[marionette_port, bidi_port] = pick_ports(2)
self.profiledir = self.homedir / "profile"
self.profiledir.mkdir()
(self.profiledir / "user.js").write_text(f"""
// https://bugzilla.mozilla.org/show_bug.cgi?id=1947402
user_pref('remote.events.async.enabled', false);
// set this to "Trace" for debugging BiDi interactions
user_pref('remote.log.level', 'Warn');
user_pref('remote.log.truncate', false);
// enable remote logs on stdout
user_pref('browser.dom.window.dump.enabled', true);
user_pref("app.update.auto", false);
user_pref("datareporting.policy.dataSubmissionEnabled", false);
user_pref("toolkit.telemetry.reportingpolicy.firstRun", false);
user_pref("dom.disable_beforeunload", true);
user_pref("browser.download.dir", "{self.download_dir}");
user_pref("browser.download.folderList", 2);
user_pref("signon.rememberSignons", false);
user_pref("dom.navigation.locationChangeRateLimit.count", 9999);
user_pref('marionette.port', {marionette_port});
""")
driver = await asyncio.create_subprocess_exec(
"firefox", "-profile", str(self.profiledir), "--marionette", "--no-remote",
f"--remote-debugging-port={bidi_port}",
*(["-headless"] if self.headless else []), "about:blank")
for _ in range(1, 30):
try:
reader, self.writer_marionette = await asyncio.open_connection("127.0.0.1", marionette_port)
break
except ConnectionRefusedError as e:
log_proto.debug("waiting for firefox marionette: %s", e)
await asyncio.sleep(1)
else:
raise WebdriverError("could not connect to firefox marionette")
reply = await reader.read(1024)
if b'"marionetteProtocol":3' not in reply:
raise WebdriverError(f"unexpected marionette reply: {reply.decode()}")
cmd = '[0,1,"WebDriver:NewSession",{"webSocketUrl":true,"acceptInsecureCerts":true}]'
self.writer_marionette.write(f"{len(cmd)}:{cmd}".encode())
await self.writer_marionette.drain()
reply = await reader.read(1024)
reply = json.loads(reply[reply.index(b":") + 1:].decode())
if not isinstance(reply, list) or len(reply) != 4 or not isinstance(reply[3], dict):
raise WebdriverError(f"unexpected marionette session request reply: {reply!r}")
log_proto.debug("marionette session request reply: %s", reply)
url = reply[3]["capabilities"]["webSocketUrl"]
self.bidi_session = BidiSession(session_url=url, ws_url=url, process=driver)
log_proto.debug("Established firefox session %r", self.bidi_session)
async def close_bidi_session(self) -> None:
self.writer_marionette.close()
await self.writer_marionette.wait_closed()
