Add `mullvad proxy api edit` command

Allow a user to edit an existing custom api proxy method

7 files changed, 213 insertions, 9 deletions

diff --git a/mullvad-cli/src/cmds/proxy.rs b/mullvad-cli/src/cmds/proxy.rs
index 583accef1e..5a4bf2fdd8 100644
--- a/mullvad-cli/src/cmds/proxy.rs
+++ b/mullvad-cli/src/cmds/proxy.rs
@@ -1,6 +1,6 @@
 use anyhow::{anyhow, Result};
 use mullvad_management_interface::MullvadProxyClient;
-use mullvad_types::api_access_method::AccessMethod;
+use mullvad_types::api_access_method::{AccessMethod, ApiAccessMethodReplace};
 use std::net::IpAddr;
 
 use clap::{Args, Subcommand};
@@ -25,13 +25,15 @@ impl Proxy {
                     //println!("Adding custom proxy");
                     Self::add(cmd).await?;
                 }
-                ApiCommands::Edit(_) => todo!(),
+                ApiCommands::Edit(cmd) => {
+                    // Transform human-readable index to 0-based indexing.
+                    let index = Self::zero_to_one_based_index(cmd.index)?;
+                    Self::edit(EditCustomCommands { index, ..cmd }).await?
+                }
                 ApiCommands::Remove(cmd) => {
                     // Transform human-readable index to 0-based indexing.
-                    match cmd.index.checked_sub(1) {
-                        Some(index) => Self::remove(RemoveCustomCommands { index, ..cmd }).await?,
-                        None => println!("Access method 0 does not exist"),
-                    }
+                    let index = Self::zero_to_one_based_index(cmd.index)?;
+                    Self::remove(RemoveCustomCommands { index }).await?
                 }
             },
         };
@@ -71,6 +73,65 @@ impl Proxy {
             .await
             .map_err(Into::<anyhow::Error>::into)
     }
+
+    async fn edit(cmd: EditCustomCommands) -> Result<()> {
+        let mut rpc = MullvadProxyClient::new().await?;
+        // Retrieve the access method to edit
+        let access_method = rpc
+            .get_api_access_methods()
+            .await?
+            .get(cmd.index)
+            .ok_or(anyhow!(format!(
+                "Access method {} does not exist",
+                cmd.index + 1
+            )))?
+            .clone();
+
+        // Create a new access method combining the new params with the previous values
+        let edited_access_method: AccessMethod = match access_method {
+            AccessMethod::Shadowsocks(shadowsocks) => {
+                let ip = cmd.params.ip.unwrap_or(shadowsocks.peer.ip()).to_string();
+                let port = cmd.params.port.unwrap_or(shadowsocks.peer.port());
+                let password = cmd.params.password.unwrap_or(shadowsocks.password);
+                let cipher = cmd.params.cipher.unwrap_or(shadowsocks.cipher);
+                mullvad_types::api_access_method::Shadowsocks::from_args(ip, port, cipher, password)
+                    .map(|x| x.into())
+            }
+            AccessMethod::Socks5(socks) => match socks {
+                mullvad_types::api_access_method::Socks5::Local(local) => {
+                    let ip = cmd.params.ip.unwrap_or(local.peer.ip()).to_string();
+                    let port = cmd.params.port.unwrap_or(local.peer.port());
+                    let local_port = cmd.params.local_port.unwrap_or(local.port);
+                    mullvad_types::api_access_method::Socks5Local::from_args(ip, port, local_port)
+                        .map(|x| x.into())
+                }
+                mullvad_types::api_access_method::Socks5::Remote(remote) => {
+                    let ip = cmd.params.ip.unwrap_or(remote.peer.ip()).to_string();
+                    let port = cmd.params.port.unwrap_or(remote.peer.port());
+                    mullvad_types::api_access_method::Socks5Remote::from_args(ip, port)
+                        .map(|x| x.into())
+                }
+            },
+        }
+        .ok_or(anyhow!(
+            "Could not edit access method {}, reverting changes.",
+            cmd.index
+        ))?;
+
+        rpc.replace_access_method(ApiAccessMethodReplace {
+            index: cmd.index,
+            access_method: edited_access_method,
+        })
+        .await?;
+
+        Ok(())
+    }
+
+    fn zero_to_one_based_index(index: usize) -> Result<usize> {
+        index
+            .checked_sub(1)
+            .ok_or(anyhow!("Access method 0 does not exist"))
+    }
 }
 
 #[derive(Subcommand, Debug, Clone)]
@@ -80,6 +141,8 @@ pub enum ApiCommands {
     /// Add a custom API proxy
     #[clap(subcommand)]
     Add(AddCustomCommands),
+    /// Edit an API proxy
+    Edit(EditCustomCommands),
     /// Remove an API proxy
     Remove(RemoveCustomCommands),
 }
@@ -107,6 +170,15 @@ pub enum AddCustomCommands {
 }
 
 #[derive(Args, Debug, Clone)]
