use super::{relay::resolve_location_constraint, relay_constraints::LocationArgs}; use anyhow::{Result, anyhow, bail}; use clap::Subcommand; use mullvad_management_interface::MullvadProxyClient; use mullvad_types::{ constraints::Constraint, relay_constraints::GeographicLocationConstraint, relay_list::RelayList, }; /// Custom list length, expressed as a number of UTF8 codepoints (i.e. chars). pub const CUSTOM_LIST_MAX_LEN: usize = 30; #[derive(Subcommand, Debug)] pub enum CustomList { /// Create a new custom list New { /// A name for the new custom list #[clap(value_parser = parse_custom_list_name)] name: String, }, /// Show all custom lists or retrieve a specific custom list List { // TODO: Would be cool to provide dynamic auto-completion: // https://github.com/clap-rs/clap/issues/1232 /// A custom list. If omitted, all custom lists are shown name: Option, }, /// Edit a custom list #[clap(subcommand)] Edit(EditCommand), /// Delete a custom list Delete { /// A custom list name: String, }, } #[derive(Subcommand, Debug)] pub enum EditCommand { /// Add a location to some custom list Add { /// A custom list name: String, #[command(flatten)] location: LocationArgs, }, /// Remove a location from some custom list Remove { /// A custom list name: String, #[command(flatten)] location: LocationArgs, }, /// Rename a custom list Rename { /// Current name of the custom list name: String, /// A new name for the custom list #[clap(value_parser = parse_custom_list_name)] new_name: String, }, } impl CustomList { pub async fn handle(self) -> Result<()> { match self { CustomList::List { name: None } => Self::list().await, CustomList::List { name: Some(name) } => Self::get(name).await, CustomList::New { name } => Self::create_list(name).await, CustomList::Delete { name } => Self::delete_list(name).await, CustomList::Edit(cmd) => match cmd { EditCommand::Add { name, location } => Self::add_location(name, location).await, EditCommand::Rename { name, new_name } => Self::rename_list(name, new_name).await, EditCommand::Remove { name, location } => { Self::remove_location(name, location).await } }, } } /// Print all custom lists. async fn list() -> Result<()> { let mut rpc = MullvadProxyClient::new().await?; let cache = rpc.get_relay_locations().await?; for custom_list in rpc.get_settings().await?.custom_lists { Self::print_custom_list(&custom_list, &cache) } Ok(()) } /// Print a specific custom list (if it exists). /// If the list does not exist, print an error. async fn get(name: String) -> Result<()> { let mut rpc = MullvadProxyClient::new().await?; let custom_list = find_list_by_name(&mut rpc, &name).await?; let cache = rpc.get_relay_locations().await?; Self::print_custom_list_content(&custom_list, &cache); Ok(()) } async fn create_list(name: String) -> Result<()> { let mut rpc = MullvadProxyClient::new().await?; rpc.create_custom_list(name).await?; Ok(()) } async fn add_location(name: String, location_args: LocationArgs) -> Result<()> { let mut rpc = MullvadProxyClient::new().await?; // Don't filter out any hosts, i.e. allow adding even inactive ones let relay_filter = |_: &_| true; let location_constraint = resolve_location_constraint(&mut rpc, location_args, relay_filter).await?; match location_constraint { Constraint::Any => bail!("\"any\" is not a valid location"), Constraint::Only(location) => { let mut list = find_list_by_name(&mut rpc, &name).await?; if list.locations.insert(location) { rpc.update_custom_list(list).await?; println!("Location added to custom-list") } else { bail!("Provided location is already present in custom-list") }; } } Ok(()) } async fn remove_location(name: String, location_args: LocationArgs) -> Result<()> { let mut rpc = MullvadProxyClient::new().await?; // Don't filter out any hosts, i.e. allow adding even inactive ones let relay_filter = |_: &_| true; let location_constraint = resolve_location_constraint(&mut rpc, location_args, relay_filter).await?; match location_constraint { Constraint::Any => bail!("\"any\" is not a valid location"), Constraint::Only(location) => { let mut list = find_list_by_name(&mut rpc, &name).await?; if list.locations.remove(&location) { rpc.update_custom_list(list).await?; println!("Location removed from custom-list") } else { bail!("Provided location was not present in custom-list") }; } } Ok(()) } async fn delete_list(name: String) -> Result<()> { let mut rpc = MullvadProxyClient::new().await?; let list = find_list_by_name(&mut rpc, &name).await?; rpc.delete_custom_list(list.id()).await?; Ok(()) } async fn rename_list(name: String, new_name: String) -> Result<()> { let mut rpc = MullvadProxyClient::new().await?; let mut list = find_list_by_name(&mut rpc, &name).await?; list.name = new_name; rpc.update_custom_list(list).await?; Ok(()) } fn print_custom_list(custom_list: &mullvad_types::custom_list::CustomList, cache: &RelayList) { println!("{}", custom_list.name); Self::print_custom_list_content(custom_list, cache); } fn print_custom_list_content( custom_list: &mullvad_types::custom_list::CustomList, cache: &RelayList, ) { for location in &custom_list.locations { println!( "\t{}", GeographicLocationConstraintFormatter::from_constraint(location, cache) ); } } } /// Struct used for pretty printing [`GeographicLocationConstraint`] with /// human-readable names for countries and cities. pub struct GeographicLocationConstraintFormatter<'a> { constraint: &'a GeographicLocationConstraint, country: Option, city: Option, } impl<'a> GeographicLocationConstraintFormatter<'a> { fn from_constraint(constraint: &'a GeographicLocationConstraint, cache: &RelayList) -> Self { use GeographicLocationConstraint::*; let (country_code, city_code) = match constraint { Country(country) => (Some(country), None), City(country, city) | Hostname(country, city, _) => (Some(country), Some(city)), }; let country = country_code.and_then(|country_code| cache.lookup_country(country_code.clone())); let city = city_code.and_then(|city_code| { country.and_then(|country| country.lookup_city(city_code.clone())) }); Self { constraint, country: country.map(|x| x.name.clone()), city: city.map(|x| x.name.clone()), } } } impl std::fmt::Display for GeographicLocationConstraintFormatter<'_> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> { let unwrap_country = |country: Option, constraint: &str| { country.unwrap_or(format!("{constraint} ")) }; let unwrap_city = |city: Option, constraint: &str| { city.unwrap_or(format!("{constraint} ")) }; match &self.constraint { GeographicLocationConstraint::Country(country) => { let rich_country = unwrap_country(self.country.clone(), country); write!(f, "{rich_country} ({country})") } GeographicLocationConstraint::City(country, city) => { let rich_country = unwrap_country(self.country.clone(), country); let rich_city = unwrap_city(self.city.clone(), city); write!(f, "{rich_city}, {rich_country} ({city}, {country})") } GeographicLocationConstraint::Hostname(country, city, hostname) => { let rich_country = unwrap_country(self.country.clone(), country); let rich_city = unwrap_city(self.city.clone(), city); write!( f, "{hostname} in {rich_city}, {rich_country} ({city}, {country})" ) } } } } pub async fn find_list_by_name( rpc: &mut MullvadProxyClient, name: &str, ) -> Result { rpc.get_settings() .await? .custom_lists .into_iter() .find(|list| list.name == name) .ok_or(anyhow!("List not found")) } /// Trim the string and validate the length against [CUSTOM_LIST_MAX_LEN]. // NOTE: should only be used when *creating* custom lists, as we don't want to make it impossible // to reference any custom lists created before the max length and whitespace restrictions were put // in place. fn parse_custom_list_name(s: &str) -> Result { let s = s.trim(); let length = s.chars().count(); if length > CUSTOM_LIST_MAX_LEN { bail!("Provided name is too long, {length}/{CUSTOM_LIST_MAX_LEN} characters."); } Ok(s.to_string()) }