Repository: wkeeling/selenium-wire

Issue:
importing uc from seleniumwire ignores proxy option
On newer versions of undetected chromedriver, I need to use this arg ``use_subprocess=True`` else I can't run it without getting a Runtime error. (https://github.com/ultrafunkamsterdam/undetected-chromedriver/issues/432) This argument seems to break the usage of proxies with selenium wire though.

```python
import seleniumwire.undetected_chromedriver.v2 as uc
from seleniumwire import webdriver

chrome_options = webdriver.ChromeOptions()
chrome_options.headless = False
seleniumwire_options = {'proxy': {'https': 'type://host:port',}} 
driver = uc.Chrome(options=chrome_options, seleniumwire_options=seleniumwire_options, use_subprocess=True)

driver.get('https://whatismyipaddress.com/')
```

Ignores passed proxy option and loads website without proxy.

Comments:

Patch:
diff --git a/README.rst b/README.rst
index 17e62f7..bb02330 100644
--- a/README.rst
+++ b/README.rst
@@ -694,21 +694,19 @@ If you wish to take advantage of this make sure you have undetected_chromedriver
 
     pip install undetected-chromedriver
 
-Then you can select the version of undetected_chromedriver you want to use by importing ``Chrome`` and ``ChromeOptions`` from the appropriate package.
-
-For undetected_chromedriver version 1:
+Then in your code, import the ``seleniumwire.undetected_chromedriver`` package:
 
 .. code:: python
 
-    from seleniumwire.undetected_chromedriver import Chrome, ChromeOptions
-
-For undetected_chromedriver version 2:
+    import seleniumwire.undetected_chromedriver as uc
 
-.. code:: python
+    chrome_options = uc.ChromeOptions()
 
-    from seleniumwire.undetected_chromedriver.v2 import Chrome, ChromeOptions
+    driver = uc.Chrome(
+        options=chrome_options,
+        seleniumwire_options={}
+    )
 
-See the `undetected_chromedriver docs <https://github.com/ultrafunkamsterdam/undetected-chromedriver>`_ for differences between the two versions.
 
 Certificates
 ~~~~~~~~~~~~
diff --git a/seleniumwire/storage.py b/seleniumwire/storage.py
index 6d1cd66..1be95da 100644
--- a/seleniumwire/storage.py
+++ b/seleniumwire/storage.py
@@ -181,12 +181,9 @@ class RequestStorage:
         request_dir = self._get_request_dir(request_id)
 
         with open(os.path.join(request_dir, 'request'), 'rb') as req:
-            try:
-                request = pickle.load(req)
-            except Exception:
-                # Errors may sometimes occur with unpickling - e.g.
-                # sometimes data hasn't been fully flushed to disk
-                # by the OS by the time we come to unpickle it.
+            request = self._unpickle(req)
+
+            if request is None:
                 return None
 
             ws_messages = self._ws_messages.get(request.id)
@@ -198,19 +195,37 @@ class RequestStorage:
             try:
                 # Attach the response if there is one.
                 with open(os.path.join(request_dir, 'response'), 'rb') as res:
-                    response = pickle.load(res)
-                    request.response = response
-
-                    # The certificate data has been stored on the response but we make
-                    # it available on the request which is a more logical location.
-                    if hasattr(response, 'cert'):
-                        request.cert = response.cert
-                        del response.cert
+                    response = self._unpickle(res)
+
+                    if response is not None:
+                        request.response = response
+
+                        # The certificate data has been stored on the response but we make
+                        # it available on the request which is a more logical location.
+                        if hasattr(response, 'cert'):
+                            request.cert = response.cert
+                            del response.cert
             except (FileNotFoundError, EOFError):
                 pass
 
         return request
 
