[py] Implement add_request_handler

Author: jkim2492

py/selenium/webdriver/common/bidi/network.py

@@ -3,6 +3,7 @@
 from dataclasses import dataclass, fields, is_dataclass
 
 from selenium.webdriver.common.bidi.cdp import import_devtools
+from selenium.webdriver.common.bidi.session import session_subscribe
 
 from . import script
 
@@ -489,3 +490,27 @@ def cmd(self):
         result = yield self.to_json()
         return result
 
+
+class Network:
+    def __init__(self, conn):
+
+        self.conn = conn
+        self.callbacks = {}
+
+    async def add_intercept(self, event, params: AddInterceptParameters):
+        await self.bidi.execute(session_subscribe(event.event_class))
+        result = await self.bidi.execute(AddIntercept(params).cmd())
+        return result
+
+    async def add_listener(self, event, callback):
+        listener = self.bidi.listen(event)
+
+        async for event in listener:
+            request_data = BeforeRequestSentParameters.from_json(
+                event.to_json()["params"]
+            )
+            await callback(request_data)
+
+    async def continue_request(self, params: ContinueRequestParameters):
+        result = await self.bidi.execute(ContinueRequest(params).cmd())
+        return result

py/selenium/webdriver/remote/network.py

@@ -0,0 +1,41 @@
+from dataclasses import fields
+
+from selenium.webdriver.common.bidi import network
+from selenium.webdriver.common.bidi.cdp import open_cdp
+from selenium.webdriver.common.bidi.network import (
+    AddInterceptParameters,
+    BeforeRequestSent,
+    BeforeRequestSentParameters,
+    ContinueRequestParameters,
+)
+
+
+def default_request_handler(params: BeforeRequestSentParameters):
+    return ContinueRequestParameters(request=params.request["request"])
+
+
+class Network:
+    def __init__(self, ws_url):
+        self.ws_url = ws_url
+        self.network = None
+
+    async def add_request_handler(self, request_filter=lambda _: True, handler=default_request_handler):
+        async with open_cdp(self.ws_url) as conn:
+            if not self.network:
+                self.network = network.Network(conn)
+            params = AddInterceptParameters(["beforeRequestSent"])
+            result = await self.network.add_intercept(event=BeforeRequestSent, params=params)
+            intercept = result["intercept"]
+            callback = self._callback(request_filter, handler)
+            await self.network.add_listener(BeforeRequestSent, callback)
+            return intercept
+
+    def _callback(self, request_filter, handler):
+        async def callback(request):
+            if request_filter(request):
+                request = handler(request)
+            else:
+                request = default_request_handler(request)
+            await self.network.continue_request(request)
+
+        return callback

py/selenium/webdriver/remote/webdriver.py

@@ -51,6 +51,7 @@
 from selenium.webdriver.common.virtual_authenticator import (
     required_virtual_authenticator,
 )
+from selenium.webdriver.remote.network import Network
 from selenium.webdriver.support.relative_locator import RelativeBy
 
 from .bidi_connection import BidiConnection
@@ -213,6 +214,7 @@ def __init__(
 
         self._websocket_connection = None
         self._script = None
+        self._network = None
 
     def __repr__(self):
         return f'<{type(self).__module__}.{type(self).__name__} (session="{self.session_id}")>'
@@ -1080,6 +1082,13 @@ def script(self):
 
         return self._script
 
+    @property
+    def network(self):
+        if not self._network:
+            self._network = Network(self.caps.get("webSocketUrl"))
+
+        return self._network
+
     def _start_bidi(self):
         if self.caps.get("webSocketUrl"):
             ws_url = self.caps.get("webSocketUrl")