Add ability to add extensions to Chrome as well as other options

Author: stevepryde

src/common/capabilities/chrome.rs

@@ -3,6 +3,8 @@ use serde_json::{from_value, json, to_value, Value};
 
 use crate::common::capabilities::desiredcapabilities::Capabilities;
 use crate::error::WebDriverResult;
+use serde::de::DeserializeOwned;
+use std::path::Path;
 
 #[derive(Debug, Clone, Serialize)]
 #[serde(transparent)]
@@ -26,6 +28,54 @@ impl ChromeCapabilities {
         ChromeCapabilities::default()
     }
 
+    /// Add the specified Chrome option. This is a helper method for `add_chrome_arg()`.
+    pub fn add_chrome_option<T>(&mut self, key: &str, value: T) -> WebDriverResult<()>
+    where
+        T: Serialize,
+    {
+        self.add_subkey("goog:chromeOptions", key, value)
+    }
+
+    /// Get the specified Chrome option.
+    pub fn get_chrome_option<T>(&self, key: &str) -> T
+    where
+        T: DeserializeOwned + Default,
+    {
+        from_value(self.capabilities["goog:chromeOptions"][key].clone()).unwrap_or_default()
+    }
+
+    /// Get the current list of command-line arguments to `chromedriver` as a vec.
+    pub fn get_args(&self) -> Vec<String> {
+        self.get_chrome_option("args")
+    }
+
+    /// Get the current list of Chrome extensions as a vec.
+    /// Each item is a base64-encoded string containing the .CRX extension file contents.
+    /// Use `add_extension()` to add a new extension file.
+    pub fn get_extensions(&self) -> Vec<String> {
+        self.get_chrome_option("extensions")
+    }
+
+    /// Get the path to the chrome binary (if one was previously set).
+    pub fn get_binary(&self) -> String {
+        self.get_chrome_option("binary")
+    }
+
+    /// Set the path to chrome binary to use.
+    pub fn set_binary(&mut self, path: &str) -> WebDriverResult<()> {
+        self.add_chrome_option("binary", path)
+    }
+
+    /// Get the current debugger address (if one was previously set).
+    pub fn get_debugger_address(&self) -> String {
+        self.get_chrome_option("debuggerAddress")
+    }
+
+    /// Set the debugger address.
+    pub fn set_debugger_address(&mut self, address: &str) -> WebDriverResult<()> {
+        self.add_chrome_option("debuggerAddress", address)
+    }
+
     /// Add the specified command-line argument to `chromedriver`. Eg. "--disable-local-storage"
     /// The full list of switches can be found here:
     /// [https://chromium.googlesource.com/chromium/src/+/master/chrome/common/chrome_switches.cc](https://chromium.googlesource.com/chromium/src/+/master/chrome/common/chrome_switches.cc)
@@ -38,6 +88,7 @@ impl ChromeCapabilities {
         self.add_chrome_option("args", to_value(args)?)
     }
 
+    /// Remove the specified Chrome command-line argument if it had been added previously.
     pub fn remove_chrome_arg(&mut self, arg: &str) -> WebDriverResult<()> {
         let mut args = self.get_args();
         if args.is_empty() {
@@ -48,24 +99,52 @@ impl ChromeCapabilities {
         }
     }
 
-    /// Add the specified chrome option. This is a helper method for `add_chrome_arg()`.
-    pub fn add_chrome_option<T>(&mut self, key: &str, value: T) -> WebDriverResult<()>
-    where
-        T: Serialize,
-    {
-        self.add_subkey("goog:chromeOptions", key, value)
+    /// Add a base64-encoded extension.
+    pub fn add_encoded_extension(&mut self, extension_base64: &str) -> WebDriverResult<()> {
+        let mut extensions = self.get_extensions();
+        let ext_string = extension_base64.to_string();
+        if !extensions.contains(&ext_string) {
+            extensions.push(ext_string);
+        }
+        self.add_chrome_option("extensions", to_value(extensions)?)
     }
 
-    /// Get the current list of command-line arguments to `chromedriver` as a vec.
-    pub fn get_args(&self) -> Vec<String> {
-        from_value(self.capabilities["goog:chromeOptions"]["args"].clone()).unwrap_or_default()
+    /// Remove the specified base64-encoded extension if it had been added previously.
+    pub fn remove_encoded_extension(&mut self, extension_base64: &str) -> WebDriverResult<()> {
+        let mut extensions = self.get_extensions();
+        if extensions.is_empty() {
+            Ok(())
+        } else {
+            extensions.retain(|v| v != extension_base64);
+            self.add_chrome_option("extensions", to_value(extensions)?)
+        }
+    }
+
+    /// Add Chrome extension file. This will be a file with a .CRX extension.
+    pub fn add_extension(&mut self, crx_file: &Path) -> WebDriverResult<()> {
+        let contents = std::fs::read(crx_file)?;
+        let b64_contents = base64::encode(contents);
+        self.add_encoded_extension(&b64_contents)
+    }
+
+    /// Remove the specified Chrome extension file if an identical extension had been added
+    /// previously.
+    pub fn remove_extension(&mut self, crx_file: &Path) -> WebDriverResult<()> {
+        let contents = std::fs::read(crx_file)?;
+        let b64_contents = base64::encode(contents);
+        self.remove_encoded_extension(&b64_contents)
     }
 
     /// Set the browser to run headless.
     pub fn set_headless(&mut self) -> WebDriverResult<()> {
         self.add_chrome_arg("--headless")
     }
 
+    /// Unset the headless option.
+    pub fn unset_headless(&mut self) -> WebDriverResult<()> {
+        self.remove_chrome_arg("--headless")
+    }
+
     /// Set disable web security.
     pub fn set_disable_web_security(&mut self) -> WebDriverResult<()> {
         self.add_chrome_arg("--disable-web-security")
@@ -76,12 +155,12 @@ impl ChromeCapabilities {
         self.remove_chrome_arg("--disable-web-security")
     }
 
-    /// Sets accept untrusted certs.
+    /// Set ignore certificate errors.
     pub fn set_ignore_certificate_errors(&mut self) -> WebDriverResult<()> {
         self.add_chrome_arg("--ignore-certificate-errors")
     }
 
-    /// Unsets accept untrusted certs.
+    /// Unset ignore certificate errors.
     pub fn unset_ignore_certificate_errors(&mut self) -> WebDriverResult<()> {
         self.remove_chrome_arg("--ignore-certificate-errors")
     }