Add possibility to pass your own MitmProxy (backend) to webdriver

seleniumwire/undetected_chromedriver/webdriver.py

@@ -12,6 +12,7 @@
 from seleniumwire.inspect import InspectRequestsMixin
 from seleniumwire.utils import urlsafe_address
 from seleniumwire.webdriver import DriverCommonMixin
+from seleniumwire.server import MitmProxy
 
 log = logging.getLogger(__name__)
 
@@ -20,16 +21,20 @@ class Chrome(InspectRequestsMixin, DriverCommonMixin, uc.Chrome):
     """Extends the undetected_chrome Chrome webdriver to provide additional
     methods for inspecting requests."""
 
-    def __init__(self, *args, seleniumwire_options=None, **kwargs):
+    def __init__(self, *args, seleniumwire_options=None, mitm_proxy: MitmProxy = None, **kwargs):
         """Initialise a new Chrome WebDriver instance.
 
         Args:
             seleniumwire_options: The seleniumwire options dictionary.
+            mitm_proxy: if you pass your own MitmProxy, seleniumwire will use it insteadof creating a new one
         """
         if seleniumwire_options is None:
             seleniumwire_options = {}
 
-        config = self._setup_backend(seleniumwire_options)
+        if mitm_proxy is None:
+            config = self._setup_backend(seleniumwire_options)
+        else:
+            config = self._set_backend(mitm_proxy, seleniumwire_options)
 
         if seleniumwire_options.get('auto_config', True):
             capabilities = kwargs.get('desired_capabilities')

seleniumwire/webdriver.py

@@ -26,6 +26,7 @@
 
 from seleniumwire import backend, utils
 from seleniumwire.inspect import InspectRequestsMixin
+from seleniumwire.server import MitmProxy
 
 SELENIUM_V4 = parse_version(getattr(selenium, '__version__', '0')) >= parse_version('4.0.0')
 
@@ -42,8 +43,22 @@ def _setup_backend(self, seleniumwire_options: Dict[str, Any]) -> Dict[str, Any]
             port=seleniumwire_options.get('port', 0),
             options=seleniumwire_options,
         )
+        return self._get_config(self.backend, seleniumwire_options)
 
