Merge branch 'adjust-cli-gui-for-wg'

3 files changed, 267 insertions, 86 deletions

diff --git a/mullvad-cli/Cargo.toml b/mullvad-cli/Cargo.toml
index efa908d363..989d27046e 100644
--- a/mullvad-cli/Cargo.toml
+++ b/mullvad-cli/Cargo.toml
@@ -18,11 +18,12 @@ name = "mullvad"
 path = "src/main.rs"
 
 [dependencies]
-clap = "2.20"
+clap = "2.32"
 error-chain = "0.12"
 env_logger = "0.5"
 serde = "1.0"
 futures = "0.1"
+base64 = "0.10"
 
 mullvad-ipc-client = { path = "../mullvad-ipc-client" }
 mullvad-types = { path = "../mullvad-types" }
diff --git a/mullvad-cli/src/cmds/relay.rs b/mullvad-cli/src/cmds/relay.rs
index 2f3d276fce..a7d571e2ce 100644
--- a/mullvad-cli/src/cmds/relay.rs
+++ b/mullvad-cli/src/cmds/relay.rs
@@ -1,15 +1,20 @@
 use crate::{new_rpc_client, Command, Result, ResultExt};
-use clap::value_t;
-use std::{net::Ipv4Addr, str::FromStr};
+use clap::{value_t, values_t};
+use std::{
+    io::{self, BufRead},
+    net::{IpAddr, Ipv4Addr, SocketAddr},
+    str::FromStr,
+};
 
 use mullvad_types::{
+    endpoint::all_of_the_internet,
     relay_constraints::{
         Constraint, LocationConstraint, OpenVpnConstraints, RelayConstraintsUpdate,
         RelaySettingsUpdate, TunnelConstraints,
     },
     ConnectionConfig, CustomTunnelEndpoint,
 };
-use talpid_types::net::{openvpn, Endpoint, TransportProtocol};
+use talpid_types::net::{openvpn, wireguard, Endpoint, TransportProtocol};
 
 pub struct Relay;
 