+    def _unpickle(self, f):
+        """Unpickle the object specified by the file f.
+
+        If unpickling fails return None.
+        """
+        try:
+            return pickle.load(f)
+        except Exception:
+            # Errors may sometimes occur with unpickling - e.g.
+            # sometimes data hasn't been fully flushed to disk
+            # by the OS by the time we come to unpickle it.
+            if log.isEnabledFor(logging.DEBUG):
+                log.exception('Error unpickling object')
+
+            return None
+
     def load_last_request(self) -> Optional[Request]:
         """Load the last saved request.
 
@@ -240,8 +255,10 @@ class RequestStorage:
 
             try:
                 with open(os.path.join(request_dir, 'har_entry'), 'rb') as f:
-                    entry = pickle.load(f)
-                    entries.append(entry)
+                    entry = self._unpickle(f)
+
+                    if entry is not None:
+                        entries.append(entry)
             except FileNotFoundError:
                 # HAR entries aren't necessarily saved with each request.
                 pass
diff --git a/seleniumwire/undetected_chromedriver/__init__.py b/seleniumwire/undetected_chromedriver/__init__.py
index a470664..8eaa1cd 100644
--- a/seleniumwire/undetected_chromedriver/__init__.py
+++ b/seleniumwire/undetected_chromedriver/__init__.py
@@ -1,12 +1,1 @@
-try:
-    import undetected_chromedriver as uc
-except ImportError as e:
-    raise ImportError(
-        'undetected_chromedriver not found. ' 'Install it with `pip install undetected_chromedriver`.'
-    ) from e
-
-from seleniumwire.webdriver import Chrome
-
-uc._Chrome = Chrome
-Chrome = uc.Chrome
-ChromeOptions = uc.ChromeOptions  # noqa: F811
+from .webdriver import Chrome, ChromeOptions
diff --git a/seleniumwire/undetected_chromedriver/v2.py b/seleniumwire/undetected_chromedriver/v2.py
index c4a42c9..e0c9dd0 100644
--- a/seleniumwire/undetected_chromedriver/v2.py
+++ b/seleniumwire/undetected_chromedriver/v2.py
@@ -1,58 +1,1 @@
-import logging
-
-import undetected_chromedriver.v2 as uc
-from selenium.webdriver import DesiredCapabilities
-
-from seleniumwire.inspect import InspectRequestsMixin
-from seleniumwire.utils import urlsafe_address
-from seleniumwire.webdriver import DriverCommonMixin
-
-log = logging.getLogger(__name__)
-
-
-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):
-        """Initialise a new Chrome WebDriver instance.
-
-        Args:
-            seleniumwire_options: The seleniumwire options dictionary.
-        """
-        if seleniumwire_options is None:
-            seleniumwire_options = {}
-
-        config = self._setup_backend(seleniumwire_options)
-
-        if seleniumwire_options.get('auto_config', True):
-            capabilities = kwargs.get('desired_capabilities')
-            if capabilities is None:
-                capabilities = DesiredCapabilities.CHROME
-            capabilities = capabilities.copy()
-
-            capabilities.update(config)
-
-            kwargs['desired_capabilities'] = capabilities
-
-        try:
-            chrome_options = kwargs['options']
-        except KeyError:
-            chrome_options = ChromeOptions()
-
-        log.info('Using undetected_chromedriver.v2')
-
-        # We need to point Chrome back to Selenium Wire since the executable
-        # will be started separately by undetected_chromedriver.
-        addr, port = urlsafe_address(self.backend.address())
-        chrome_options.add_argument(f'--proxy-server={addr}:{port}')
-        chrome_options.add_argument(
-            f"--proxy-bypass-list={','.join(seleniumwire_options.get('exclude_hosts', ['<-loopback>']))}"
-        )
-
-        kwargs['options'] = chrome_options
-
-        super().__init__(*args, **kwargs)
-
-
-ChromeOptions = uc.ChromeOptions  # noqa: F811
+from .webdriver import Chrome, ChromeOptions  # noqa: F401
diff --git a/seleniumwire/undetected_chromedriver/webdriver.py b/seleniumwire/undetected_chromedriver/webdriver.py
new file mode 100644
index 0000000..c261346
--- /dev/null
+++ b/seleniumwire/undetected_chromedriver/webdriver.py
@@ -0,0 +1,64 @@
+import logging
+
+from selenium.webdriver import DesiredCapabilities
+
+try:
+    import undetected_chromedriver as uc
+except ImportError as e:
+    raise ImportError(
+        'undetected_chromedriver not found. ' 'Install it with `pip install undetected_chromedriver`.'
+    ) from e
+
+from seleniumwire.inspect import InspectRequestsMixin
+from seleniumwire.utils import urlsafe_address
+from seleniumwire.webdriver import DriverCommonMixin
+
+log = logging.getLogger(__name__)
+
+
+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):
+        """Initialise a new Chrome WebDriver instance.
+
+        Args:
+            seleniumwire_options: The seleniumwire options dictionary.
+        """
+        if seleniumwire_options is None:
+            seleniumwire_options = {}
+
+        config = self._setup_backend(seleniumwire_options)
+
+        if seleniumwire_options.get('auto_config', True):
+            capabilities = kwargs.get('desired_capabilities')
+            if capabilities is None:
+                capabilities = DesiredCapabilities.CHROME
+            capabilities = capabilities.copy()
+
+            capabilities.update(config)
+
+            kwargs['desired_capabilities'] = capabilities
+
+        try:
+            chrome_options = kwargs['options']
+        except KeyError:
+            chrome_options = ChromeOptions()
+
+        log.info('Using undetected_chromedriver')
+
+        # We need to point Chrome back to Selenium Wire since the executable
+        # will be started separately by undetected_chromedriver.
+        addr, port = urlsafe_address(self.backend.address())
+        chrome_options.add_argument(f'--proxy-server={addr}:{port}')
+        chrome_options.add_argument(
+            f"--proxy-bypass-list={','.join(seleniumwire_options.get('exclude_hosts', ['<-loopback>']))}"
+        )
+
+        kwargs['options'] = chrome_options
+
+        super().__init__(*args, **kwargs)
+
+
+ChromeOptions = uc.ChromeOptions  # noqa: F811
diff --git a/seleniumwire/webdriver.py b/seleniumwire/webdriver.py
index cb876e2..115a0ea 100644
--- a/seleniumwire/webdriver.py
+++ b/seleniumwire/webdriver.py
@@ -26,7 +26,7 @@ SELENIUM_V4 = parse_version(getattr(selenium, '__version__', '0')) >= parse_vers
 
 
 class DriverCommonMixin:
-    """Operations common to all webdriver types."""
+    """Attributes common to all webdriver types."""
 
     def _setup_backend(self, seleniumwire_options: Dict[str, Any]) -> Dict[str, Any]:
         """Create the backend proxy server and return its configuration
