azirevpn/src/main.rs

use std::io::Write;
use std::net::{AddrParseError, IpAddr, ToSocketAddrs};
use std::process;

use log::debug;

use clap::Clap;
use serde::{Deserialize, Serialize};

const BASE_URL: &str = "https://api.azirevpn.com/v1";

/// AzireVPN client
#[derive(Clap, Debug)]
#[clap(version=clap::crate_version!())]
struct Opts {
    /// Enables JSON output
    #[clap(short, long)]
    json: bool,

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

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

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

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

#[derive(Clap, Debug)]
enum Command {
    /// Prints the list of VPN endpoints
    Endpoints,

    /// Prints WireGuard config
    Config(ConfigOpts),

    /// Checks connection status
    Check,
}

#[derive(Serialize, Deserialize, Debug)]
struct Endpoints {
    wireguard: String,
}

#[derive(Serialize, Deserialize, Debug)]
struct Location {
    name: String,
    city: String,
    country: String,
    iso: String,
    endpoints: Endpoints,
}

#[derive(Serialize, Deserialize, Debug)]
struct Locations {
    locations: Vec<Location>,
}

#[derive(Serialize, Deserialize, Debug)]
struct CheckResult {
    connected: bool,
    ip: String, // XXX
    location: Option<String>,
}

#[derive(Serialize, Deserialize, Debug)]
struct WireguardConfig {
    status: String,
    data: WireguardConfigData,
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(rename_all = "PascalCase")]
struct WireguardConfigData {
    public_key: String,
    address: String,
    endpoint: String,

    #[serde(rename = "DNS")]
    dns: String,
}

#[derive(Debug)]
struct WireguardKeyPair {
    public_key: String,
    private_key: String,
}

impl WireguardConfigData {
    fn addresses(&self) -> Result<Vec<ipnet::IpNet>, ipnet::AddrParseError> {
        self.address
            .split(',')
            .map(|s: &str| -> Result<ipnet::IpNet, ipnet::AddrParseError> { s.trim().parse() })
            .collect()
    }

    fn dns(&self) -> Result<Vec<IpAddr>, AddrParseError> {
        self.dns
            .split(',')
            .map(|s: &str| -> Result<IpAddr, AddrParseError> { s.trim().parse() })
            .collect()
    }
}

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

fn list_endpoints(opts: &Opts) -> Result<(), anyhow::Error> {
    let locations: Locations = 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!("WireGuard endpoint: {}", location.endpoints.wireguard);
        }
    }
    Ok(())
}

fn get_config(_opts: &Opts, config_opts: &ConfigOpts) -> Result<(), anyhow::Error> {
    let locations = 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 = generage_keys()?;
    debug!("keys = {:?}", &keys);
    let config: WireguardConfig = ureq::post(&location.endpoints.wireguard)
        .send_form(&[
            ("username", &config_opts.username),
            ("token", &config_opts.token),
            ("pubkey", &keys.public_key),
        ])?
        .into_json()?;
    debug!("config = {:?}", &config);

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

fn write_config(
    output: &mut dyn Write,
    config_opts: &ConfigOpts,
    config: &WireguardConfig,
    keys: &WireguardKeyPair,
) -> Result<(), anyhow::Error> {
    writeln!(output, "[Interface]")?;
    writeln!(output, "PrivateKey = {}", &keys.private_key)?;
    let addresses = config.data.addresses()?;
    let allowed_addresses = addresses
        .iter()
        .filter(|addr| addr.addr().is_ipv4() || !config_opts.no_ipv6);
    write_list(output, "Address = ", allowed_addresses)?;
    if !config_opts.no_dns {
        let dns_addrs = config.data.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 = {}", &config.data.public_key)?;

    let endpoint_addrs = config.data.endpoint.to_socket_addrs()?;
    write_list(output, "Endpoint = ", endpoint_addrs)?;

    let allowed_ips: &[&str] = if config_opts.no_ipv6 {
        &["0.0.0.0"]
    } else {
        &["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(())
}

fn check(opts: &Opts) -> Result<(), anyhow::Error> {
    let url = format!("{}/check", BASE_URL);
    let result: CheckResult = ureq::get(&url).call()?.into_json()?;
    if opts.json {
        println!("{}", serde_json::to_string(&result)?);
    } else {
        println!("Connected: {:?}", result.connected);
        println!("IP: {}", result.ip);
        if let Some(location) = result.location {
            println!("Location: {}", location);
        }
    }
    Ok(())
}

fn get_locations() -> Result<Locations, anyhow::Error> {
    let url = format!("{}/locations", BASE_URL);
    let locations: Locations = ureq::get(&url).call()?.into_json()?;
    Ok(locations)
}

fn generage_keys() -> Result<WireguardKeyPair, anyhow::Error> {
    let privkey = process::Command::new("wg").arg("genkey").output()?.stdout;
    let mut pubkey_cmd = process::Command::new("wg")
        .arg("pubkey")
        .stdin(process::Stdio::piped())
        .stdout(process::Stdio::piped())
        .spawn()?;
    pubkey_cmd
        .stdin
        .as_mut()
        .expect("no stdin")
        .write_all(&privkey)?;
    let pubkey = pubkey_cmd.wait_with_output()?.stdout;
    Ok(WireguardKeyPair {
        private_key: String::from_utf8(privkey)?.trim_end().to_string(),
        public_key: String::from_utf8(pubkey)?.trim_end().to_string(),
    })
}