-        addr, port = utils.urlsafe_address(self.backend.address())
+    def _set_backend(self, mitm_proxy: MitmProxy, seleniumwire_options: Dict[str, Any]) -> Dict[str, Any]:
+        """Set the backend proxy server and return its configuration
+        in a dictionary
+
+        :param mitm_proxy: proxy server
+        :param seleniumwire_options:
+        :return:
+        """
+        self.backend = mitm_proxy
+        return self._get_config(self.backend, seleniumwire_options)
+
+    @staticmethod
+    def _get_config(mitm_proxy: MitmProxy, seleniumwire_options: Dict[str, Any]) -> Dict[str, Any]:
+        addr, port = utils.urlsafe_address(mitm_proxy.address())
 
         config = {
             'proxy': {
@@ -128,11 +143,12 @@ def proxy(self, proxy_conf: Dict[str, Any]):
 class Firefox(InspectRequestsMixin, DriverCommonMixin, _Firefox):
     """Extend the Firefox webdriver to provide additional methods for inspecting requests."""
 
-    def __init__(self, *args, seleniumwire_options=None, **kwargs):
+    def __init__(self, *args, seleniumwire_options=None, mitm_proxy: MitmProxy = None, **kwargs):
         """Initialise a new Firefox WebDriver instance.
 
         Args:
             seleniumwire_options: The seleniumwire options dictionary.
+            mitm_proxy: if you pass your own MitmProxy, seleniumwire will use it insteadof creating a new one
         """
         if seleniumwire_options is None:
             seleniumwire_options = {}
@@ -148,7 +164,10 @@ def __init__(self, *args, seleniumwire_options=None, **kwargs):
         firefox_options.set_preference('network.proxy.allow_hijacking_localhost', True)
         firefox_options.accept_insecure_certs = True
 
-        config = self._setup_backend(seleniumwire_options)
+        if mitm_proxy is None:
+            config = self._setup_backend(seleniumwire_options)
+        else:
+            config = self._set_backend(mitm_proxy, seleniumwire_options)
 
         if seleniumwire_options.get('auto_config', True):
             if SELENIUM_V4:
@@ -181,11 +200,12 @@ def __init__(self, *args, seleniumwire_options=None, **kwargs):
 class Chrome(InspectRequestsMixin, DriverCommonMixin, _Chrome):
     """Extend the Chrome webdriver to provide additional methods for inspecting requests."""
 
-    def __init__(self, *args, seleniumwire_options=None, **kwargs):
+    def __init__(self, *args, seleniumwire_options=None, mitm_proxy: MitmProxy = None, **kwargs):
         """Initialise a new Chrome WebDriver instance.
 
         Args:
             seleniumwire_options: The seleniumwire options dictionary.
+            mitm_proxy: if you pass your own MitmProxy, seleniumwire will use it insteadof creating a new one
         """
         if seleniumwire_options is None:
             seleniumwire_options = {}
@@ -202,7 +222,10 @@ def __init__(self, *args, seleniumwire_options=None, **kwargs):
         chrome_options.add_argument('--proxy-bypass-list=<-loopback>')
         kwargs['options'] = chrome_options
 
-        config = self._setup_backend(seleniumwire_options)
+        if mitm_proxy is None:
+            config = self._setup_backend(seleniumwire_options)
+        else:
+            config = self._set_backend(mitm_proxy, seleniumwire_options)
 
         if seleniumwire_options.get('auto_config', True):
             for key, value in config.items():
@@ -214,11 +237,12 @@ def __init__(self, *args, seleniumwire_options=None, **kwargs):
 class Safari(InspectRequestsMixin, DriverCommonMixin, _Safari):
     """Extend the Safari webdriver to provide additional methods for inspecting requests."""
 
-    def __init__(self, seleniumwire_options=None, *args, **kwargs):
+    def __init__(self, seleniumwire_options=None, *args, mitm_proxy: MitmProxy = None, **kwargs):
         """Initialise a new Safari WebDriver instance.
 
         Args:
             seleniumwire_options: The seleniumwire options dictionary.
+            mitm_proxy: if you pass your own MitmProxy, seleniumwire will use it insteadof creating a new one
         """
         if seleniumwire_options is None:
             seleniumwire_options = {}
@@ -227,21 +251,27 @@ def __init__(self, seleniumwire_options=None, *args, **kwargs):
         # DesiredCapabilities API, and thus has to be configured manually.
         # Whatever port number is chosen for that manual configuration has to
         # be passed in the options.
-        assert 'port' in seleniumwire_options, 'You must set a port number in the seleniumwire_options'
+        assert mitm_proxy is not None or 'port' in seleniumwire_options,\
+            'You must set a port number in the seleniumwire_options ' \
+            '(or pass your own MitmProxy in mitm_proxy param)'
 
-        self._setup_backend(seleniumwire_options)
+        if mitm_proxy is None:
+            self._setup_backend(seleniumwire_options)
+        else:
+            self._set_backend(mitm_proxy, seleniumwire_options)
 
         super().__init__(*args, **kwargs)
 
 
 class Edge(InspectRequestsMixin, DriverCommonMixin, _Edge):
     """Extend the Edge webdriver to provide additional methods for inspecting requests."""
 
-    def __init__(self, seleniumwire_options=None, *args, **kwargs):
+    def __init__(self, seleniumwire_options=None, *args, mitm_proxy: MitmProxy = None, **kwargs):
         """Initialise a new Edge WebDriver instance.
 
         Args:
             seleniumwire_options: The seleniumwire options dictionary.
+            mitm_proxy: if you pass your own MitmProxy, seleniumwire will use it insteadof creating a new one
         """
         if seleniumwire_options is None:
             seleniumwire_options = {}
@@ -252,24 +282,31 @@ def __init__(self, seleniumwire_options=None, *args, **kwargs):
         # be passed in the options.
         assert 'port' in seleniumwire_options, 'You must set a port number in the seleniumwire_options'
 
-        self._setup_backend(seleniumwire_options)
+        if mitm_proxy is None:
+            self._setup_backend(seleniumwire_options)
+        else:
+            self._set_backend(mitm_proxy, seleniumwire_options)
 
         super().__init__(*args, **kwargs)
 
 
 class Remote(InspectRequestsMixin, DriverCommonMixin, _Remote):
     """Extend the Remote webdriver to provide additional methods for inspecting requests."""
 
-    def __init__(self, *args, seleniumwire_options=None, **kwargs):
+    def __init__(self, *args, seleniumwire_options=None, mitm_proxy: MitmProxy = None, **kwargs):
         """Initialise a new Firefox WebDriver instance.
 
         Args:
             seleniumwire_options: The seleniumwire options dictionary.
+            mitm_proxy: if you pass your own MitmProxy, seleniumwire will use it insteadof creating a new one
         """
         if seleniumwire_options is None:
             seleniumwire_options = {}
 
-        config = self._setup_backend(seleniumwire_options)
+        if mitm_proxy is None:
+            config = self._setup_backend(seleniumwire_options)
+        else:
+            config = self._set_backend(mitm_proxy, seleniumwire_options)
 
         if seleniumwire_options.get('auto_config', True):
             capabilities = kwargs.get('desired_capabilities')

tests/seleniumwire/test_webdriver.py

@@ -1,9 +1,10 @@
 from unittest.mock import patch
 
 import pytest
+from selenium.webdriver import chrome
 from selenium.webdriver.common.proxy import ProxyType
 
-from seleniumwire.webdriver import Chrome, ChromeOptions, Firefox
+from seleniumwire.webdriver import Chrome, ChromeOptions, Firefox, Safari
 
 
 @pytest.fixture(autouse=True)
@@ -37,6 +38,17 @@ def test_create_backend(self, mock_backend):
         mock_backend.create.assert_called_once_with(addr='127.0.0.1', port=0, options={})
         mock_backend.create.return_value.address.assert_called_once_with()
 
+    def test_pass_backend(self):
+        from seleniumwire import backend
+        with patch.object(backend, "create") as mock_create:
+            mock_create.return_value.address.return_value = ('127.0.0.1', 12345)
+            mitm_proxy = backend.create(addr='127.0.0.1', port=0, options={})
+            firefox = Firefox(mitm_proxy=mitm_proxy)
+
+        assert firefox.backend
+        mock_create.assert_called_once_with(addr='127.0.0.1', port=0, options={})
+        mock_create.return_value.address.assert_called_once_with()
+
     def test_allow_hijacking_localhost(self, firefox_super_kwargs):
         Firefox()
 
@@ -128,6 +140,17 @@ def test_create_backend(self, mock_backend):
         mock_backend.create.assert_called_once_with(addr='127.0.0.1', port=0, options={})
         mock_backend.create.return_value.address.assert_called_once_with()
 
+    def test_pass_backend(self):
+        from seleniumwire import backend
+        with patch.object(backend, "create") as mock_create:
+            mock_create.return_value.address.return_value = ('127.0.0.1', 12345)
+            mitm_proxy = backend.create(addr='127.0.0.1', port=0, options={})
+            chrome = Chrome(mitm_proxy=mitm_proxy)
+
+        assert chrome.backend
+        mock_create.assert_called_once_with(addr='127.0.0.1', port=0, options={})
+        mock_create.return_value.address.assert_called_once_with()
+
     def test_proxy_bypass_list(self, chrome_super_kwargs):
         Chrome()