Added proxy support

Author: shaunpud

selenium_extensions/core.py

@@ -221,6 +221,7 @@ class SeleniumDriver:
         executable_path (str): path to the browser's webdriver binary. If set to ``None`` selenium will serach for browser's webdriver in ``$PATH``.
         run_headless (bool): boolean flag that indicates if webdriver has to be headless (without GUI).
         load_images (bool): boolean flag that indicates if webdriver has to render images.
+        use_proxy (str): use http proxy in <host:port> format.
 
     Raises:
         selenium_extensions.exceptions.SeleniumExtensionsException: ``browser`` is not supported by ``selenium_extensions``.
@@ -254,12 +255,13 @@ def goto_google(self):
         Firefox doesn't support native headless mode. We use ``pyvirtualdisplay`` to simulate it. In order ``pyvirtualdisplay`` to work you need to install ``Xvfb`` package: ``sudo apt install xvfb``.
     '''
 
-    def __init__(self, browser=None, executable_path=None, run_headless=False, load_images=True):
+    def __init__(self, browser=None, executable_path=None, run_headless=False, load_images=True, use_proxy=None):
         self._initialize_driver(browser, executable_path,
-                                run_headless, load_images)
+                                run_headless, load_images,
+                                use_proxy)
         self._initialize_methods()
 
-    def _initialize_driver(self, browser, executable_path, run_headless, load_images):
+    def _initialize_driver(self, browser, executable_path, run_headless, load_images, use_proxy):
         if browser is None:
             self.driver = chrome_driver()
             browser = browser.lower()
@@ -272,11 +274,13 @@ def _initialize_driver(self, browser, executable_path, run_headless, load_images
         if browser == 'chrome':
             self.driver = chrome_driver(executable_path=executable_path,
                                         run_headless=run_headless,
-                                        load_images=load_images)
+                                        load_images=load_images,
+                                        use_proxy=use_proxy)
         if browser == 'firefox':
             self.driver = firefox_driver(executable_path=executable_path,
                                          run_headless=run_headless,
-                                         load_images=load_images)
+                                         load_images=load_images,
+                                         use_proxy=use_proxy)
 
     def _initialize_methods(self):
         self.shut_down = partial(shut_down, self.driver)

selenium_extensions/drivers.py

@@ -5,13 +5,14 @@
 
 
 def chrome_driver(executable_path=None, run_headless=False,
-                  load_images=True):
+                  load_images=True, use_proxy=None):
     '''Function to initialize ``selenium.webdriver.Chrome`` with extended options
 
     Args:
         executable_path (str): path to the chromedriver binary. If set to ``None`` selenium will serach for ``chromedriver`` in ``$PATH``.
         run_headless (bool): boolean flag that indicates if chromedriver has to be headless (without GUI).
         load_images (bool): boolean flag that indicates if Chrome has to render images.
+        use_proxy (str): use http proxy in <host:port> format.
 
     Returns:
         selenium.webdriver.Chrome: created driver.
@@ -28,6 +29,8 @@ def chrome_driver(executable_path=None, run_headless=False,
     if not load_images:
         prefs = {'profile.managed_default_content_settings.images': 2}
         chrome_options.add_experimental_option('prefs', prefs)
+    if use_proxy:
+        chrome_options.add_argument('proxy-server=' + use_proxy)
     if executable_path:
         driver = webdriver.Chrome(chrome_options=chrome_options,
                                   executable_path=executable_path)
@@ -37,13 +40,14 @@ def chrome_driver(executable_path=None, run_headless=False,
 
 
 def firefox_driver(executable_path=None, run_headless=False,
-                   load_images=True):
+                   load_images=True, use_proxy=None):
     '''Function to initialize ``selenium.webdriver.Firefox`` with extended options
 
     Args:
         executable_path (str): path to the ``geckdriver`` binary. If set to ``None`` selenium will search for ``geckdriver`` in ``$PATH``.
         run_headless (bool): boolean flag that indicates if ``geckodriver`` has to be headless (without GUI). ``geckodriver`` doesn't support native headless mode, that's why ``pyvirtualdisplay`` is used.
         load_images (bool): boolean flag that indicates if Firefox has to render images.
+        use_proxy (str): use http proxy in <host:port> format.
 
     Returns:
         selenium.webdriver.Firefox: created driver.
@@ -73,6 +77,13 @@ def firefox_driver(executable_path=None, run_headless=False,
         # Turns animated images off
         firefox_profile.set_preference(
             'thatoneguydotnet.QuickJava.startupStatus.AnimatedImage', 2)
+    if use_proxy:
+        _ = use_proxy.split(':')
+        firefox_profile.set_preference('network.proxy.type', 1)
+        firefox_profile.set_preference('network.proxy.http', _[0])
+        firefox_profile.set_preference('network.proxy.http_port', int(_[1]))
+        firefox_profile.set_preference('network.proxy.ssl', _[0])
+        firefox_profile.set_preference('network.proxy.ssl_port', int(_[1]))
     if executable_path:
         driver = webdriver.Firefox(
             firefox_profile, executable_path=executable_path)