path: root/src/main.rs

use std::fs::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::{arg, command, Command};
use wg::config::WireguardConfig;

mod wg;
mod configuration;

fn main() {
    let matches = command!()
        .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 = match configuration::find_configuration_file(matches.get_one::<String>("config")) {
        Ok(c) => c,
        Err(e) => {
            eprintln!("Could not find wgmgr configuration file. Place it at /etc/wgmgr.toml or use the --config flag.");
            eprintln!("{:?}", e);
            exit(1);
        }
    };
    

    // Load the Wireguard configuration file
    let mut wg_conf = match wg::config::WireguardConfig::new(&wgmgr_conf.wgconf)  {
        Ok(c) => c,
        Err(e) => {
            eprintln!("Error loading the Wireguard configuration file '{}'.", wgmgr_conf.wgconf);
            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, wgmgr_conf.wgconf, &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, wgmgr_conf.wgconf, &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 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(())
}