Merge branch 'override-relay-ip' into main

CHANGELOG.md

@@ -31,6 +31,7 @@ Line wrap the file at 100 chars.                                              Th
   `mullvad api-access`, and the initially supported network protocols are `Shadowsocks` and
   `SOCKS5`.
 - Add social media content blocker.
+- Add ability to override server IPs to the CLI.
 
 ### Changed
 - Update Electron from 25.2.0 to 26.3.0.

mullvad-cli/src/cmds/mod.rs

@@ -1,4 +1,5 @@
 use clap::builder::{PossibleValuesParser, TypedValueParser, ValueParser};
+use std::io::stdin;
 use std::ops::Deref;
 
 pub mod account;
@@ -80,3 +81,28 @@ impl std::fmt::Display for BooleanOption {
         }
     }
 }
+
+async fn receive_confirmation(msg: &'static str, default: bool) -> bool {
+    let helper_str = match default {
+        true => "[Y/n]",
+        false => "[y/N]",
+    };
+
+    println!("{msg} {helper_str}");
+
+    tokio::task::spawn_blocking(move || loop {
+        let mut buf = String::new();
+        if let Err(e) = stdin().read_line(&mut buf) {
+            eprintln!("Couldn't read from STDIN: {e}");
+            return false;
+        }
+        match buf.trim().to_ascii_lowercase().as_str() {
+            "" => return default,
+            "y" | "ye" | "yes" => return true,
+            "n" | "no" => return false,
+            _ => eprintln!("Unexpected response. Please enter \"y\" or \"n\""),
+        }
+    })
+    .await
+    .unwrap()
+}

mullvad-cli/src/cmds/relay.rs

@@ -3,16 +3,17 @@ use clap::Subcommand;
 use itertools::Itertools;
 use mullvad_management_interface::MullvadProxyClient;
 use mullvad_types::{
-    location::Location,
+    location::{CountryCode, Location},
     relay_constraints::{
         Constraint, GeographicLocationConstraint, LocationConstraint, LocationConstraintFormatter,
-        Match, OpenVpnConstraints, Ownership, Provider, Providers, RelayConstraints, RelaySettings,
-        TransportPort, WireguardConstraints,
+        Match, OpenVpnConstraints, Ownership, Provider, Providers, RelayConstraints, RelayOverride,
+        RelaySettings, TransportPort, WireguardConstraints,
     },
     relay_list::{RelayEndpointData, RelayListCountry},
     ConnectionConfig, CustomTunnelEndpoint,
 };
 use std::{
+    collections::HashMap,
     io::BufRead,
     net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
 };
@@ -21,7 +22,7 @@ use talpid_types::net::{
 };
 
 use super::{relay_constraints::LocationArgs, BooleanOption};
-use crate::print_option;
+use crate::{cmds::receive_confirmation, print_option};
 
 #[derive(Subcommand, Debug)]
 pub enum Relay {
@@ -37,6 +38,10 @@ pub enum Relay {
 
     /// Update the relay list
     Update,
+
+    /// Override options for individual relays/servers
+    #[clap(subcommand)]
+    Override(OverrideCommands),
 }
 
 #[derive(Subcommand, Debug, Clone)]
@@ -201,13 +206,58 @@ pub enum SetCustomCommands {
     },
 }
 
+#[derive(Subcommand, Debug, Clone)]
+pub enum OverrideCommands {
+    /// Show current custom fields for servers
+    Get,
+    /// Set a custom field for a server
+    #[clap(subcommand)]
+    Set(OverrideSetCommands),
+    /// Unset a custom field for a server
+    #[clap(subcommand)]
+    Unset(OverrideUnsetCommands),
+    /// Unset custom IPs for all servers
+    ClearAll {
+        /// Clear overrides without asking for confirmation
+        #[arg(long, short = 'y', default_value_t = false)]
+        confirm: bool,
+    },
+}
+
+#[derive(Subcommand, Debug, Clone)]
+pub enum OverrideSetCommands {
+    /// Override entry IPv4 address for a given relay
+    Ipv4 {
+        /// The unique hostname for the server to set the override on
+        hostname: String,
+        /// The IPv4 address to use to connect to this server
+        address: Ipv4Addr,
+    },
+    /// Override entry IPv6 address for a given relay
+    Ipv6 {
+        /// The unique hostname for the server to set the override on
+        hostname: String,
+        /// The IPv6 address to use to connect to this server
+        address: Ipv6Addr,
+    },
+}
+
+#[derive(Subcommand, Debug, Clone)]
+pub enum OverrideUnsetCommands {
+    /// Remove overridden entry IPv4 address for the given server
+    Ipv4 { hostname: String },
+    /// Remove overridden entry IPv6 address for the given server
+    Ipv6 { hostname: String },
+}
+
 impl Relay {
     pub async fn handle(self) -> Result<()> {
         match self {
             Relay::Get => Self::get().await,
             Relay::List => Self::list().await,
             Relay::Update => Self::update().await,
             Relay::Set(subcmd) => Self::set(subcmd).await,
+            Relay::Override(subcmd) => Self::r#override(subcmd).await,
         }
     }
 
@@ -662,6 +712,164 @@ impl Relay {
         })
         .await
     }