@@ -31,41 +36,72 @@ impl Command for Relay {
                     .subcommand(
                         clap::SubCommand::with_name("custom")
                             .about("Set a custom VPN relay")
-                            .arg(
-                                clap::Arg::with_name("tunnel")
-                                    .required(true)
-                                    .index(1)
-                                    .possible_values(&["openvpn", "wireguard"]),
+                            .subcommand(clap::SubCommand::with_name("wireguard")
+                                .arg(
+                                    clap::Arg::with_name("host")
+                                        .help("Hostname or IP")
+                                        .required(true)
+                                        .index(1),
+                                )
+                                .arg(
+                                    clap::Arg::with_name("port")
+                                        .help("Remote network port")
+                                        .required(true)
+                                        .index(2),
+                                )
+                                .arg(
+                                    clap::Arg::with_name("peer-key")
+                                        .help("Base64 encoded peer public key")
+                                        .index(3)
+                                        .required(false),
+                                )
+                                .arg(
+                                    clap::Arg::with_name("gateway")
+                                        .help("Gateway address")
+                                        .long("gateway")
+                                        .index(4)
+                                        .required(false),
+                                )
+                                .arg(
+                                    clap::Arg::with_name("addr")
+                                        .help("Local address of wireguard tunnel")
+                                        .long("addr")
+                                        .takes_value(true)
+                                        .multiple(true)
+                                        .required(false),
+                                ),
                             )
-                            .arg(
-                                clap::Arg::with_name("host")
-                                    .help("Hostname or IP")
-                                    .required(true)
-                                    .index(2),
+                            .subcommand(clap::SubCommand::with_name("openvpn")
+                                .arg(
+                                    clap::Arg::with_name("host")
+                                        .help("Hostname or IP")
+                                        .required(true)
+                                        .index(1),
+                                )
+                                .arg(
+                                    clap::Arg::with_name("port")
+                                        .help("Remote network port")
+                                        .required(true)
+                                        .index(2),
+                                )
+                                .arg(
+                                    clap::Arg::with_name("protocol")
+                                        .help("Transport protocol. For Wireguard this is ignored.")
+                                        .index(3)
+                                        .default_value("udp")
+                                        .possible_values(&["udp", "tcp"]),
+                                )
+                                .arg(
+                                    clap::Arg::with_name("username")
+                                        .help("Username to be used with the OpenVpn relay")
+                                        .index(4),
+                                )
+                                .arg(
+                                    clap::Arg::with_name("password")
+                                        .help("Password to be used with the OpenVpn relay")
+                                        .index(5),
+                                )
                             )
-                            .arg(
-                                clap::Arg::with_name("port")
-                                    .help("Remote network port")
-                                    .required(true)
-                                    .index(3),
-                            )
-                            .arg(
-                                clap::Arg::with_name("protocol")
-                                    .help("Transport protocol. For Wireguard this is ignored.")
-                                    .index(4)
-                                    .default_value("udp")
-                                    .possible_values(&["udp", "tcp"]),
-                            )
-                            .arg(
-                                clap::Arg::with_name("username")
-                                    .help("Username to be used with the OpenVpn relay")
-                                    .index(5),
-                            )
-                            .arg(
-                                clap::Arg::with_name("password")
-                                    .help("Password to be used with the OpenVpn relay")
-                                    .index(6),
-                            ),
                     )
                     .subcommand(
                         clap::SubCommand::with_name("location")
@@ -152,29 +188,83 @@ impl Relay {
     }
 
     fn set_custom(&self, matches: &clap::ArgMatches) -> Result<()> {
+        let custom_endpoint = match matches.subcommand() {
+            ("openvpn", Some(openvpn_matches)) => Self::read_custom_openvpn_relay(openvpn_matches),
+            ("wireguard", Some(wg_matches)) => Self::read_custom_wireguard_relay(wg_matches),
+            (_unknown_tunnel, _) => unreachable!("No set relay command given"),
+        };
+        self.update_constraints(RelaySettingsUpdate::CustomTunnelEndpoint(custom_endpoint))
+    }
+
+    fn read_custom_openvpn_relay(matches: &clap::ArgMatches) -> CustomTunnelEndpoint {
         let host = value_t!(matches.value_of("host"), String).unwrap_or_else(|e| e.exit());
         let port = value_t!(matches.value_of("port"), u16).unwrap_or_else(|e| e.exit());
-        let config = match matches.value_of("tunnel").unwrap() {
-            "openvpn" => {
-                let username =
-                    value_t!(matches.value_of("username"), String).unwrap_or_else(|e| e.exit());
-                let password =
-                    value_t!(matches.value_of("password"), String).unwrap_or_else(|e| e.exit());
-                let protocol = value_t!(matches.value_of("protocol"), TransportProtocol)
-                    .unwrap_or_else(|e| e.exit());
-                ConnectionConfig::OpenVpn(openvpn::ConnectionConfig {
-                    endpoint: Endpoint::new(Ipv4Addr::UNSPECIFIED, port, protocol),
-                    username,
-                    password,
-                })
-            }
-            // TODO: Gather all the data to build a WireguardEndpointData properly.
-            // "wireguard" => TunnelEndpointData::Wireguard(WireguardEndpointData { port }),
-            _ => unreachable!("Invalid tunnel protocol"),
-        };
-        self.update_constraints(RelaySettingsUpdate::CustomTunnelEndpoint(
-            CustomTunnelEndpoint::new(host, config),
-        ))
+        let username = value_t!(matches.value_of("username"), String).unwrap_or_else(|e| e.exit());
+        let password = value_t!(matches.value_of("password"), String).unwrap_or_else(|e| e.exit());
+        let protocol =
+            value_t!(matches.value_of("protocol"), TransportProtocol).unwrap_or_else(|e| e.exit());
+        CustomTunnelEndpoint::new(
+            host,
+            ConnectionConfig::OpenVpn(openvpn::ConnectionConfig {
+                endpoint: Endpoint::new(Ipv4Addr::UNSPECIFIED, port, protocol),
+                username,
+                password,
+            }),
+        )
+    }
+
+    fn read_custom_wireguard_relay(matches: &clap::ArgMatches) -> CustomTunnelEndpoint {
+        let host = value_t!(matches.value_of("host"), String).unwrap_or_else(|e| e.exit());
+        let port = value_t!(matches.value_of("port"), u16).unwrap_or_else(|e| e.exit());
+        let addresses = values_t!(matches.values_of("addr"), IpAddr).unwrap_or_else(|e| e.exit());
+        println!("addresses - {:?}", addresses);
+        let peer_key_str =
+            value_t!(matches.value_of("peer-key"), String).unwrap_or_else(|e| e.exit());
+        let gateway = value_t!(matches.value_of("gateway"), IpAddr).unwrap_or_else(|e| e.exit());
+        let mut private_key_str = String::new();
+        println!("Reading private key from standard input");
+        let _ = io::stdin().lock().read_line(&mut private_key_str);
+        if private_key_str.trim().len() == 0 {
+            eprintln!("Expected to read private key from standard input");
+        }
+        let private_key = Self::validate_wireguard_key(&private_key_str).into();
+        let peer_public_key = Self::validate_wireguard_key(&peer_key_str).into();
+
+
+        CustomTunnelEndpoint::new(
+            host,
+            ConnectionConfig::Wireguard(wireguard::ConnectionConfig {
+                tunnel: wireguard::TunnelConfig {
+                    private_key,
+                    addresses,
+                },
+                peer: wireguard::PeerConfig {
+                    public_key: peer_public_key,
+                    allowed_ips: all_of_the_internet(),
+                    endpoint: SocketAddr::new(Ipv4Addr::UNSPECIFIED.into(), port),
+                },
+                gateway,
+            }),
+        )
+    }
+
+    fn validate_wireguard_key(key_str: &str) -> [u8; 32] {
+        let key_bytes = base64::decode(key_str.trim()).unwrap_or_else(|e| {
+            eprintln!("Failed to decode wireguard key: {}", e);
+            ::std::process::exit(1);
+        });
+
+        let mut key = [0u8; 32];
+        if key_bytes.len() != 32 {
+            eprintln!(
+                "Expected key length to be 32 bytes, got {}",
+                key_bytes.len()
+            );
+            ::std::process::exit(1);
+        }
+
+        key.copy_from_slice(&key_bytes);
+        key
     }
 
     fn set_location(&self, matches: &clap::ArgMatches) -> Result<()> {
diff --git a/mullvad-cli/src/cmds/tunnel.rs b/mullvad-cli/src/cmds/tunnel.rs
index 33e75d26d4..5887148abe 100644
--- a/mullvad-cli/src/cmds/tunnel.rs
+++ b/mullvad-cli/src/cmds/tunnel.rs
@@ -18,20 +18,55 @@ impl Command for Tunnel {
             .about("Manage tunnel specific options")
             .setting(clap::AppSettings::SubcommandRequired)
             .subcommand(create_openvpn_subcommand())
+            .subcommand(create_wireguard_subcommand())
             .subcommand(create_ipv6_subcommand())
     }
 
     fn run(&self, matches: &clap::ArgMatches) -> Result<()> {
-        if let Some(openvpn_matches) = matches.subcommand_matches("openvpn") {
-            Self::handle_openvpn_cmd(openvpn_matches)
-        } else if let Some(ipv6_matches) = matches.subcommand_matches("ipv6") {
-            Self::handle_ipv6_cmd(ipv6_matches)
-        } else {
-            unreachable!("unhandled command");
+        match matches.subcommand() {
+            ("openvpn", Some(openvpn_matches)) => Self::handle_openvpn_cmd(openvpn_matches),
+            ("wireguard", Some(wg_matches)) => Self::handle_wireguard_cmd(wg_matches),
+            ("ipv6", Some(ipv6_matches)) => Self::handle_ipv6_cmd(ipv6_matches),
+            _ => {
+                unreachable!("unhandled comand");
+            }
         }
     }
 }
 
+fn create_wireguard_subcommand() -> clap::App<'static, 'static> {
+    let app = clap::SubCommand::with_name("wireguard")
+        .about("Manage options for Wireguard tunnels")
+        .setting(clap::AppSettings::SubcommandRequired)
+        .subcommand(create_wireguard_mtu_subcommand());
+    if cfg!(target_os = "linux") {
+        app.subcommand(create_wireguard_fwmark_subcommand())
+    } else {
+        app
+    }
+}
+
+fn create_wireguard_mtu_subcommand() -> clap::App<'static, 'static> {
+    clap::SubCommand::with_name("mtu")
+        .about("Configure the MTU of the wireguard tunnel")
+        .setting(clap::AppSettings::SubcommandRequired)
+        .subcommand(clap::SubCommand::with_name("get"))
+        .subcommand(clap::SubCommand::with_name("unset"))
+        .subcommand(
+            clap::SubCommand::with_name("set").arg(clap::Arg::with_name("mtu").required(true)),
+        )
+}
+
+fn create_wireguard_fwmark_subcommand() -> clap::App<'static, 'static> {
+    clap::SubCommand::with_name("fwmark")
+        .about("Configure the firewall mark used to direct traffic through Wireguard tunnel")
+        .setting(clap::AppSettings::SubcommandRequired)
+        .subcommand(clap::SubCommand::with_name("get"))
+        .subcommand(
+            clap::SubCommand::with_name("set").arg(clap::Arg::with_name("fwmark").required(true)),
+        )
+}
+
 fn create_openvpn_subcommand() -> clap::App<'static, 'static> {
     clap::SubCommand::with_name("openvpn")
         .about("Manage options for OpenVPN tunnels")
@@ -129,39 +164,94 @@ fn create_ipv6_subcommand() -> clap::App<'static, 'static> {
 
 impl Tunnel {
     fn handle_openvpn_cmd(matches: &clap::ArgMatches) -> Result<()> {
-        if let Some(m) = matches.subcommand_matches("mssfix") {
-            Self::handle_openvpn_mssfix_cmd(m)
-        } else if let Some(m) = matches.subcommand_matches("proxy") {
-            Self::handle_openvpn_proxy_cmd(m)
-        } else {
-            unreachable!("unhandled command");
+        match matches.subcommand() {
+            ("mssfix", Some(mssfix_matches)) => Self::handle_openvpn_mssfix_cmd(mssfix_matches),
+            ("proxy", Some(proxy_matches)) => Self::handle_openvpn_proxy_cmd(proxy_matches),
+            _ => unreachable!("unhandled command"),
         }
     }
 
     fn handle_openvpn_mssfix_cmd(matches: &clap::ArgMatches) -> Result<()> {
-        if matches.subcommand_matches("get").is_some() {
-            Self::process_openvpn_mssfix_get()
-        } else if matches.subcommand_matches("unset").is_some() {
-            Self::process_openvpn_mssfix_unset()
-        } else if let Some(m) = matches.subcommand_matches("set") {
-            Self::process_openvpn_mssfix_set(m)
-        } else {
-            unreachable!("unhandled command");
+        match matches.subcommand() {
+            ("get", Some(_)) => Self::process_openvpn_mssfix_get(),
+            ("unset", Some(_)) => Self::process_openvpn_mssfix_unset(),
+            ("set", Some(set_matches)) => Self::process_openvpn_mssfix_set(set_matches),
+            _ => unreachable!("unhandled command"),
         }
     }
 
     fn handle_openvpn_proxy_cmd(matches: &clap::ArgMatches) -> Result<()> {
-        if matches.subcommand_matches("get").is_some() {
-            Self::process_openvpn_proxy_get()
-        } else if matches.subcommand_matches("unset").is_some() {
-            Self::process_openvpn_proxy_unset()
-        } else if let Some(m) = matches.subcommand_matches("set") {
-            Self::process_openvpn_proxy_set(m)
-        } else {
-            unreachable!("unhandled command");
+        match matches.subcommand() {
+            ("get", Some(_)) => Self::process_openvpn_proxy_get(),
+            ("unset", Some(_)) => Self::process_openvpn_proxy_unset(),
+            ("set", Some(set_matches)) => Self::process_openvpn_proxy_set(set_matches),
+            _ => unreachable!("unhandled command"),
         }
     }
 
+    fn handle_wireguard_cmd(matches: &clap::ArgMatches) -> Result<()> {
+        match matches.subcommand() {
+            ("mtu", Some(matches)) => match matches.subcommand() {
+                ("get", _) => Self::process_wireguard_mtu_get(),
+                ("set", Some(matches)) => Self::process_wireguard_mtu_set(matches),
+                ("unset", _) => Self::process_wireguard_mtu_unset(),
+                _ => unreachable!("unhandled command"),
+            },
+
+            #[cfg(target_os = "linux")]
+            ("fwmark", Some(matches)) => match matches.subcommand() {
+                ("get", _) => Self::process_wireguard_fwmark_get(),
+                ("set", Some(fwmark_matches)) => Self::process_wireguard_fwmark_set(fwmark_matches),
+                _ => unreachable!("unhandled command"),
+            },
+            _ => unreachable!("unhandled command"),
+        }
+    }
+
+    fn process_wireguard_mtu_get() -> Result<()> {
+        let tunnel_options = Self::get_tunnel_options()?;
+        println!(
+            "mtu: {}",
+            tunnel_options
+                .wireguard
+                .mtu
+                .map(|mtu| mtu.to_string())
+                .unwrap_or("unset".into())
+        );
+        Ok(())
+    }
+
+    fn process_wireguard_mtu_set(matches: &clap::ArgMatches) -> Result<()> {
+        let mtu = value_t!(matches.value_of("mtu"), u16).unwrap_or_else(|e| e.exit());
+        let mut rpc = new_rpc_client()?;
+        rpc.set_wireguard_mtu(Some(mtu))?;
+        println!("Wireguard MTU has been updated");
+        Ok(())
+    }
+
+    fn process_wireguard_mtu_unset() -> Result<()> {
+        let mut rpc = new_rpc_client()?;
+        rpc.set_wireguard_mtu(None)?;
+        println!("Wireguard MTU has been unset");
+        Ok(())
+    }
+
+    #[cfg(target_os = "linux")]
+    fn process_wireguard_fwmark_get() -> Result<()> {
+        let tunnel_options = Self::get_tunnel_options()?;
+        println!("fwmark: {}", tunnel_options.wireguard.fwmark);
+        Ok(())
+    }
+
+    #[cfg(target_os = "linux")]
+    fn process_wireguard_fwmark_set(matches: &clap::ArgMatches) -> Result<()> {
+        let fwmark = value_t!(matches.value_of("fwmark"), i32).unwrap_or_else(|e| e.exit());
+        let mut rpc = new_rpc_client()?;
+        rpc.set_wireguard_fwmark(fwmark)?;
+        println!("Firewall mark parameter has been updated");
+        Ok(())
+    }
+
     fn handle_ipv6_cmd(matches: &clap::ArgMatches) -> Result<()> {
         if matches.subcommand_matches("get").is_some() {
             Self::process_ipv6_get()