azirevpn/src/main.rs

mod api;
use std::net::IpAddr;
mod dirs;
mod keys;

use std::io::Write;
use std::path::PathBuf;

use clap::{Parser, Subcommand};
use log::debug;

use crate::keys::{get_keys, WireguardKeyPair};

/// AzireVPN client
#[derive(Parser, Debug)]
#[command(version)]
struct Opts {
    /// Enables JSON output
    #[arg(short, long)]
    json: bool,

    #[arg(short, long)]
    machine: Option<PathBuf>,

    #[command(subcommand)]
    command: Command,
}

#[derive(Parser, Debug)]
struct ConfigOpts {
    location: String,
    username: String,
    token: String,

    #[arg(short, long)]
    no_dns: bool,

    #[arg(short = '4', long)]
    no_ipv6: bool,
}

#[derive(Subcommand, Debug)]
enum Command {
    /// Checks connection status
    Check,

    /// Prints the list of VPN endpoints
    Locations,

    /// Prints WireGuard config
    Config(ConfigOpts),
}

fn main() -> Result<(), anyhow::Error> {
    env_logger::init();
    let opts = Opts::parse();
    match &opts.command {
        Command::Check => check(&opts)?,
        Command::Locations => list_locations(&opts)?,
        Command::Config(get_config_opts) => get_config(&opts, get_config_opts)?,
    }
    Ok(())
}

fn check(opts: &Opts) -> Result<(), anyhow::Error> {
    let check_result = api::check()?;
    if opts.json {
        println!("{}", serde_json::to_string(&check_result)?);
    } else {
        println!("Connected: {}", check_result.connected);
        println!("Ip: {}", check_result.ip);
        if let Some(location) = check_result.location {
            println!("Location: {}", location);
        }
    }
    Ok(())
}

fn list_locations(opts: &Opts) -> Result<(), anyhow::Error> {
    let locations = api::get_locations()?;
    if opts.json {
        println!("{}", serde_json::to_string(&locations)?);
    } else {
        for (i, location) in locations.locations.iter().enumerate() {
            if i > 0 {
                println!()
            };
            println!("Name: {}", location.name);
            println!("City: {}", location.city);
            println!("Country: {}", location.country);
            println!("Country code: {}", location.iso);
            println!("Pool: {}", location.pool);
            println!("Pubkey: {}", location.pubkey);
        }
    }
    Ok(())
}

fn get_config(opts: &Opts, config_opts: &ConfigOpts) -> Result<(), anyhow::Error> {
    let locations = api::get_locations()?;
    let location = locations
        .locations
        .iter()
        .find(|location| location.name == config_opts.location)
        .ok_or_else(|| anyhow::anyhow!("no such location: {}", config_opts.location))?;
    debug!("location = {:?}", &location);
    let keys = get_keys(opts.machine.as_ref())?;
    debug!("keys = {:?}", &keys);
    let addresses = api::add_ip(&config_opts.username, &config_opts.token, &keys.public_key)?;

    write_config(
        &mut std::io::stdout().lock(),
        config_opts,
        location,
        &addresses,
        &keys,
    )
}

fn write_config(
    output: &mut dyn Write,
    config_opts: &ConfigOpts,
    location: &api::Location,
    config: &api::Addresses,
    keys: &WireguardKeyPair,
) -> Result<(), anyhow::Error> {
    writeln!(output, "[Interface]")?;
    writeln!(output, "PrivateKey = {}", &keys.private_key.to_base64())?;
    let allowed_addresses = if config_opts.no_ipv6 {
        vec![IpAddr::V4(config.ipv4.address)]
    } else {
        vec![
            IpAddr::V4(config.ipv4.address),
            IpAddr::V6(config.ipv6.address),
        ]
    };
    write_list(output, "Address = ", allowed_addresses)?;
    if !config_opts.no_dns {
        let dns_addrs = &config.dns;
        let allowed_dns_addrs = dns_addrs
            .iter()
            .filter(|addr| addr.is_ipv4() || !config_opts.no_ipv6);
        write_list(output, "DNS = ", allowed_dns_addrs)?;
    }
    writeln!(output)?;

    writeln!(output, "[Peer]")?;
    writeln!(output, "PublicKey = {}", &location.pubkey)?;
    writeln!(output, "Endpoint = {}", &location.get_endpoint()?)?;

    let allowed_ips: &[&str] = if config_opts.no_ipv6 {
        &["0.0.0.0/0"]
    } else {
        &["0.0.0.0/0", "::/0"]
    };
    write_list(output, "AllowedIPs = ", allowed_ips)?;

    Ok(())
}

fn write_list<I, T>(output: &mut dyn Write, prefix: &str, values: I) -> Result<(), std::io::Error>
where
    I: IntoIterator<Item = T>,
    T: std::fmt::Display,
{
    write!(output, "{}", prefix)?;
    for (i, value) in values.into_iter().enumerate() {
        if i != 0 {
            write!(output, ", ")?;
        }
        write!(output, "{}", value)?;
    }
    writeln!(output)?;
    Ok(())
}