Add 'rm' and 'wg' commands

3 files changed, 93 insertions, 2 deletions

diff --git a/src/main.rs b/src/main.rs
index a0e38ca..05b05a2 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,8 +4,9 @@ use std::net::Ipv4Addr;
 use std::str::FromStr;
 use std::{env, process::exit};
 use std::io::Write;
+use std::process::Command as processCommand;
 
-use anyhow::{Result, anyhow};
+use anyhow::{Result, anyhow, Context};
 use clap::{command, ArgMatches};
 use clap::Command;
 use clap::arg;
@@ -59,6 +60,24 @@ fn main() {
                 .required(false)
             )
         )
+
+        // "rm"
+        .subcommand(
+            Command::new("rm")
+            .about("Remove a client from your VPN")
+            .arg(
+                arg!(
+                    <NAME> "Name of the client to remove"
+                )
+                .required(true)
+            )
+        )
+
+        // "wg"
+        .subcommand(
+            Command::new("wg")
+            .about("Run 'wg', but with the client names")
+        )
         .get_matches();
 
     let conf_path = match find_config_file(&matches) {
@@ -118,6 +137,23 @@ fn main() {
             else {
                 println!("Peer '{}' added successfully. Don't forget to restart your WireGuard server.", new_name);
             }
+        },
+
+        Some(("rm", args)) => {
+            let rm_name = args.get_one::<String>("NAME").unwrap().to_string();
+            if let Err(e) = do_rm(&mut conf, conf_path, &rm_name) {
+                eprintln!("Error removing peer: {}", e);
+                exit(1);
+            }
+
+            println!("Client '{}' removed successfully. Don't forget to restart your WireGuard server.", rm_name);
+        },
+
+        Some(("wg", _)) => {
+            if let Err(e) = do_wg(&conf) {
+                eprintln!("{}", e);
+                exit(1);
+            }
         }
 
         _ => {
@@ -238,3 +274,55 @@ fn do_add(conf: &mut WireguardConfig, conf_path: String, peer_name: &String, ip:
 
     Ok(())
 }
+
+fn do_rm(conf: &mut WireguardConfig, conf_path: String, peer_name: &String) -> Result<()> {
+    let mut del_index = 0;
+    let mut found = false;
+    let mut pk_path = String::new();
+
+    for (i, peer) in conf.peers.iter().enumerate() {
+        if &peer.name == peer_name {
+            del_index = i;
+            found = true;
+            pk_path = peer.private_key_path().context("could not get private key path")?;
+            break;
+        }
+    }
+
+    if !found {
+        return Err(anyhow!("No such peer: {}", peer_name));
+    }
+
+    conf.peers.remove(del_index);
+
+    let mut f = File::create(conf_path.clone()).context(format!("error opening configuration file: {}", conf_path))?;
+    let data = conf.gen_config().context("error generating configuration")?;
+
+    f.write_all(data.as_bytes()).context(format!("error writing to file: {}", conf_path))?;
+
+    std::fs::remove_file(pk_path).context(format!("could not remove file: {}", conf_path))?;
+    
+    Ok(())
+}
+
+fn do_wg(conf: &WireguardConfig) -> Result<()> {
+    let wg = processCommand::new("wg")
+        .env("WG_COLOR_MODE", "always")
+        .output()
+        .context("could not run 'wg', is it installed?")?;
+
+    if !wg.status.success() {
+        let err = String::from_utf8(wg.stderr)?;
+        return Err(anyhow!("error running 'wg': {}", err));
+    }   
+
+    let mut out = String::from_utf8(wg.stdout).context("error parsing stdout")?;
+
+    for peer in conf.peers.iter() {
+        out = out.replace(peer.public_key.as_str(), peer.name.as_str());
+    }
+
+    println!("{}", out);
+    
+    Ok(())
+}
diff --git a/src/wg/config.rs b/src/wg/config.rs
index 4e8c205..b354626 100644
--- a/src/wg/config.rs
+++ b/src/wg/config.rs
@@ -81,6 +81,9 @@ impl WireguardConfig {
                     let addr = config_value(line)?;
                     conf.network = Ipv4Network::from_str(addr)?;
                 }
+                else if line.starts_with("#") {
+                    // We ignore comments, because they are used for peer identification
+                }
                 else {
                     conf.other_lines.push(String::from(line)); 
                 }
diff --git a/src/wg/peer.rs b/src/wg/peer.rs
index 1d2002f..7b703fc 100644
--- a/src/wg/peer.rs
+++ b/src/wg/peer.rs
@@ -74,7 +74,7 @@ impl Peer {
         Ok(pk_folder)
     }
 
-    fn private_key_path(&self) -> Result<String> {
+    pub fn private_key_path(&self) -> Result<String> {
         let folder_name = self.private_key_folder()?;
         let pk_folder = Path::new(folder_name.as_str());
         let pk_path = pk_folder.join(self.name.clone());