@@ -121,7 +121,7 @@ class DriverCommonMixin:
 
 
 class Firefox(InspectRequestsMixin, DriverCommonMixin, _Firefox):
-    """Extends the Firefox webdriver to provide additional methods for inspecting requests."""
+    """Extend the Firefox webdriver to provide additional methods for inspecting requests."""
 
     def __init__(self, *args, seleniumwire_options=None, **kwargs):
         """Initialise a new Firefox WebDriver instance.
@@ -174,7 +174,7 @@ class Firefox(InspectRequestsMixin, DriverCommonMixin, _Firefox):
 
 
 class Chrome(InspectRequestsMixin, DriverCommonMixin, _Chrome):
-    """Extends the Chrome webdriver to provide additional methods for inspecting requests."""
+    """Extend the Chrome webdriver to provide additional methods for inspecting requests."""
 
     def __init__(self, *args, seleniumwire_options=None, **kwargs):
         """Initialise a new Chrome WebDriver instance.
@@ -207,7 +207,7 @@ class Chrome(InspectRequestsMixin, DriverCommonMixin, _Chrome):
 
 
 class Safari(InspectRequestsMixin, DriverCommonMixin, _Safari):
-    """Extends the Safari webdriver to provide additional methods for inspecting requests."""
+    """Extend the Safari webdriver to provide additional methods for inspecting requests."""
 
     def __init__(self, seleniumwire_options=None, *args, **kwargs):
         """Initialise a new Safari WebDriver instance.
@@ -230,7 +230,7 @@ class Safari(InspectRequestsMixin, DriverCommonMixin, _Safari):
 
 
 class Edge(InspectRequestsMixin, DriverCommonMixin, _Edge):
-    """Extends the Edge webdriver to provide additional methods for inspecting requests."""
+    """Extend the Edge webdriver to provide additional methods for inspecting requests."""
 
     def __init__(self, seleniumwire_options=None, *args, **kwargs):
         """Initialise a new Edge WebDriver instance.
@@ -253,7 +253,7 @@ class Edge(InspectRequestsMixin, DriverCommonMixin, _Edge):
 
 
 class Remote(InspectRequestsMixin, DriverCommonMixin, _Remote):
-    """Extends the Remote webdriver to provide additional methods for inspecting requests."""
+    """Extend the Remote webdriver to provide additional methods for inspecting requests."""
 
     def __init__(self, *args, seleniumwire_options=None, **kwargs):
         """Initialise a new Firefox WebDriver instance.
diff --git a/setup.py b/setup.py
index 40a722f..013c914 100644
--- a/setup.py
+++ b/setup.py
@@ -61,7 +61,7 @@ setup(
             'gunicorn',
             'httpbin',
             'isort',
-            'mitmproxy',  # Needed for end2end tests
+            "mitmproxy>5.3.0; python_version>='3.8.0'",  # Needed for end2end tests
             'mypy',
             'pre-commit',
             'pytest',


Test patch:
diff --git a/tests/seleniumwire/test_storage.py b/tests/seleniumwire/test_storage.py
index a31ab9b..67f83a8 100644
--- a/tests/seleniumwire/test_storage.py
+++ b/tests/seleniumwire/test_storage.py
@@ -212,6 +212,18 @@ class RequestStorageTest(TestCase):
 
         self.assertIsNotNone(requests[0].response)
 
+    @patch('seleniumwire.storage.pickle')
+    def test_load_response_unpickle_error(self, mock_pickle):
+        request = self._create_request()
+        self.storage.save_request(request)
+        response = self._create_response()
+        self.storage.save_response(request.id, response)
+        mock_pickle.load.side_effect = [request, Exception]
+
+        requests = self.storage.load_requests()
+
+        self.assertIsNone(requests[0].response)
+
     def test_load_last_request(self):
         request_1 = self._create_request()
         request_2 = self._create_request()