1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177

use serde::Serialize;
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)]
pub struct ChromeCapabilities {
    capabilities: Value,
}

impl Default for ChromeCapabilities {
    fn default() -> Self {
        ChromeCapabilities {
            capabilities: json!({
                "browserName": "chrome"
            }),
        }
    }
}

impl ChromeCapabilities {
    /// Create a new ChromeCapabilities struct.
    pub fn new() -> Self {
        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)
    pub fn add_chrome_arg(&mut self, arg: &str) -> WebDriverResult<()> {
        let mut args = self.get_args();
        let arg_string = arg.to_string();
        if !args.contains(&arg_string) {
            args.push(arg_string);
        }
        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() {
            Ok(())
        } else {
            args.retain(|v| v != arg);
            self.add_chrome_option("args", to_value(args)?)
        }
    }

    /// 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)?)
    }

    /// 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")
    }

    /// Unset disable web security.
    pub fn unset_disable_web_security(&mut self) -> WebDriverResult<()> {
        self.remove_chrome_arg("--disable-web-security")
    }

    /// Set ignore certificate errors.
    pub fn set_ignore_certificate_errors(&mut self) -> WebDriverResult<()> {
        self.add_chrome_arg("--ignore-certificate-errors")
    }

    /// Unset ignore certificate errors.
    pub fn unset_ignore_certificate_errors(&mut self) -> WebDriverResult<()> {
        self.remove_chrome_arg("--ignore-certificate-errors")
    }
}

impl Capabilities for ChromeCapabilities {
    fn get(&self) -> &Value {
        &self.capabilities
    }

    fn get_mut(&mut self) -> &mut Value {
        &mut self.capabilities
    }
}