diff --git a/seleniumwire/undetected_chromedriver/webdriver.py b/seleniumwire/undetected_chromedriver/webdriver.py index c261346..ee3e402 100644 --- a/seleniumwire/undetected_chromedriver/webdriver.py +++ b/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') diff --git a/seleniumwire/webdriver.py b/seleniumwire/webdriver.py index 14137bd..586507c 100644 --- a/seleniumwire/webdriver.py +++ b/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,9 +251,14 @@ 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) @@ -237,11 +266,12 @@ def __init__(self, seleniumwire_options=None, *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,7 +282,10 @@ 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) @@ -260,16 +293,20 @@ def __init__(self, seleniumwire_options=None, *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') diff --git a/tests/seleniumwire/test_webdriver.py b/tests/seleniumwire/test_webdriver.py index e78c032..0f3a38b 100644 --- a/tests/seleniumwire/test_webdriver.py +++ b/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()