+pub struct EditCustomCommands {
+    /// Which API proxy to edit
+    index: usize,
+    /// Editing parameters
+    #[clap(flatten)]
+    params: EditParams,
+}
+
+#[derive(Args, Debug, Clone)]
 pub struct RemoveCustomCommands {
     /// Which API proxy to remove
     index: usize,
@@ -138,6 +210,28 @@ pub enum Socks5AddCommands {
     },
 }
 
+#[derive(Args, Debug, Clone)]
+pub struct EditParams {
+    /// Username for authentication [Shadowsocks]
+    #[arg(long)]
+    username: Option<String>,
+    /// Password for authentication [Shadowsocks]
+    #[arg(long)]
+    password: Option<String>,
+    /// Cipher to use [Shadowsocks]
+    #[arg(value_parser = SHADOWSOCKS_CIPHERS, long)]
+    cipher: Option<String>,
+    /// The IP of the remote proxy server [Socks5 (Local & Remote proxy), Shadowsocks]
+    #[arg(long)]
+    ip: Option<IpAddr>,
+    /// The port of the remote proxy server [Socks5 (Local & Remote proxy), Shadowsocks]
+    #[arg(long)]
+    port: Option<u16>,
+    /// The port that the server on localhost is listening on [Socks5 (Local proxy)]
+    #[arg(long)]
+    local_port: Option<u16>,
+}
+
 /// Implement conversions from CLI types to Daemon types.
 ///
 /// Since these are not supposed to be used outside of the CLI,
diff --git a/mullvad-daemon/src/access_methods.rs b/mullvad-daemon/src/access_methods.rs
index 12d8f10f26..e778cd2eeb 100644
--- a/mullvad-daemon/src/access_methods.rs
+++ b/mullvad-daemon/src/access_methods.rs
@@ -1,5 +1,5 @@
 use crate::{new_selector_config, settings, Daemon, EventListener};
-use mullvad_types::api_access_method::AccessMethod;
+use mullvad_types::api_access_method::{AccessMethod, ApiAccessMethodReplace};
 
 #[derive(err_derive::Error, Debug)]
 pub enum Error {
@@ -54,4 +54,26 @@ where
             })
             .map_err(Error::Settings)
     }
