Add CLI parameters for disable/enable peer (#248)

* Add CLI parameters for disable/enable peer Fixes #214. * Formatting * Remove redundant clones * Require name for yes param Yes param only makes sense if name is provided. * Formatting
tonarino · Jan 15, 2024 · aa1ac51 · aa1ac51
1 parent 489f505
commit aa1ac51
Show file tree

Hide file tree

Showing 4 changed files with 86 additions and 34 deletions.
diff --git a/client/src/main.rs b/client/src/main.rs
@@ -10,9 +10,9 @@ use shared::{
     prompts,
     wg::{DeviceExt, PeerInfoExt},
     AddCidrOpts, AddDeleteAssociationOpts, AddPeerOpts, Association, AssociationContents, Cidr,
-    CidrTree, DeleteCidrOpts, Endpoint, EndpointContents, InstallOpts, Interface, IoErrorContext,
-    ListenPortOpts, NatOpts, NetworkOpts, OverrideEndpointOpts, Peer, RedeemContents,
-    RenamePeerOpts, State, WrappedIoError, REDEEM_TRANSITION_WAIT,
+    CidrTree, DeleteCidrOpts, EnableDisablePeerOpts, Endpoint, EndpointContents, InstallOpts,
+    Interface, IoErrorContext, ListenPortOpts, NatOpts, NetworkOpts, OverrideEndpointOpts, Peer,
+    RedeemContents, RenamePeerOpts, State, WrappedIoError, REDEEM_TRANSITION_WAIT,
 };
 use std::{
     fmt, io,
@@ -211,10 +211,20 @@ enum Command {
     },
 
     /// Disable an enabled peer
-    DisablePeer { interface: Interface },
+    DisablePeer {
+        interface: Interface,
+
+        #[clap(flatten)]
+        sub_opts: EnableDisablePeerOpts,
+    },
 
     /// Enable a disabled peer
-    EnablePeer { interface: Interface },
+    EnablePeer {
+        interface: Interface,
+
+        #[clap(flatten)]
+        sub_opts: EnableDisablePeerOpts,
+    },
 
     /// Add an association between CIDRs
     AddAssociation {
@@ -817,6 +827,7 @@ fn rename_peer(
 fn enable_or_disable_peer(
     interface: &InterfaceName,
     opts: &Opts,
+    sub_opts: EnableDisablePeerOpts,
     enable: bool,
 ) -> Result<(), Error> {
     let InterfaceConfig { server, .. } =
@@ -826,7 +837,7 @@ fn enable_or_disable_peer(
     log::info!("Fetching peers.");
     let peers: Vec<Peer> = api.http("GET", "/admin/peers")?;
 
-    if let Some(peer) = prompts::enable_or_disable_peer(&peers[..], enable)? {
+    if let Some(peer) = prompts::enable_or_disable_peer(&peers[..], &sub_opts, enable)? {
         let Peer { id, mut contents } = peer;
         contents.is_disabled = !enable;
         api.http_form("PUT", &format!("/admin/peers/{id}"), contents)?;
@@ -1245,8 +1256,14 @@ fn run(opts: &Opts) -> Result<(), Error> {
             sub_opts,
         } => delete_cidr(&interface, opts, sub_opts)?,
         Command::ListCidrs { interface, tree } => list_cidrs(&interface, opts, tree)?,
-        Command::DisablePeer { interface } => enable_or_disable_peer(&interface, opts, false)?,
-        Command::EnablePeer { interface } => enable_or_disable_peer(&interface, opts, true)?,
+        Command::DisablePeer {
+            interface,
+            sub_opts,
+        } => enable_or_disable_peer(&interface, opts, sub_opts, false)?,
+        Command::EnablePeer {
+            interface,
+            sub_opts,
+        } => enable_or_disable_peer(&interface, opts, sub_opts, true)?,
         Command::AddAssociation {
             interface,
             sub_opts,

diff --git a/server/src/main.rs b/server/src/main.rs
@@ -9,8 +9,8 @@ use parking_lot::{Mutex, RwLock};
 use rusqlite::Connection;
 use serde::{Deserialize, Serialize};
 use shared::{
-    get_local_addrs, AddCidrOpts, AddPeerOpts, DeleteCidrOpts, Endpoint, IoErrorContext,
-    NetworkOpts, PeerContents, RenamePeerOpts, INNERNET_PUBKEY_HEADER,
+    get_local_addrs, AddCidrOpts, AddPeerOpts, DeleteCidrOpts, EnableDisablePeerOpts, Endpoint,
+    IoErrorContext, NetworkOpts, PeerContents, RenamePeerOpts, INNERNET_PUBKEY_HEADER,
 };
 use std::{
     collections::{HashMap, VecDeque},
@@ -95,10 +95,20 @@ enum Command {
     },
 
     /// Disable an enabled peer
-    DisablePeer { interface: Interface },
+    DisablePeer {
+        interface: Interface,
+
+        #[clap(flatten)]
+        args: EnableDisablePeerOpts,
+    },
 
     /// Enable a disabled peer
-    EnablePeer { interface: Interface },
+    EnablePeer {
+        interface: Interface,
+
+        #[clap(flatten)]
+        args: EnableDisablePeerOpts,
+    },
 
     /// Rename an existing peer.
     RenamePeer {
@@ -271,11 +281,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         } => serve(*interface, &conf, routing).await?,
         Command::AddPeer { interface, args } => add_peer(&interface, &conf, args, opts.network)?,
         Command::RenamePeer { interface, args } => rename_peer(&interface, &conf, args)?,
-        Command::DisablePeer { interface } => {
-            enable_or_disable_peer(&interface, &conf, false, opts.network)?
+        Command::DisablePeer { interface, args } => {
+            enable_or_disable_peer(&interface, &conf, false, opts.network, args)?
         },
-        Command::EnablePeer { interface } => {
-            enable_or_disable_peer(&interface, &conf, true, opts.network)?
+        Command::EnablePeer { interface, args } => {
+            enable_or_disable_peer(&interface, &conf, true, opts.network, args)?
         },
         Command::AddCidr { interface, args } => add_cidr(&interface, &conf, args)?,
         Command::DeleteCidr { interface, args } => delete_cidr(&interface, &conf, args)?,
@@ -384,14 +394,15 @@ fn enable_or_disable_peer(
     conf: &ServerConfig,
     enable: bool,
     network: NetworkOpts,
+    opts: EnableDisablePeerOpts,
 ) -> Result<(), Error> {
     let conn = open_database_connection(interface, conf)?;
     let peers = DatabasePeer::list(&conn)?
         .into_iter()
         .map(|dp| dp.inner)
         .collect::<Vec<_>>();
 
-    if let Some(peer) = prompts::enable_or_disable_peer(&peers[..], enable)? {
+    if let Some(peer) = prompts::enable_or_disable_peer(&peers[..], &opts, enable)? {
         let mut db_peer = DatabasePeer::get(&conn, peer.id)?;
         db_peer.update(
             &conn,

diff --git a/shared/src/prompts.rs b/shared/src/prompts.rs
@@ -1,8 +1,8 @@
 use crate::{
     interface_config::{InterfaceConfig, InterfaceInfo, ServerInfo},
     AddCidrOpts, AddDeleteAssociationOpts, AddPeerOpts, Association, Cidr, CidrContents, CidrTree,
-    DeleteCidrOpts, Endpoint, Error, Hostname, IpNetExt, ListenPortOpts, OverrideEndpointOpts,
-    Peer, PeerContents, RenamePeerOpts, PERSISTENT_KEEPALIVE_INTERVAL_SECS,
+    DeleteCidrOpts, EnableDisablePeerOpts, Endpoint, Error, Hostname, IpNetExt, ListenPortOpts,
+    OverrideEndpointOpts, Peer, PeerContents, RenamePeerOpts, PERSISTENT_KEEPALIVE_INTERVAL_SECS,
 };
 use anyhow::anyhow;
 use colored::*;
@@ -394,28 +394,41 @@ pub fn rename_peer(
 
 /// Presents a selection and confirmation of eligible peers for either disabling or enabling,
 /// and returns back the ID of the selected peer.
-pub fn enable_or_disable_peer(peers: &[Peer], enable: bool) -> Result<Option<Peer>, Error> {
+pub fn enable_or_disable_peer(
+    peers: &[Peer],
+    args: &EnableDisablePeerOpts,
+    enable: bool,
+) -> Result<Option<Peer>, Error> {
     let enabled_peers: Vec<_> = peers
         .iter()
         .filter(|peer| enable && peer.is_disabled || !enable && !peer.is_disabled)
         .collect();
 
-    let peer_selection: Vec<_> = enabled_peers
-        .iter()
-        .map(|peer| format!("{} ({})", &peer.name, &peer.ip))
-        .collect();
-    let (index, _) = select(
-        &format!("Peer to {}able", if enable { "en" } else { "dis" }),
-        &peer_selection,
-    )?;
-    let peer = enabled_peers[index];
+    let peer = if let Some(ref name) = args.name {
+        enabled_peers
+            .into_iter()
+            .find(|p| &p.name == name)
+            .ok_or_else(|| anyhow!("Peer '{}' does not exist", name))?
+    } else {
+        let peer_selection: Vec<_> = enabled_peers
+            .iter()
+            .map(|peer| format!("{} ({})", &peer.name, &peer.ip))
+            .collect();
+        let (index, _) = select(
+            &format!("Peer to {}able", if enable { "en" } else { "dis" }),
+            &peer_selection,
+        )?;
+        enabled_peers[index]
+    };
 
     Ok(
-        if confirm(&format!(
-            "{}able peer {}?",
-            if enable { "En" } else { "Dis" },
-            peer.name.yellow()
-        ))? {
+        if args.yes
+            || confirm(&format!(
+                "{}able peer {}?",
+                if enable { "En" } else { "Dis" },
+                peer.name.yellow()
+            ))?
+        {
             Some(peer.clone())
         } else {
             None

diff --git a/shared/src/types.rs b/shared/src/types.rs
@@ -351,6 +351,17 @@ pub struct RenamePeerOpts {
     pub yes: bool,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Args)]
+pub struct EnableDisablePeerOpts {
+    /// Name of peer to enable/disable
+    #[clap(long)]
+    pub name: Option<Hostname>,
+
+    /// Bypass confirmation
+    #[clap(long, requires("name"))]
+    pub yes: bool,
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Args)]
 pub struct AddCidrOpts {
     /// The CIDR name (eg. 'engineers')