path: root/src/main.rs

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::io::Write;
use std::process::Command as processCommand;

use anyhow::{Result, anyhow, Context};
use clap::{command, ArgMatches};
use clap::Command;
use clap::arg;
use wg::config::WireguardConfig;

mod wg;
mod configuration;

fn main() {
    let matches = command!()
        .arg(
            arg!(
                -w --wgconfig <WGCONFIG> "Wireguard configuration file name"
            )
            .required(false)
        )
        .arg(
            arg!(
                -c --config <CONFIG> "wgmgr configuration file path"
            )
            .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")
            .arg(
                arg!(
                    -t --type <TYPE> "Type of configuration: \"split\" or \"full\""
                )
                .required(false)
            )
            .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)
            )
        )

        // "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();
    
    
    // Find the configuration file for wgmgr
    let wgmgr_conf_path = match configuration::find_configuration_file(matches.get_one::<String>("config")) {
        Ok(s) => s,
        Err(e) => {
            eprintln!("Could not find wgmgr configuration file. Place it at /etc/wgmgr.toml or use the --config flag.");
            eprintln!("{:?}", e);
            exit(1);
        }
    };
    
    let wgmgr_conf: configuration::Configuration = toml::from_str(&read_to_string(wgmgr_conf_path).unwrap()).unwrap();

    // Find the Wireguard configuration file
    let wg_conf_path = match find_wg_config_file(&matches, &wgmgr_conf.wgconf) {
        Ok(s) => s,
        Err(e) => {
            eprintln!("Wireguard configuration file error: {:?}", e);
            exit(1);
        }
    };

    let mut wg_conf = match wg::config::WireguardConfig::new(&wg_conf_path) {
        Ok(c) => c,
        Err(e) => {
            eprintln!("Error loading the Wireguard configuration file '{}'.", wg_conf_path);
            eprintln!("{:?}", e);
            exit(1);
        }
    };

    match matches.subcommand() {
        Some(("ls", _)) => {
            do_list(&wg_conf);
        }

        Some(("config", args)) => {
            match args.get_one::<String>("PEER") {
                Some(peer_name) => {
                    let is_full = match args.get_one::<String>("type") {
                        Some(t) => {
                            match t.to_string().as_str() {
                                "split" => false,
                                "full" => true,
                                _ => {
                                    eprintln!("Error: '--type' must be either 'split' or 'full'");
                                    exit(3);
                                }
                            }
                        },
                        None => false
                    };

                    if let Err(e) = do_config(&wg_conf, &wgmgr_conf, peer_name.to_string(), is_full) {
                        eprintln!("Error generating configuration for {}: {:?}", peer_name.to_string(), e);
                        exit(1);
                    }
                },
                None => {}
            }
        },

        Some(("add", args)) => {
            let new_name = args.get_one::<String>("NAME").unwrap().to_string();
            if let Err(e) = do_add(&mut wg_conf, wg_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);
            }
        },

        Some(("rm", args)) => {
            let rm_name = args.get_one::<String>("NAME").unwrap().to_string();
            if let Err(e) = do_rm(&mut wg_conf, wg_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(&wg_conf) {
                eprintln!("{:?}", e);
                exit(1);
            }
        }

        _ => {
            unimplemented!();
        }
    }

}

fn find_wg_config_file(matches: &ArgMatches, wgconf: &Option<String>) -> Result<String> {
    // Top priority goes to the command-line argument
    match matches.get_one::<String>("wgconfig") {
        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(_) => {}
    };

    // Then, is there is a "wgconf" in the wgmgr configuration file, we take that
    match wgconf {
        Some(s) => return Ok(s.clone()),
        None => {}
    };

    // 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(wg_conf: &wg::config::WireguardConfig) {
    let mut max_length = 0;
    for p in wg_conf.peers.iter() {
        if p.name.len() > max_length {
            max_length = p.name.len();
        }
    }

    for p in wg_conf.peers.iter() {
        println!("{:max_length$} | {}", p.name, p.ip);
    }
}

fn do_config(wg_conf: &wg::config::WireguardConfig, conf: &configuration::Configuration, peer_name: String, is_full: bool) -> Result<()> {
    let peer = match wg_conf.get_peer(peer_name.as_str()) {
        Some(p) => p,
        None => {
            return Err(anyhow!("No such peer: {}", peer_name));
        }
    };

    println!("{}", peer.gen_config(wg_conf, conf.dns.clone(), conf.endpoint.clone(), conf.port, is_full)?);

    Ok(())
}

fn do_add(wg_conf: &mut WireguardConfig, wg_conf_path: String, peer_name: &String, ip: Option<&String>) -> Result<()> {
    let ip = match ip {
        Some(s) => {
            Ipv4Addr::from_str(s.as_str())?
        },
        None => {
            match wg_conf.next_free_ip() {
                Ok(i) => i,
                Err(e) => {
                    return Err(e);
                }
            }
        }
    };

    match wg_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)?;
    wg_conf.peers.push(p);

    let mut f = File::create(wg_conf_path)?;
    write!(f, "{}", wg_conf.gen_config()?)?;

    Ok(())
}

fn do_rm(wg_conf: &mut WireguardConfig, wg_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 wg_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));
    }

    wg_conf.peers.remove(del_index);

    let mut f = File::create(wg_conf_path.clone()).context(format!("error opening configuration file: {}", wg_conf_path))?;
    let data = wg_conf.gen_config().context("error generating configuration")?;

    f.write_all(data.as_bytes()).context(format!("error writing to file: {}", wg_conf_path))?;

    std::fs::remove_file(pk_path).context(format!("could not remove file: {}", wg_conf_path))?;
    
    Ok(())
}

fn do_wg(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 wg_conf.peers.iter() {
        out = out.replace(peer.public_key.as_str(), peer.name.as_str());
    }

    println!("{}", out);
    
    Ok(())
}