From 1e123b9a90209415de8db55039659f719ecbe2d6 Mon Sep 17 00:00:00 2001
From: Mark Collin <mark.collin@lazeryattack.com>
Date: Mon, 11 May 2015 08:20:44 +0100
Subject: [PATCH] Add proxy support

---
 README.md                                     |  6 +++
 pom.xml                                       |  6 +++
 .../selenium/config/DriverType.java           | 48 ++++++++++++++++++-
 3 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 87a20822..b1b1e30e 100644
--- a/README.md
+++ b/README.md
@@ -43,6 +43,12 @@ You can even specify multiple threads (you can do it on a grid as well!):
 
 - -Dthreads=2
 
+You can also specify a proxy to use
+
+- -DproxyEnabled=true
+- -DproxyHost=localhost
+- -DproxyPort=8080
+
 If the tests fail screenshots will be saved in ${project.basedir}/target/screenshots
 
 If you need to force a binary overwrite you can do:
diff --git a/pom.xml b/pom.xml
index ce1b8a4b..609ca084 100644
--- a/pom.xml
+++ b/pom.xml
@@ -48,6 +48,9 @@
         <seleniumGridURL/>
         <platform/>
         <browserVersion/>
+        <proxyEnabled>false</proxyEnabled>
+        <proxyHost/>
+        <proxyPort/>
     </properties>
 
     <dependencies>
@@ -130,6 +133,9 @@
                                 <gridURL>${seleniumGridURL}</gridURL>
                                 <desiredPlatform>${platform}</desiredPlatform>
                                 <desiredBrowserVersion>${browserVersion}</desiredBrowserVersion>
+                                <proxyEnabled>${proxyEnabled}</proxyEnabled>
+                                <proxyHost>${proxyHost}</proxyHost>
+                                <proxyPort>${proxyPort}</proxyPort>
 
                                 <!--Set properties passed in by the driver binary downloader-->
                                 <phantomjs.binary.path>${phantomjs.binary.path}</phantomjs.binary.path>
diff --git a/src/test/java/com/lazerycode/selenium/config/DriverType.java b/src/test/java/com/lazerycode/selenium/config/DriverType.java
index 36152af3..0959becb 100644
--- a/src/test/java/com/lazerycode/selenium/config/DriverType.java
+++ b/src/test/java/com/lazerycode/selenium/config/DriverType.java
@@ -1,6 +1,7 @@
 package com.lazerycode.selenium.config;
 
 import org.openqa.selenium.Platform;
+import org.openqa.selenium.Proxy;
 import org.openqa.selenium.WebDriver;
 import org.openqa.selenium.chrome.ChromeDriver;
 import org.openqa.selenium.firefox.FirefoxDriver;
@@ -14,14 +15,24 @@
 
 import java.net.MalformedURLException;
 import java.net.URL;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.List;
+
+import static java.lang.Boolean.getBoolean;
+import static org.openqa.selenium.Proxy.ProxyType.MANUAL;
+import static org.openqa.selenium.remote.CapabilityType.PROXY;
 
 public enum DriverType implements DriverSetup {
 
     FIREFOX {
         public DesiredCapabilities getDesiredCapabilities() {
-            return DesiredCapabilities.firefox();
+            DesiredCapabilities capabilities = DesiredCapabilities.firefox();
+            if (proxyEnabled) {
+                capabilities.setCapability(PROXY, getSeleniumProxyDetails());
+            }
+            return capabilities;
         }
 
         public WebDriver getWebDriverObject(DesiredCapabilities capabilities) {
@@ -35,6 +46,9 @@ public DesiredCapabilities getDesiredCapabilities() {
             HashMap<String, String> chromePreferences = new HashMap<String, String>();
             chromePreferences.put("profile.password_manager_enabled", "false");
             capabilities.setCapability("chrome.prefs", chromePreferences);
+            if (proxyEnabled) {
+                capabilities.setCapability(PROXY, getSeleniumProxyDetails());
+            }
             return capabilities;
         }
 
@@ -53,6 +67,9 @@ public DesiredCapabilities getDesiredCapabilities() {
             capabilities.setCapability(CapabilityType.ForSeleniumServer.ENSURING_CLEAN_SESSION, true);
             capabilities.setCapability(InternetExplorerDriver.ENABLE_PERSISTENT_HOVERING, true);
             capabilities.setCapability("requireWindowFocus", true);
+            if (proxyEnabled) {
+                capabilities.setCapability(PROXY, getSeleniumProxyDetails());
+            }
             return capabilities;
         }
 
@@ -69,6 +86,9 @@ public String getWebDriverSystemPropertyKey() {
         public DesiredCapabilities getDesiredCapabilities() {
             DesiredCapabilities capabilities = DesiredCapabilities.safari();
             capabilities.setCapability("safari.cleanSession", true);
+            if (proxyEnabled) {
+                capabilities.setCapability(PROXY, getSeleniumProxyDetails());
+            }
             return capabilities;
         }
 
@@ -78,7 +98,11 @@ public WebDriver getWebDriverObject(DesiredCapabilities capabilities) {
     },
     OPERA {
         public DesiredCapabilities getDesiredCapabilities() {
-            return DesiredCapabilities.operaBlink();
+            DesiredCapabilities capabilities = DesiredCapabilities.operaBlink();
+            if (proxyEnabled) {
+                capabilities.setCapability(PROXY, getSeleniumProxyDetails());
+            }
+            return capabilities;
         }
 
         public WebDriver getWebDriverObject(DesiredCapabilities capabilities) {
@@ -94,6 +118,16 @@ public String getWebDriverSystemPropertyKey() {
         public DesiredCapabilities getDesiredCapabilities() {
             DesiredCapabilities capabilities = DesiredCapabilities.phantomjs();
             capabilities.setCapability("takesScreenshot", true);
+            final List<String> cliArguments = new ArrayList<String>();
+            cliArguments.add("--web-security=false");
+            cliArguments.add("--ssl-protocol=any");
+            cliArguments.add("--ignore-ssl-errors=true");
+            if (proxyEnabled) {
+                cliArguments.add("--proxy-type=http");
+                cliArguments.add("--proxy=" + proxyDetails);
+            } else {
+                cliArguments.add("--proxy-type=none");
+            }
             return capabilities;
         }
 
@@ -111,6 +145,16 @@ public String getWebDriverSystemPropertyKey() {
     public static final boolean useRemoteWebDriver = Boolean.valueOf(System.getProperty("remoteDriver"));
     private static final String operatingSystem = System.getProperties().getProperty("os.name").toUpperCase();
     private static final String systemArchitecture = System.getProperties().getProperty("os.arch");
+    public static boolean proxyEnabled = getBoolean("proxyEnabled");
+    private static final String proxyDetails = String.format("%s:%s", System.getProperty("proxy"), System.getProperty("proxyPort"));
+
+    public Proxy getSeleniumProxyDetails() {
+        Proxy proxy = new Proxy();
+        proxy.setProxyType(MANUAL);
+        proxy.setHttpProxy(proxyDetails);
+        proxy.setSslProxy(proxyDetails);
+        return proxy;
+    }
 
     public String getWebDriverSystemPropertyKey() {
         return null;