nazrin/client/src/main.rs

use clap::{CommandFactory, FromArgMatches, Parser, Subcommand};
use nzr_api::config;
use nzr_api::hickory_proto::rr::Name;
use nzr_api::model;
use nzr_api::net::cidr::CidrV4;
use std::any::{Any, TypeId};
use std::path::PathBuf;
use std::str::FromStr;
use tokio::net::UnixStream;

mod table;

#[derive(Debug, clap::Args)]
pub struct NewInstanceArgs {
    /// Name of the instance to be created
    name: String,
    /// Subnet the instance will initially run on
    #[arg(short, long)]
    subnet: String,
    /// Long description of the instance
    #[arg(long)]
    description: Option<String>,
    /// Base image to use for the instance
    #[arg(short, long)]
    base: String,
    ///How many cores to assign to the instance
    #[arg(short, long, default_value_t = 2)]
    cores: u8,
    /// Memory to assign, in MiB
    #[arg(short, long, default_value_t = 1024)]
    mem: u32,
    /// Primary HDD size, in GiB
    #[arg(short, long, default_value_t = 20)]
    primary_size: u32,
    /// Secndary HDD size, in GiB
    #[arg(long)]
    secondary_size: Option<u32>,
    /// Path to cloud-init userdata, if any
    #[arg(long)]
    ci_userdata: Option<PathBuf>,
}

#[derive(Debug, Subcommand)]
enum InstanceCmd {
    /// Create a new instance
    New(NewInstanceArgs),
    /// Delete an existing instance
    Delete { name: String },
    /// List all instances
    List,
    /// Deletes all invalid instances from the database
    Prune,
    /// Shows information on an instance
    Dump {
        name: Option<String>,
        #[arg(short, long)]
        quick: bool,
    },
}

#[derive(Debug, clap::Args)]
pub struct AddNetArgs {
    /// Short name for the subnet (e.g., `servers')
    pub name: String,
    /// Name of the bridge interface VMs will attach to
    pub interface: String,
    /// Subnet associated with the bridge interface, in CIDR notation (x.x.x.x/y)
    pub network: String,
    /// Default gateway for the subnet
    #[arg(long)]
    pub gateway: Option<std::net::Ipv4Addr>,
    /// Default DNS address for the subnet
    #[arg(long)]
    pub dns_server: Option<std::net::Ipv4Addr>,
    /// Start address for IP assignment
    #[arg(short, long)]
    pub start_addr: Option<std::net::Ipv4Addr>,
    /// End address for IP assignment
    #[arg(short, long)]
    pub end_addr: Option<std::net::Ipv4Addr>,
    #[arg(short, long)]
    pub domain_name: Option<Name>,
    /// VLAN ID for the VM, if any
    #[arg(short, long)]
    pub vlan_id: Option<u32>,
}

#[derive(Debug, clap::Args)]
pub struct EditNetArgs {
    /// Short name of the subnet
    pub name: String,
    /// Default gateway for the subnet
    #[arg(long)]
    pub gateway: Option<std::net::Ipv4Addr>,
    /// Default DNS address for the subnet
    #[arg(long)]
    pub dns_server: Option<std::net::Ipv4Addr>,
    /// Start address for IP assignment
    #[arg(short, long)]
    pub start_addr: Option<std::net::Ipv4Addr>,
    /// End address for IP assignment
    #[arg(short, long)]
    pub end_addr: Option<std::net::Ipv4Addr>,
    /// Domain name associated with the subnet
    #[arg(short, long)]
    pub domain_name: Option<Name>,
    /// VLAN ID for the VM, if any
    #[arg(short, long)]
    pub vlan_id: Option<u32>,
}

#[derive(Debug, Subcommand)]
enum NetCmd {
    /// Add a new network to the database
    Add(AddNetArgs),
    /// Edit an existing network
    Edit(EditNetArgs),
    /// List all networks in the database
    List,
    /// Delete a network from the database
    Delete { name: String },
    /// Shows information on a subnet
    Dump { name: Option<String> },
}

#[derive(Debug, Subcommand)]
enum Commands {
    /// Commands for managing instances
    Instance {
        #[command(subcommand)]
        command: InstanceCmd,
    },
    /// Commands for managing network assignments
    Net {
        #[command(subcommand)]
        command: NetCmd,
    },
}

#[derive(Parser, Debug)]
#[command(version, about, long_about = "")]
struct Args {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Debug)]
struct CommandError {
    inner: Option<Box<dyn std::error::Error>>,
    message: String,
}