+
+    pub async fn replace_access_method(
+        &mut self,
+        access_method_replace: ApiAccessMethodReplace,
+    ) -> Result<(), Error> {
+        self.settings
+            .update(|settings| {
+                let access_methods = &mut settings.api_access_methods.api_access_methods;
+                access_methods.push(access_method_replace.access_method);
+                access_methods.swap_remove(access_method_replace.index);
+            })
+            .await
+            .map(|changed| {
+                if changed {
+                    self.event_listener
+                        .notify_settings(self.settings.to_settings());
+                    self.relay_selector
+                        .set_config(new_selector_config(&self.settings));
+                };
+            })
+            .map_err(Error::Settings)
+    }
 }
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 6ebb7f828e..9528aa1d5e 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -40,7 +40,7 @@ use mullvad_relay_selector::{
 };
 use mullvad_types::{
     account::{AccountData, AccountToken, VoucherSubmission},
-    api_access_method::AccessMethod,
+    api_access_method::{AccessMethod, ApiAccessMethodReplace},
     auth_failed::AuthFailed,
     custom_list::CustomList,
     device::{Device, DeviceEvent, DeviceEventCause, DeviceId, DeviceState, RemoveDeviceEvent},
@@ -266,6 +266,8 @@ pub enum DaemonCommand {
     AddApiAccessMethod(ResponseTx<(), Error>, AccessMethod),
     /// Remove an API access method
     RemoveApiAccessMethod(ResponseTx<(), Error>, AccessMethod),
+    /// Edit an API access method
+    ReplaceApiAccessMethod(ResponseTx<(), Error>, ApiAccessMethodReplace),
     /// Get information about the currently running and latest app versions
     GetVersionInfo(oneshot::Sender<Option<AppVersionInfo>>),
     /// Return whether the daemon is performing post-upgrade tasks
@@ -1044,6 +1046,9 @@ where
             GetApiAccessMethods(tx) => self.on_get_api_access_methods(tx),
             AddApiAccessMethod(tx, method) => self.on_add_api_access_method(tx, method).await,
             RemoveApiAccessMethod(tx, method) => self.on_remove_api_access_method(tx, method).await,
+            ReplaceApiAccessMethod(tx, method) => {
+                self.on_replace_api_access_method(tx, method).await
+            }
             IsPerformingPostUpgrade(tx) => self.on_is_performing_post_upgrade(tx),
             GetCurrentVersion(tx) => self.on_get_current_version(tx),
             #[cfg(not(target_os = "android"))]
@@ -2243,6 +2248,18 @@ where
         Self::oneshot_send(tx, result, "remove_api_access_method response");
     }
 
+    async fn on_replace_api_access_method(
+        &mut self,
+        tx: ResponseTx<(), Error>,
+        method: ApiAccessMethodReplace,
+    ) {
+        let result = self
+            .replace_access_method(method)
+            .await
+            .map_err(Error::AccessMethodError);
+        Self::oneshot_send(tx, result, "edit_api_access_method response");
+    }
+
     fn on_get_settings(&self, tx: oneshot::Sender<Settings>) {
         Self::oneshot_send(tx, self.settings.to_settings(), "get_settings response");
     }
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index 9491541dc4..2c1061a741 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -81,6 +81,9 @@ service ManagementService {
   rpc RemoveApiAccessMethod(ApiAccessMethod) returns (google.protobuf.Empty) {
     // Can I return something useful here instead of Empty?
   }
+  rpc ReplaceApiAccessMethod(ApiAccessMethodReplace) returns (google.protobuf.Empty) {
+    // Can I return something useful here instead of Empty?
+  }
 
   // Split tunneling (Linux)
   rpc GetSplitTunnelProcesses(google.protobuf.Empty) returns (stream google.protobuf.Int32Value) {}
diff --git a/mullvad-management-interface/src/client.rs b/mullvad-management-interface/src/client.rs
index d4edabc391..892f4ea5cd 100644
--- a/mullvad-management-interface/src/client.rs
+++ b/mullvad-management-interface/src/client.rs
@@ -487,6 +487,17 @@ impl MullvadProxyClient {
             .map(drop)
     }
 
+    pub async fn replace_access_method(
+        &mut self,
+        access_method_replace: ApiAccessMethodReplace,
+    ) -> Result<()> {
+        self.0
+            .replace_api_access_method(types::ApiAccessMethodReplace::from(access_method_replace))
+            .await
+            .map_err(Error::Rpc)
+            .map(drop)
+    }
+
     #[cfg(target_os = "linux")]
     pub async fn get_split_tunnel_processes(&mut self) -> Result<Vec<i32>> {
         use futures::TryStreamExt;
diff --git a/mullvad-management-interface/src/types/conversions/api_access_method.rs b/mullvad-management-interface/src/types/conversions/api_access_method.rs
index a91923e9e9..b9217b88c9 100644
--- a/mullvad-management-interface/src/types/conversions/api_access_method.rs
+++ b/mullvad-management-interface/src/types/conversions/api_access_method.rs
@@ -32,6 +32,25 @@ mod settings {
             }
         }
     }
+
+    impl From<api_access_method::ApiAccessMethodReplace> for proto::ApiAccessMethodReplace {
+        fn from(value: api_access_method::ApiAccessMethodReplace) -> Self {
+            proto::ApiAccessMethodReplace {
+                index: value.index as u32,
+                access_method: Some(value.access_method.into()),
+            }
+        }
+    }
+
+    impl From<proto::ApiAccessMethodReplace> for api_access_method::ApiAccessMethodReplace {
+        // TODO: Implement `TryFrom` instead, and skip the `unwrap`.
+        fn from(value: proto::ApiAccessMethodReplace) -> Self {
+            api_access_method::ApiAccessMethodReplace {
+                index: value.index as usize,
+                access_method: value.access_method.unwrap().into(),
+            }
+        }
+    }
 }
 
 /// Implements conversions for the 'main' AccessMethod data type.
diff --git a/mullvad-types/src/api_access_method.rs b/mullvad-types/src/api_access_method.rs
index 5a4d32551f..0cb417dda1 100644
--- a/mullvad-types/src/api_access_method.rs
+++ b/mullvad-types/src/api_access_method.rs
@@ -61,7 +61,7 @@ impl Shadowsocks {
     /// If `ip` or `port` are valid [`Some(Socks5Local)`] is returned, otherwise [`None`].
     pub fn from_args(ip: String, port: u16, cipher: String, password: String) -> Option<Self> {
         let peer = SocketAddrV4::new(Ipv4Addr::from_str(&ip).ok()?, port).into();
-        Some(Self::new(peer, password, cipher))
+        Some(Self::new(peer, cipher, password))
     }
 }
 
@@ -92,3 +92,41 @@ impl Socks5Remote {
         Some(Self::new(peer))
     }
 }
+
+impl From<Shadowsocks> for AccessMethod {
+    fn from(value: Shadowsocks) -> Self {
+        AccessMethod::Shadowsocks(value)
+    }
+}
+
+impl From<Socks5Remote> for AccessMethod {
+    fn from(value: Socks5Remote) -> Self {
+        AccessMethod::Socks5(value.into())
+    }
+}
+
+impl From<Socks5Local> for AccessMethod {
+    fn from(value: Socks5Local) -> Self {
+        AccessMethod::Socks5(value.into())
+    }
+}
+
+impl From<Socks5Remote> for Socks5 {
+    fn from(value: Socks5Remote) -> Self {
+        Socks5::Remote(value)
+    }
+}
+
+impl From<Socks5Local> for Socks5 {
+    fn from(value: Socks5Local) -> Self {
+        Socks5::Local(value)
+    }
+}
+
+/// TODO: Document why this is needed.
+/// Hint: Argument to protobuf rpc `ApiAccessMethodReplace`.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct ApiAccessMethodReplace {
+    pub index: usize,
+    pub access_method: AccessMethod,
+}