+
+    async fn update_override(
+        hostname: &str,
+        update_fn: impl FnOnce(&mut RelayOverride),
+        warn_non_existent_hostname: bool,
+    ) -> Result<()> {
+        let mut rpc = MullvadProxyClient::new().await?;
+        let settings = rpc.get_settings().await?;
+
+        if warn_non_existent_hostname {
+            let countries = get_filtered_relays_with_client(&mut rpc).await?;
+            if find_relay_by_hostname(&countries, hostname).is_none() {
+                eprintln!("Warning: Setting overrides for an unrecognized server");
+            }
+        }
+
+        let mut relay_overrides = settings.relay_overrides;
+        let mut element = relay_overrides
+            .iter()
+            .position(|elem| elem.hostname == hostname)
+            .map(|index| relay_overrides.swap_remove(index))
+            .unwrap_or_else(|| RelayOverride::empty(hostname.to_owned()));
+        update_fn(&mut element);
+
+        rpc.set_relay_override(element).await?;
+        println!("Updated override options for {hostname}");
+        Ok(())
+    }
+
+    async fn r#override(subcmd: OverrideCommands) -> Result<()> {
+        match subcmd {
+            OverrideCommands::Get => {
+                let mut rpc = MullvadProxyClient::new().await?;
+                let settings = rpc.get_settings().await?;
+
+                let mut overrides = HashMap::new();
+                for relay_override in settings.relay_overrides {
+                    overrides.insert(relay_override.hostname.clone(), relay_override);
+                }
+
+                struct Country {
+                    name: String,
+                    code: CountryCode,
+                    cities: Vec<City>,
+                }
+                struct City {
+                    name: String,
+                    code: CountryCode,
+                    overrides: Vec<RelayOverride>,
+                }
+
+                let mut countries_with_overrides = vec![];
+                for country in get_filtered_relays().await? {
+                    let mut country_with_overrides = Country {
+                        name: country.name,
+                        code: country.code,
+                        cities: vec![],
+                    };
+
+                    for city in country.cities {
+                        let mut city_with_overrides = City {
+                            name: city.name,
+                            code: city.code,
+                            overrides: vec![],
+                        };
+
+                        for relay in city.relays {
+                            if let Some(relay_override) = overrides.remove(&relay.hostname) {
+                                city_with_overrides.overrides.push(relay_override);
+                            }
+                        }
+
+                        if !city_with_overrides.overrides.is_empty() {
+                            country_with_overrides.cities.push(city_with_overrides);
+                        }
+                    }
+
+                    if !country_with_overrides.cities.is_empty() {
+                        countries_with_overrides.push(country_with_overrides);
+                    }
+                }
+
+                let print_relay_override = |relay_override: RelayOverride| {
+                    println!("{:<8}{}:", " ", relay_override.hostname);
+                    if let Some(ipv4) = relay_override.ipv4_addr_in {
+                        println!("{:<12}ipv4: {ipv4}", " ");
+                    }
+                    if let Some(ipv6) = relay_override.ipv6_addr_in {
+                        println!("{:<12}ipv6: {ipv6}", " ");
+                    }
+                };
+
+                for country in countries_with_overrides {
+                    println!("{} ({})", country.name, country.code);
+                    for city in country.cities {
+                        println!("{:<4}{} ({})", " ", city.name, city.code);
+                        for relay_override in city.overrides {
+                            print_relay_override(relay_override);
+                        }
+                    }
+                }
+
+                if !overrides.is_empty() {
+                    println!("Overrides for unrecognized servers. Consider removing these!");
+                    for relay_override in overrides.into_values() {
+                        print_relay_override(relay_override);
+                    }
+                }
+            }
+            OverrideCommands::Set(set_cmds) => match set_cmds {
+                OverrideSetCommands::Ipv4 { hostname, address } => {
+                    Self::update_override(
+                        &hostname,
+                        |relay_override| relay_override.ipv4_addr_in = Some(address),
+                        true,
+                    )
+                    .await?;
+                }
+                OverrideSetCommands::Ipv6 { hostname, address } => {
+                    Self::update_override(
+                        &hostname,
+                        |relay_override| relay_override.ipv6_addr_in = Some(address),
+                        true,
+                    )
+                    .await?;
+                }
+            },
+            OverrideCommands::Unset(cmds) => match cmds {
+                OverrideUnsetCommands::Ipv4 { hostname } => {
+                    Self::update_override(
+                        &hostname,
+                        |relay_override| relay_override.ipv4_addr_in = None,
+                        false,
+                    )
+                    .await?;
+                }
+                OverrideUnsetCommands::Ipv6 { hostname } => {
+                    Self::update_override(
+                        &hostname,
+                        |relay_override| relay_override.ipv6_addr_in = None,
+                        false,
+                    )
+                    .await?;
+                }
+            },
+            OverrideCommands::ClearAll { confirm } => {
+                if confirm
+                    || receive_confirmation("Are you sure you want to clear all overrides?", true)
+                        .await
+                {
+                    let mut rpc = MullvadProxyClient::new().await?;
+                    rpc.clear_all_relay_overrides().await?;
+                    println!("All overrides unset");
+                }
+            }
+        }
+        Ok(())
+    }
 }
 
 fn parse_transport_port(
@@ -714,8 +922,16 @@ pub fn find_relay_by_hostname(
         })
 }
 