impl std::fmt::Display for CommandError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", &self.message)?;
        if let Some(inner) = &self.inner {
            write!(f, "\n  inner: {:?}", &inner)?;
        }
        Ok(())
    }
}

impl std::error::Error for CommandError {}

impl From<String> for CommandError {
    fn from(value: String) -> Self {
        Self {
            inner: None,
            message: value,
        }
    }
}

impl From<&str> for CommandError {
    fn from(value: &str) -> Self {
        Self {
            inner: None,
            message: value.to_owned(),
        }
    }
}

impl CommandError {
    fn new<S, E>(message: S, inner: E) -> Self
    where
        S: AsRef<str>,
        E: std::error::Error + 'static,
    {
        Self {
            message: message.as_ref().to_owned(),
            inner: Some(Box::new(inner)),
        }
    }
}

async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
    env_logger::init();

    let mut matches = Args::command().infer_subcommands(true).get_matches();
    let cli = Args::from_arg_matches_mut(&mut matches)?;
    let config: config::Config = nzr_api::config::Config::figment().extract()?;
    let conn = UnixStream::connect(&config.rpc.socket_path).await?;
    let client = nzr_api::new_client(conn);

    match cli.command {
        Commands::Instance { command } => match command {
            InstanceCmd::Dump { name, quick } => {
                let instances = (client.get_instances(nzr_api::default_ctx(), !quick).await?)?;
                if let Some(name) = name {
                    if let Some(inst) = instances.iter().find(|f| f.name == name) {
                        println!("{}", serde_json::to_string(inst)?);
                    }
                } else {
                    println!("{}", serde_json::to_string(&instances)?);
                }
            }
            InstanceCmd::New(args) => {
                /*
                let ssh_keys: Vec<String> = {
                    let key_file = args.sshkey_file.map_or_else(
                        || {
                            home::home_dir().map_or_else(
                                || {
                                    Err(CommandError::from(
                                        "SSH keyfile not defined, and couldn't find home directory",
                                    ))
                                },
                                |hd| Ok(hd.join(".ssh/authorized_keys")),
                            )
                        },
                        Ok,
                    )?;

                    if !key_file.exists() {
                        Err("SSH keyfile doesn't exist".into())
                    } else {
                        match std::fs::read_to_string(&key_file) {
                            Ok(data) => {
                                let keys: Vec<String> =
                                    data.split('\n').map(|s| s.trim().to_owned()).collect();
                                Ok(keys)
                            }
                            Err(err) => Err(CommandError::new(
                                format!("Couldn't read {} for SSH keys", &key_file.display()),
                                err,
                            )),
                        }
                    }
                }?;
                */
                let ci_userdata = {
                    if let Some(path) = &args.ci_userdata {
                        if !path.exists() {
                            return Err("cloud-init userdata file doesn't exist".into());
                        } else {
                            Some(
                                std::fs::read(path)
                                    .map_err(|e| format!("Couldn't read userdata file: {e}"))?,
                            )
                        }
                    } else {
                        None
                    }
                };

                let build_args = nzr_api::args::NewInstance {
                    name: args.name,
                    title: None,
                    description: args.description,
                    subnet: args.subnet,
                    base_image: args.base,
                    cores: args.cores,
                    memory: args.mem,
                    disk_sizes: (args.primary_size, args.secondary_size),
                    ci_userdata,
                };
                let task_id = (client
                    .new_instance(nzr_api::default_ctx(), build_args)
                    .await?)?;

                const MAX_RETRIES: i32 = 5;
                let mut retries = 0;
                let mut current_pct: f32 = 0.0;
                loop {
                    let status = client
                        .poll_new_instance(nzr_api::default_ctx(), task_id)
                        .await;
                    match status {
                        Ok(Some(status)) => {
                            if let Some(result) = status.result {
                                match result {
                                    Ok(instance) => {
                                        println!("Instance {} created!", &instance.name);
                                        println!(
                                            "You should be able to reach it with: ssh root@{}",
                                            instance.lease.addr.addr,
                                        );
                                    }
                                    Err(err) => {
                                        log::error!("Error while creating instance: {}", err);
                                    }
                                }
                                break;
                            } else if status.completion != current_pct {
                                println!("[remote] {}", &status.status_text);
                                current_pct = status.completion;
                            }
                        }
                        Ok(None) => {
                            log::error!("Task ID {} went AWOL??", task_id);
                            break;
                        }
                        Err(err) => {
                            log::error!("Got RPC error: {}", err);
                            retries += 1;
                            if retries >= MAX_RETRIES {
                                break;
                            } else {
                                log::error!("Retrying (attempt {}/{})...", retries, MAX_RETRIES);
                            }
                        }
                    }
                }
            }
            InstanceCmd::Delete { name } => {
                (client.delete_instance(nzr_api::default_ctx(), name).await?)?;
            }
            InstanceCmd::List => {
                let instances = client.get_instances(nzr_api::default_ctx(), true).await?;

                let tabular: Vec<table::Instance> =
                    instances?.iter().map(table::Instance::from).collect();
                let mut table = tabled::Table::new(tabular);
                println!("{}", table.with(tabled::settings::Style::psql()));
            }
            InstanceCmd::Prune => (client.garbage_collect(nzr_api::default_ctx()).await?)?,
        },
        Commands::Net { command } => match command {
            NetCmd::Add(args) => {
                let net_arg = CidrV4::from_str(&args.network)?;
                let build_args = model::Subnet {
                    name: args.name,
                    data: model::SubnetData {
                        ifname: args.interface.clone(),
                        network: net_arg,
                        start_host: args.start_addr.unwrap_or(net_arg.make_ip(10)?),
                        end_host: args
                            .end_addr
                            .unwrap_or((u32::from(net_arg.broadcast()) - 1u32).into()),
                        gateway4: Some(args.gateway.unwrap_or(net_arg.make_ip(1)?)),
                        dns: args.dns_server.map_or(Vec::new(), |d| vec![d]),
                        domain_name: args.domain_name,
                        vlan_id: args.vlan_id,
                    },
                };
                (client
                    .new_subnet(nzr_api::default_ctx(), build_args)
                    .await?)?;
            }
            NetCmd::Edit(args) => {
                let mut net = client
                    .get_subnets(nzr_api::default_ctx())
                    .await
                    .map_err(|e| e.to_string())
                    .and_then(|res| {
                        res?.iter()
                            .find_map(|ent| {
                                if ent.name == args.name {
                                    Some(ent.clone())
                                } else {
                                    None
                                }
                            })
                            .ok_or_else(|| format!("Couldn't find network {}", &args.name))
                    })?;

                // merge in the new args
                net.data.gateway4 = args.gateway;

                if let Some(dns_server) = args.dns_server {
                    net.data.dns = vec![dns_server]
                }
                if let Some(start_addr) = args.start_addr {
                    net.data.start_host = start_addr;
                }
                if let Some(end_addr) = args.end_addr {
                    net.data.end_host = end_addr;
                }
                if let Some(domain_name) = args.domain_name {
                    net.data.domain_name = Some(domain_name);
                }
                if let Some(vlan_id) = args.vlan_id {
                    net.data.vlan_id = Some(vlan_id);
                }

                // run the update
                client
                    .modify_subnet(nzr_api::default_ctx(), net)
                    .await
                    .map_err(|err| format!("RPC error: {}", err))
                    .and_then(|res| {
                        res.map(|e| {
                            println!("Subnet {} updated.", e.name);
                        })
                    })?;
            }
            NetCmd::Dump { name } => {
                let subnets = (client.get_subnets(nzr_api::default_ctx()).await?)?;
                if let Some(name) = name {
                    if let Some(net) = subnets.iter().find(|s| s.name == name) {
                        println!("{}", serde_json::to_string(net)?);
                    }
                } else {
                    println!("{}", serde_json::to_string(&subnets)?);
                }
            }
            NetCmd::Delete { name } => {
                (client.delete_subnet(nzr_api::default_ctx(), name).await?)?;
            }
            NetCmd::List => {
                let subnets = client.get_subnets(nzr_api::default_ctx()).await?;

                let tabular: Vec<table::Subnet> =
                    subnets?.iter().map(table::Subnet::from).collect();
                let mut table = tabled::Table::new(tabular);
                println!("{}", table.with(tabled::settings::Style::psql()));
            }
        },
    };
    Ok(())
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    if let Err(err) = handle_command().await {
        if std::any::Any::type_id(&*err).type_id() == TypeId::of::<nzr_api::RpcError>() {
            log::error!("Error communicating with server: {}", err);
        } else {
            log::error!("{}", err);
        }
    }
    Ok(())
}