[py] Implement add/remove_request_handler

Author: jman-sketch

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

@@ -19,8 +19,14 @@
 import typing
 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, session_unsubscribe
+
 from . import script
 
+devtools = import_devtools("")
+event_class = devtools.util.event_class
+
 
 @dataclass
 class StringValue:
@@ -500,3 +506,32 @@ def from_json(cls, json):
     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.conn.execute(session_subscribe(event.event_class))
+        result = await self.conn.execute(AddIntercept(params).cmd())
+        return result
+
+    async def add_listener(self, event, callback):
+        listener = self.conn.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.conn.execute(ContinueRequest(params).cmd())
+        return result
+
+    async def remove_intercept(self, event, params: RemoveInterceptParameters):
+        await self.conn.execute(session_unsubscribe(event.event_class))
+        await self.conn.execute(RemoveIntercept(params).cmd())

py/selenium/webdriver/remote/network.py

@@ -0,0 +1,76 @@
+# 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.
+from dataclasses import fields
+
+import trio
+
+from selenium.webdriver.common.bidi import network
+from selenium.webdriver.common.bidi.browsing_context import Navigate
+from selenium.webdriver.common.bidi.browsing_context import NavigateParameters
+from selenium.webdriver.common.bidi.network import AddInterceptParameters
+from selenium.webdriver.common.bidi.network import BeforeRequestSent
+from selenium.webdriver.common.bidi.network import BeforeRequestSentParameters
+from selenium.webdriver.common.bidi.network import ContinueRequestParameters
+
+
+def default_request_handler(params: BeforeRequestSentParameters):
+    return ContinueRequestParameters(request=params.request["request"])
+
+
+class Network:
+    def __init__(self, driver):
+        self.network = None
+        self.driver = driver
+        self.intercept = None
+        self.scope = None
+
+    async def add_request_handler(
+        self, request_filter=lambda _: True, handler=default_request_handler, conn=None
+    ):
+        with trio.CancelScope() as scope:
+            self.scope = scope
+            self.network = network.Network(conn)
+            params = AddInterceptParameters(["beforeRequestSent"])
+            callback = self._callback(request_filter, handler)
+            result = await self.network.add_intercept(
+                event=BeforeRequestSent, params=params
+            )
+            intercept = result["intercept"]
+            self.intercept = intercept
+            await self.network.add_listener(event=BeforeRequestSent, callback=callback)
+            return intercept
+
+    async def get(self, url, conn):
+        params = NavigateParameters(context=self.driver.current_window_handle, url=url)
+        await conn.execute(Navigate(params).cmd())
+
+    async def remove_request_handler(self):
+        await self.network.remove_intercept(
+            event=BeforeRequestSent,
+            params=network.RemoveInterceptParameters(self.intercept),
+        )
+        self.scope.cancel()
+
+    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

@@ -59,6 +59,7 @@
 from .file_detector import FileDetector
 from .file_detector import LocalFileDetector
 from .mobile import Mobile
+from .network import Network
 from .remote_connection import RemoteConnection
 from .script_key import ScriptKey
 from .shadowroot import ShadowRoot
@@ -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)
+
+        return self._network
+
     def _start_bidi(self):
         if self.caps.get("webSocketUrl"):
             ws_url = self.caps.get("webSocketUrl")