Use anyhow, +add +config

8 files changed, 297 insertions, 68 deletions

diff --git a/.gitignore b/.gitignore
index ea8c4bf..55e5a90 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,3 @@
 /target
+private_keys
+*.conf
\ No newline at end of file
diff --git a/Cargo.lock b/Cargo.lock
index 7e505bb..6c862da 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -52,6 +52,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "anyhow"
+version = "1.0.70"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7de8ce5e0f9f8d88245311066a578d72b7af3e7088f32783804676302df237e4"
+
+[[package]]
 name = "bitflags"
 version = "1.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -211,6 +217,7 @@ checksum = "711b9620af191e0cdc7468a8d14e709c3dcdb115b36f838e601583af800a370a"
 name = "wgmgr"
 version = "0.1.0"
 dependencies = [
+ "anyhow",
  "clap",
  "ipnetwork",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index 48f40cd..8b25ae0 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -6,5 +6,6 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+anyhow = "1.0.70"
 clap = { version = "4.2.2", features = ["cargo"] }
 ipnetwork = "0.20.0"
diff --git a/src/main.rs b/src/main.rs
index 6c4787e..4fdfddd 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,19 +1,35 @@
+use std::ffi::OsStr;
+use std::fs::{read_to_string, self, File};
+use std::net::Ipv4Addr;
+use std::str::FromStr;
 use std::{env, process::exit};
-use std::error::Error;
+use std::io::Write;
 
-use clap::command;
+use anyhow::{Result, anyhow};
+use clap::{command, ArgMatches};
 use clap::Command;
 use clap::arg;
+use wg::config::WireguardConfig;
 
 mod wg;
 
 fn main() {
     let matches = command!()
+        .arg(
+            arg!(
+                -c --config <CONFIG> "Configuration file name"
+            )
+            .required(false)
+        )
         .subcommand_required(true)
+
+        // "ls"
         .subcommand(
             Command::new("ls")
             .about("List known clients")
         )
+
+        // "config"
         .subcommand(
             Command::new("config")
             .about("Generate the configuration file for a client")
@@ -25,14 +41,35 @@ fn main() {
             )
             .arg(arg!(<PEER> "Name of the peer"))
         )
+
+        // "add"
+        .subcommand(
+            Command::new("add")
+            .about("Add a new client to your VPN")
+            .arg(
+                arg!(
+                    <NAME> "Name of the new client"
+                )
+                .required(true)
+            )
+            .arg(
+                arg!(
+                    -i --ip <IP> "IP address (auto-assigned)"
+                )
+                .required(false)
+            )
+        )
         .get_matches();
 
-    let conf_path = match env::var("WG_CONF") {
+    let conf_path = match find_config_file(&matches) {
         Ok(s) => s,
-        Err(_) => String::from("/etc/wireguard/wg0.conf")
+        Err(e) => {
+            eprintln!("Configuration file error: {}", e);
+            exit(1);
+        }
     };
 
-    let conf = match wg::config::WireguardConfig::new(&conf_path) {
+    let mut conf = match wg::config::WireguardConfig::new(&conf_path) {
         Ok(c) => c,
         Err(e) => {
             eprintln!("Error loading the configuration file '{}'", conf_path);
@@ -67,6 +104,17 @@ fn main() {
                 },
                 None => {}
             }
+        },
+
+        Some(("add", args)) => {
+            let new_name = args.get_one::<String>("NAME").unwrap().to_string();
+            if let Err(e) = do_add(&mut conf, conf_path, &new_name, args.get_one::<String>("ip")) {
+                eprintln!("Error adding peer: {}", e);
+                exit(1);
+            }
+            else {
+                println!("Peer '{}' added successfully. Don't forget to restart your WireGuard server.", new_name);
+            }
         }
 
         _ => {
@@ -76,6 +124,63 @@ fn main() {
 
 }
 
+fn find_config_file(matches: &ArgMatches) -> Result<String> {
+    // Top priority goes to the command-line argument
+    match matches.get_one::<String>("config") {
+        Some(s) => {
+            if s.starts_with("/") {
+                return Ok(s.to_string())
+            }
+            else {
+                return Ok(format!("/etc/wireguard/{}", s))
+            }
+        },
+        None => {}
+    }
+
+    // Then, if the environment variable exists, we take it
+    match env::var("WG_CONF") {
+        Ok(s) => return Ok(s),
+        Err(_) => {}
+    };
+
+    // Otherwise, we will first try a simple "wg0.conf"
+    match read_to_string("/etc/wireguard/wg0.conf") {
+        Ok(_) => return Ok(String::from_str("/etc/wireguard/wg0.conf").unwrap()),
+        Err(_) => {}
+    }
+
+    // Finally, we can try to see if there is a single ".conf" file in /etc/wireguard
+    match fs::read_dir("/etc/wireguard/") {
+        Ok(d) => {
+            let conf_files: Vec<fs::DirEntry> = d.filter(|e| {
+                if let Ok(e) = e {
+                    if e.file_type().unwrap().is_dir() {
+                        return false;
+                    }
+    
+                    match e.path().extension().and_then(OsStr::to_str) {
+                        Some("conf") => { return true} ,
+                        _ => {return false},
+                    }
+                }
+                
+                return false;
+            }).map(|e| {
+                return e.unwrap();
+            }).collect();
+
+            if conf_files.len() == 1 {
+                return Ok(String::from_str(conf_files[0].path().to_str().unwrap()).unwrap());
+            }
+            else {
+                return Err(anyhow!("Could not determine the path to your WireGuard configuration file. Set the WG_CONF environment variable, or pass the '--config' parameter."))
+            }
+        },
+        Err(e) => return Err(anyhow!("Error listing /etc/wireguard/: {}", e))
+    }
+}
+
 fn do_list(conf: &wg::config::WireguardConfig) {
     let mut max_length = 0;
     for p in conf.peers.iter() {
@@ -89,30 +194,44 @@ fn do_list(conf: &wg::config::WireguardConfig) {
     }
 }
 
-fn do_config(conf: &wg::config::WireguardConfig, peer_name: String, is_full: bool) -> Result<(), Box<dyn Error>> {
-    let peer = match conf.peers.iter().filter(|p| { p.name == peer_name }).nth(0) {
+fn do_config(conf: &wg::config::WireguardConfig, peer_name: String, is_full: bool) -> Result<()> {
+    let peer = match conf.get_peer(peer_name.as_str()) {
         Some(p) => p,
         None => {
-            eprintln!("No such peer: {}", peer_name);
-            exit(1);
+            return Err(anyhow!("No such peer: {}", peer_name));
         }
     };
 
-    println!("[Interface]");
-    println!("PrivateKey = {}", peer.private_key()?);
-    println!("Address = {}/32", peer.ip);
-    println!("DNS = TODO\n");
+    println!("{}", peer.gen_config(conf, is_full)?);
 
-    println!("[Peer]");
-    println!("PublicKey = TODO");
+    Ok(())
+}
 
-    let allowed_ips = match is_full { 
-        true => String::from("0.0.0.0/0"), 
-        false => conf.network.to_string()
+fn do_add(conf: &mut WireguardConfig, conf_path: String, peer_name: &String, ip: Option<&String>) -> Result<()> {
+    let ip = match ip {
+        Some(s) => {
+            Ipv4Addr::from_str(s.as_str())?
+        },
+        None => {
+            match conf.next_free_ip() {
+                Ok(i) => i,
+                Err(e) => {
+                    return Err(e);
+                }
+            }
+        }
     };
-    println!("AllowedIPs = {}", allowed_ips);
-    println!("Endpoint = TODO");
-    println!("PersistentKeepAlive = 25");
+
+    match conf.get_peer(peer_name.as_str()) {
+        Some(_) => { return Err(anyhow!("There is already a peer named {}", peer_name)); },
+        None => {}
+    }
+
+    let p = wg::peer::Peer::new(peer_name.clone(), ip);
+    conf.peers.push(p);
+
+    let mut f = File::create(conf_path)?;
+    write!(f, "{}", conf.gen_config()?)?;
 
     Ok(())
 }
\ No newline at end of file
diff --git a/src/wg.rs b/src/wg.rs
index 39cfc0c..a7d9b85 100644
--- a/src/wg.rs
+++ b/src/wg.rs
@@ -1,3 +1,2 @@
-mod error;
 pub mod config;
-mod peer;
\ No newline at end of file
+pub mod peer;
\ No newline at end of file
diff --git a/src/wg/config.rs b/src/wg/config.rs
index 03d2ce0..4a07287 100644
--- a/src/wg/config.rs
+++ b/src/wg/config.rs
@@ -1,9 +1,10 @@
 use std::str::FromStr;
-use std::{net::Ipv4Addr, fs, error::Error};
+use std::{net::Ipv4Addr, fs};
+use std::fmt::Write;
 
+use anyhow::{Result, anyhow};
 use ipnetwork::Ipv4Network;
 
-use crate::wg::error::WgMgrError;
 use crate::wg::peer::Peer;
 
 #[derive(Debug)]
@@ -19,7 +20,7 @@ pub struct WireguardConfig {
 
 
 impl WireguardConfig {
-    pub fn new(config_path: &str) -> Result<WireguardConfig, Box<dyn Error>> {
+    pub fn new(config_path: &str) -> Result<WireguardConfig> {
         let f = fs::read_to_string(config_path)?;
 
         let mut conf = WireguardConfig{
@@ -96,8 +97,7 @@ impl WireguardConfig {
                     let ip = Ipv4Network::from_str(addr)?;
 
                     if ip.prefix() != 32 {
-                        let msg = format!("Peer '{}' has invalid net range {}", current_peer.name, ip.prefix());
-                        return Err(Box::new(WgMgrError::new(msg)));
+                        return Err(anyhow!("Peer '{}' has invalid net range {}", current_peer.name, ip.prefix()));
                     }
 
                     current_peer.ip = ip.ip();
@@ -118,7 +118,7 @@ impl WireguardConfig {
     }
 
 
-    pub fn next_free_ip(&self) -> Result<Ipv4Addr, WgMgrError> {
+    pub fn next_free_ip(&self) -> Result<Ipv4Addr> {
         let mut iter = self.network.iter();
         iter.next(); // Skip the first IP (identification)
 
@@ -140,19 +140,57 @@ impl WireguardConfig {
             }
         }
 
-        Err(WgMgrError::new(String::from("No more free IP addresses")))
+        Err(anyhow!("No more free IP addresses"))
+    }
+
+    pub fn get_peer(&self, name: &str) -> Option<&Peer> {
+        self.peers.iter().filter(|p| { p.name.as_str() == name }).nth(0)
+    }
+
+    pub fn gen_config(&self) -> Result<String, std::fmt::Error> {
+        let mut res = String::new();
+
+        writeln!(res, "[Interface]")?;
+        writeln!(res, "PrivateKey = {}", self.private_key)?;
+
+        if self.listen_port != 0 {
+            writeln!(res, "ListenPort = {}", self.listen_port)?;
+        }
+
+        writeln!(res, "Address = {}", self.network)?;
+        
+        for pre_up in self.pre_ups.iter() {
+            writeln!(res, "PreUp = {}", pre_up)?;
+        }
+
+        for post_down in self.post_downs.iter() {
+            writeln!(res, "PostDown = {}", post_down)?;
+        }
+
+        writeln!(res, "")?;
+
+        // Now, add each peer
+        for peer in self.peers.iter() {
+            writeln!(res, "# {}", peer.name)?;
+            writeln!(res, "[Peer]")?;
+            writeln!(res, "PublicKey = {}", peer.public_key)?;
+            writeln!(res, "AllowedIPs = {}/32", peer.ip)?;
+            writeln!(res, "")?;
+        }
+
+        Ok(res)
     }
 }
 
 
 
-fn config_value(line: &str) -> Result<&str, &str> {
+fn config_value(line: &str) -> Result<&str> {
     match line.find("=") {
         Some(i) => {
             Ok(line[i+1..].trim())
         },
         None => {
-            Err("line does not seem to contain a key-value pair")
+            Err(anyhow!("line does not seem to contain a key-value pair"))
         }
     }
 }
\ No newline at end of file
diff --git a/src/wg/error.rs b/src/wg/error.rs
deleted file mode 100644
index 0167b22..0000000
--- a/src/wg/error.rs
+++ /dev/null
@@ -1,31 +0,0 @@
-use std::fmt::{Formatter, Display};
-use std::error::Error;
-
-#[derive(Debug)]
-pub struct WgMgrError {
-    msg: String
-}
-
-impl WgMgrError {
-    pub fn new(msg: String) -> WgMgrError {
-        WgMgrError { 
-            msg
-        }
-    }
-}
-
-impl Error for WgMgrError {
-    fn description(&self) -> &str {
-        return self.msg.as_str()
-    }
-
-    fn cause(&self) -> Option<&dyn Error> {
-        None
-    }
-}
-
-impl Display for WgMgrError {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-            write!(f, "{}", self.msg)
-    }
-}
\ No newline at end of file
diff --git a/src/wg/peer.rs b/src/wg/peer.rs
index 0a25841..17722e5 100644
--- a/src/wg/peer.rs
+++ b/src/wg/peer.rs
@@ -1,7 +1,14 @@
-use std::error::Error;
 use std::net::Ipv4Addr;
 use std::env::join_paths;
-use std::fs::read_to_string;
+use std::fs::{read_to_string, self, File};
+use std::fmt::Write;
+use std::path::{Path};
+use std::process::{Command, Stdio};
+use std::io::Write as ioWrite;
+
+use anyhow::Result;
+
+use super::config::WireguardConfig;
 
 #[derive(Debug)]
 pub struct Peer {
@@ -11,16 +18,103 @@ pub struct Peer {
 }
 
 impl Peer {
-    pub fn private_key(&self) -> Result<String, Box<dyn Error>> {
-        // TODO: do not hardcode where private keys are stored
+    pub fn new(name: String, ip: Ipv4Addr) -> Peer {
+        let private = Command::new("wg")
+            .arg("genkey")
+            .output()
+            .expect("Could not generate a private key. Is 'wg' installed?");
+
+        let mut public = Command::new("wg")
+            .arg("pubkey")
+            .stdin(Stdio::piped())
+            .stdout(Stdio::piped())
+            .spawn()
+            .expect("Could not run 'wg'");
+
+        let private_clone = private.stdout.clone();
+
+        let mut stdin = public.stdin.take().expect("could not open wg stdin");
+        std::thread::spawn(move || {
+            stdin.write_all(&private.stdout).unwrap();
+        });
+
+        let output = public.wait_with_output().expect("could not read from 'wg pubkey'");
+    
+        let private_key = String::from_utf8(private_clone).expect("could not decode private key");
+        let public_key = String::from_utf8(output.stdout).expect("could not decode public key");
+
+        let private_key = private_key.trim();
+        let public_key = public_key.trim();
 
+        let p = Peer{
+            name,
+            ip,
+            public_key: String::from(public_key)
+        };
+
+        let mut f = match File::create(p.private_key_path()) {
+            Ok(f) => f,
+            Err(e) => {
+                panic!("Error creating private key file: {}", e);
+            }
+        };
+        write!(f, "{}", private_key).expect("could not write private key to file");
+
+        p
+    }
+
+    fn private_key_folder(&self) -> String {
         let pk_folder = String::from("./private_keys/");
-        let pk_path = join_paths(&[pk_folder, self.name.clone()])?;
 
+        let pk_path = Path::new(pk_folder.as_str());
+        if !pk_path.exists() {
+            fs::create_dir(pk_path).expect("Could not create private keys folder");
+        }
+
+        if !pk_path.is_dir() {
+            panic!("Error: the private key folder exists but is not a directory")
+        }
+
+        pk_folder
+    }
+
+    fn private_key_path(&self) -> 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());
+
+        String::from(pk_path.to_str().unwrap())
+    }
+
+    pub fn private_key(&self) -> Result<String> {
+        let pk_folder = self.private_key_folder();
+        let pk_path = join_paths(&[pk_folder, self.name.clone()])?;
         let pk = read_to_string(pk_path)?;
 
         Ok(pk)
     }
+
+    pub fn gen_config(&self, conf: &WireguardConfig, is_full: bool) -> Result<String>{
+        let mut res = String::new();
+        
+        writeln!(res, "[Interface]")?;
+        writeln!(res, "PrivateKey = {}", self.private_key()?)?;
+        writeln!(res, "Address = {}/32", self.ip)?;
+        writeln!(res, "DNS = TODO\n")?;
+
+        writeln!(res, "[Peer]")?;
+        writeln!(res, "PublicKey = TODO")?;
+
+        let allowed_ips = match is_full { 
+            true => String::from("0.0.0.0/0"), 
+            false => conf.network.to_string()
+        };
+        writeln!(res, "AllowedIPs = {}", allowed_ips)?;
+        writeln!(res, "Endpoint = TODO")?;
+        writeln!(res, "PersistentKeepAlive = 25")?;
+
+        Ok(res)
+    }
 }
 
 impl PartialOrd for Peer {