+/// Return a list of all active non-bridge relays
 pub async fn get_filtered_relays() -> Result<Vec<RelayListCountry>> {
     let mut rpc = MullvadProxyClient::new().await?;
+    get_filtered_relays_with_client(&mut rpc).await
+}
+
+/// Return a list of all active non-bridge relays
+async fn get_filtered_relays_with_client(
+    rpc: &mut MullvadProxyClient,
+) -> Result<Vec<RelayListCountry>> {
     let relay_list = rpc.get_relay_locations().await?;
 
     let mut countries = vec![];

mullvad-cli/src/cmds/reset.rs

@@ -1,32 +1,14 @@
+use super::receive_confirmation;
 use anyhow::Result;
 use mullvad_management_interface::MullvadProxyClient;
-use std::io::stdin;
 
 pub async fn handle() -> Result<()> {
-    if receive_confirmation().await {
-        let mut rpc = MullvadProxyClient::new().await?;
-        rpc.factory_reset().await?;
-        #[cfg(target_os = "linux")]
-        println!("If you're running systemd, to remove all logs, you must use journalctl");
+    if !receive_confirmation("Are you sure you want to disconnect, log out, delete all settings, logs and cache files for the Mullvad VPN system service?", false).await {
+        return Ok(());
     }
+    let mut rpc = MullvadProxyClient::new().await?;
+    rpc.factory_reset().await?;
+    #[cfg(target_os = "linux")]
+    println!("If you're running systemd, to remove all logs, you must use journalctl");
     Ok(())
 }
-
-async fn receive_confirmation() -> bool {
-    println!("Are you sure you want to disconnect, log out, delete all settings, logs and cache files for the Mullvad VPN system service? [Yes/No (default)]");
-
-    tokio::task::spawn_blocking(|| loop {
-        let mut buf = String::new();
-        if let Err(e) = stdin().read_line(&mut buf) {
-            eprintln!("Couldn't read from STDIN: {e}");
-            return false;
-        }
-        match buf.trim() {
-            "Yes" => return true,
-            "No" | "no" | "" => return false,
-            _ => eprintln!("Unexpected response. Please enter \"Yes\" or \"No\""),
-        }
-    })
-    .await
-    .unwrap()
-}