Handle an edge case when using a proxy with auth

Author: mdmintz

seleniumbase/core/browser_launcher.py

@@ -1,4 +1,5 @@
 import os
+import random
 import re
 import sys
 import threading
@@ -15,6 +16,7 @@
 from seleniumbase import drivers  # webdriver storage folder for SeleniumBase
 DRIVER_DIR = os.path.dirname(os.path.realpath(drivers.__file__))
 PROXY_ZIP_PATH = proxy_helper.PROXY_ZIP_PATH
+PROXY_ZIP_PATH_2 = proxy_helper.PROXY_ZIP_PATH_2
 PLATFORM = sys.platform
 IS_WINDOWS = False
 LOCAL_CHROMEDRIVER = None
@@ -57,22 +59,24 @@ def _add_chrome_proxy_extension(
         chrome_options, proxy_string, proxy_user, proxy_pass):
     """ Implementation of https://stackoverflow.com/a/35293284 for
         https://stackoverflow.com/questions/12848327/
-        (Run Selenium on a proxy server that requires authentication.)
-        The retry_on_exception is only needed for multithreaded runs
-        because proxy.zip is a common file shared between all tests
-        in a single run. """
+        (Run Selenium on a proxy server that requires authentication.) """
     if not "".join(sys.argv) == "-c":
         # Single-threaded
         proxy_helper.create_proxy_zip(proxy_string, proxy_user, proxy_pass)
     else:
         # Pytest multi-threaded test
         lock = threading.Lock()
         with lock:
+            time.sleep(random.uniform(0.02, 0.15))
             if not os.path.exists(PROXY_ZIP_PATH):
                 proxy_helper.create_proxy_zip(
                     proxy_string, proxy_user, proxy_pass)
-            time.sleep(0.3)
-    chrome_options.add_extension(PROXY_ZIP_PATH)
+            time.sleep(random.uniform(0.1, 0.2))
+    proxy_zip = PROXY_ZIP_PATH
+    if not os.path.exists(PROXY_ZIP_PATH):
+        # Handle "Permission denied" on the default proxy.zip path
+        proxy_zip = PROXY_ZIP_PATH_2
+    chrome_options.add_extension(proxy_zip)
     return chrome_options
 
 
@@ -102,7 +106,6 @@ def _set_chrome_options(
         if proxy_auth:
             chrome_options = _add_chrome_proxy_extension(
                 chrome_options, proxy_string, proxy_user, proxy_pass)
-            chrome_options.add_extension(DRIVER_DIR + "/proxy.zip")
         chrome_options.add_argument('--proxy-server=%s' % proxy_string)
     return chrome_options
 
@@ -209,9 +212,6 @@ def get_driver(browser_name, headless=False, use_grid=False,
                     "either Chrome or Firefox may be used.)")
         proxy_string = validate_proxy_string(proxy_string)
         if proxy_string and proxy_user and proxy_pass:
-            if not os.path.exists(PROXY_ZIP_PATH):
-                proxy_helper.create_proxy_zip(
-                    proxy_string, proxy_user, proxy_pass)
             proxy_auth = True
     if use_grid:
         return get_remote_driver(

seleniumbase/core/proxy_helper.py

@@ -1,9 +1,12 @@
 import os
 import threading
 import zipfile
+from seleniumbase.fixtures import constants
 from seleniumbase import drivers
 DRIVER_DIR = os.path.dirname(os.path.realpath(drivers.__file__))
 PROXY_ZIP_PATH = "%s/%s" % (DRIVER_DIR, "proxy.zip")
+DOWNLOADS_DIR = constants.Files.DOWNLOADS_FOLDER
+PROXY_ZIP_PATH_2 = "%s/%s" % (DOWNLOADS_DIR, "proxy.zip")
 
 
 def create_proxy_zip(proxy_string, proxy_user, proxy_pass):
@@ -15,7 +18,6 @@ def create_proxy_zip(proxy_string, proxy_user, proxy_pass):
     """
     proxy_host = proxy_string.split(':')[0]
     proxy_port = proxy_string.split(':')[1]
-    proxy_zip = DRIVER_DIR + '/proxy.zip'
     background_js = (
         """var config = {\n"""
         """    mode: "fixed_servers",\n"""
@@ -64,7 +66,15 @@ def create_proxy_zip(proxy_string, proxy_user, proxy_pass):
         '''}''')
     lock = threading.RLock()  # Support multi-threaded test runs with Pytest
     with lock:
-        zf = zipfile.ZipFile(proxy_zip, mode='w')
+        try:
+            zf = zipfile.ZipFile(PROXY_ZIP_PATH, mode='w')
+        except IOError:
+            # Handle "Permission denied" on the default proxy.zip path
+            abs_path = os.path.abspath('.')
+            downloads_path = os.path.join(abs_path, DOWNLOADS_DIR)
+            if not os.path.exists(downloads_path):
+                os.mkdir(downloads_path)
+            zf = zipfile.ZipFile(PROXY_ZIP_PATH_2, mode='w')
         zf.writestr("background.js", background_js)
         zf.writestr("manifest.json", manifest_json)
         zf.close()
@@ -78,5 +88,7 @@ def remove_proxy_zip_if_present():
     try:
         if os.path.exists(PROXY_ZIP_PATH):
             os.remove(PROXY_ZIP_PATH)
+        elif os.path.exists(PROXY_ZIP_PATH_2):
+            os.remove(PROXY_ZIP_PATH_2)
     except Exception:
         pass