Feat/intercept (#157)

nathanfallet · stephanlensky · web-flow · commit b6f879c7ecb4 · 2025-07-17T20:10:54.000-04:00
* feat: intercept

* fix fetch domain + changelogs

* using example file for intercept test

---------

Co-authored-by: Stephan Lensky &lt;8302875+stephanlensky@users.noreply.github.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added `Tab.intercept` @nathanfallet
+
 ### Changed
 
 ### Removed
diff --git a/tests/core/test_tab.py b/tests/core/test_tab.py
@@ -4,6 +4,8 @@
 
 import zendriver as zd
 from tests.sample_data import sample_file
+from zendriver.cdp.fetch import RequestStage
+from zendriver.cdp.network import ResourceType
 
 
 async def test_set_user_agent_sets_navigator_values(browser: zd.Browser):
@@ -219,3 +221,20 @@ async def test_expect_download(browser: zd.Browser):
         download = await asyncio.wait_for(download_ex.value, timeout=3)
         assert type(download) is zd.cdp.browser.DownloadWillBegin
         assert download.url is not None
+
+
+async def test_intercept(browser: zd.Browser):
+    tab = browser.main_tab
+
+    async with tab.intercept(
+        "*/user-data.json",
+        RequestStage.RESPONSE,
+        ResourceType.XHR,
+    ) as interception:
+        await tab.get(sample_file("profile.html"))
+        body, _ = await interception.response_body
+        await interception.continue_request()
+
+        assert body is not None
+        # original_response = loads(body)
+        # assert original_response["name"] == "Zendriver"
diff --git a/tests/sample_data/profile.html b/tests/sample_data/profile.html
@@ -0,0 +1,23 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <title>Fetch Example</title>
+    <script>
+        window.onload = function () {
+            fetch('https://cdpdriver.github.io/examples/user-data.json')
+                .then(response => response.json())
+                .then(data => {
+                    document.getElementById('result').textContent = JSON.stringify(data, null, 2);
+                })
+                .catch(error => {
+                    document.getElementById('result').textContent = 'Error: ' + error;
+                });
+        };
+    </script>
+</head>
+<body>
+<h1>Fetch Result</h1>
+<pre id="result">Loading...</pre>
+</body>
+</html>
diff --git a/zendriver/core/connection.py b/zendriver/core/connection.py
@@ -188,6 +188,7 @@ class Connection(metaclass=CantTouchThis):
     attached: bool
     websocket: websockets.asyncio.client.ClientConnection | None = None
     _target: cdp.target.TargetInfo | None
+    _current_id_mutex: asyncio.Lock = asyncio.Lock()
 
     def __init__(
         self,
@@ -468,7 +469,8 @@ async def send(
             tx.connection = self
             if not self.mapper:
                 self.__count__ = itertools.count(0)
-            tx.id = next(self.__count__)
+            async with self._current_id_mutex:
+                tx.id = next(self.__count__)
             self.mapper.update({tx.id: tx})
             if not _is_update:
                 await self._register_handlers()
diff --git a/zendriver/core/intercept.py b/zendriver/core/intercept.py
@@ -0,0 +1,160 @@
+import asyncio
+import typing
+
+from zendriver import cdp
+from zendriver.cdp.fetch import HeaderEntry, RequestStage, RequestPattern
+from zendriver.cdp.network import ResourceType
+from zendriver.core.connection import Connection
+
+
+class BaseFetchInterception:
+    """
+    Base class to wait for a Fetch response matching a URL pattern.
+    Use this to collect and decode a paused fetch response, while keeping
+    the use block clean and returning its own result.
+
+    :param tab: The Tab instance to monitor.
+    :param url_pattern: The URL pattern to match requests and responses.
+    :param request_stage: The stage of the fetch request to intercept (e.g., request or response).
+    :param resource_type: The type of resource to intercept (e.g., document, script, etc.).
+    """
+
+    def __init__(
+        self,
+        tab: Connection,
+        url_pattern: str,
+        request_stage: RequestStage,
+        resource_type: ResourceType,
+    ):
+        self.tab = tab
+        self.url_pattern = url_pattern
+        self.request_stage = request_stage
+        self.resource_type = resource_type
+        self.response_future: asyncio.Future[cdp.fetch.RequestPaused] = asyncio.Future()
+
+    async def _response_handler(self, event: cdp.fetch.RequestPaused):
+        """
+        Internal handler for response events.
+        :param event: The response event.
+        :type event: cdp.fetch.RequestPaused
+        """
+        self._remove_response_handler()
+        self.response_future.set_result(event)
+
+    def _remove_response_handler(self):
+        """
+        Remove the response event handler.
+        """
+        self.tab.remove_handlers(cdp.fetch.RequestPaused, self._response_handler)
+
+    async def __aenter__(self):
+        """
+        Enter the context manager, adding request and response handlers.
+        """
+        await self.tab.send(
+            cdp.fetch.enable(
+                [
+                    RequestPattern(
+                        url_pattern=self.url_pattern,
+                        request_stage=self.request_stage,
+                        resource_type=self.resource_type,
+                    )
+                ]
+            )
+        )
+        self.tab.enabled_domains.append(
+            cdp.fetch
+        )  # trick to avoid another `fetch.enable` call by _register_handlers
+        self.tab.add_handler(cdp.fetch.RequestPaused, self._response_handler)
+        return self
+
+    async def __aexit__(self, *args):
+        """
+        Exit the context manager, removing request and response handlers.
+        """
+        self._remove_response_handler()
+        await self.tab.send(cdp.fetch.disable())
+
+    @property
+    async def request(self):
+        """
+        Get the matched request.
+        :return: The matched request.
+        :rtype: cdp.network.request
+        """
+        return (await self.response_future).request
+
+    @property
+    async def response_body(self) -> tuple[str, bool]:
+        """
+        Get the body of the matched response.
+        :return: The response body.
+        :rtype: str
+        """
+        request_id = (await self.response_future).request_id
+        body = await self.tab.send(cdp.fetch.get_response_body(request_id=request_id))
+        return body
+
+    async def fail_request(self, error_reason: cdp.network.ErrorReason) -> None:
+        request_id = (await self.response_future).request_id
+        await self.tab.send(
+            cdp.fetch.fail_request(request_id=request_id, error_reason=error_reason)
+        )
+
+    async def continue_request(
+        self,
+        url: typing.Optional[str] = None,
+        method: typing.Optional[str] = None,
+        post_data: typing.Optional[str] = None,
+        headers: typing.Optional[typing.List[HeaderEntry]] = None,
+        intercept_response: typing.Optional[bool] = None,
+    ) -> None:
+        request_id = (await self.response_future).request_id
+        await self.tab.send(
+            cdp.fetch.continue_request(
+                request_id=request_id,
+                url=url,
+                method=method,
+                post_data=post_data,
+                headers=headers,
+                intercept_response=intercept_response,
+            )
+        )
+
+    async def fulfill_request(
+        self,
+        response_code: int,
+        response_headers: typing.Optional[typing.List[HeaderEntry]] = None,
+        binary_response_headers: typing.Optional[str] = None,
+        body: typing.Optional[str] = None,
+        response_phrase: typing.Optional[str] = None,
+    ) -> None:
+        request_id = (await self.response_future).request_id
+        await self.tab.send(
+            cdp.fetch.fulfill_request(
+                request_id=request_id,
+                response_code=response_code,
+                response_headers=response_headers,
+                binary_response_headers=binary_response_headers,
+                body=body,
+                response_phrase=response_phrase,
+            )
+        )
+
+    async def continue_response(
+        self,
+        response_code: typing.Optional[int] = None,
+        response_phrase: typing.Optional[str] = None,
+        response_headers: typing.Optional[typing.List[HeaderEntry]] = None,
+        binary_response_headers: typing.Optional[str] = None,
+    ) -> None:
+        request_id = (await self.response_future).request_id
+        await self.tab.send(
+            cdp.fetch.continue_response(
+                request_id=request_id,
+                response_code=response_code,
+                response_phrase=response_phrase,
+                response_headers=response_headers,
+                binary_response_headers=binary_response_headers,
+            )
+        )
diff --git a/zendriver/core/tab.py b/zendriver/core/tab.py
@@ -13,11 +13,14 @@
 import webbrowser
 from typing import TYPE_CHECKING, Any, List, Literal, Optional, Tuple, Union
 
+from .intercept import BaseFetchInterception
 from .. import cdp
 from . import element, util
 from .config import PathLike
 from .connection import Connection, ProtocolException
 from .expect import DownloadExpectation, RequestExpectation, ResponseExpectation
+from ..cdp.fetch import RequestStage
+from ..cdp.network import ResourceType
 
 if TYPE_CHECKING:
     from .browser import Browser
@@ -1233,6 +1236,28 @@ def expect_download(self) -> DownloadExpectation:
         """
         return DownloadExpectation(self)
 
+    def intercept(
+        self,
+        url_pattern: str,
+        request_stage: RequestStage,
+        resource_type: ResourceType,
+    ) -> BaseFetchInterception:
+        """
+        Sets up interception for network requests matching a URL pattern, request stage, and resource type.
+
+        :param url_pattern: URL string or regex pattern to match requests.
+        :type url_pattern: Union[str, re.Pattern[str]]
+        :param request_stage: Stage of the request to intercept (e.g., request, response).
+        :type request_stage: RequestStage
+        :param resource_type: Type of resource (e.g., Document, Script, Image).
+        :type resource_type: ResourceType
+        :return: A BaseFetchInterception instance for further configuration or awaiting intercepted requests.
+        :rtype: BaseFetchInterception
+
+        Use this to block, modify, or inspect network traffic for specific resources during browser automation.
+        """
+        return BaseFetchInterception(self, url_pattern, request_stage, resource_type)
+
     async def download_file(self, url: str, filename: Optional[PathLike] = None):
         """
         downloads file by given url.
