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, /// 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, /// Path to cloud-init userdata, if any #[arg(long)] ci_userdata: Option, } #[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, #[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, /// Default DNS address for the subnet #[arg(long)] pub dns_server: Option, /// Start address for IP assignment #[arg(short, long)] pub start_addr: Option, /// End address for IP assignment #[arg(short, long)] pub end_addr: Option, #[arg(short, long)] pub domain_name: Option, /// VLAN ID for the VM, if any #[arg(short, long)] pub vlan_id: Option, } #[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, /// Default DNS address for the subnet #[arg(long)] pub dns_server: Option, /// Start address for IP assignment #[arg(short, long)] pub start_addr: Option, /// End address for IP assignment #[arg(short, long)] pub end_addr: Option, /// Domain name associated with the subnet #[arg(short, long)] pub domain_name: Option, /// VLAN ID for the VM, if any #[arg(short, long)] pub vlan_id: Option, } #[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 }, } #[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>, 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 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(message: S, inner: E) -> Self where S: AsRef, E: std::error::Error + 'static, { Self { message: message.as_ref().to_owned(), inner: Some(Box::new(inner)), } } } async fn handle_command() -> Result<(), Box> { 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 = { 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 = 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 = 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 = 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> { if let Err(err) = handle_command().await { if std::any::Any::type_id(&*err).type_id() == TypeId::of::() { log::error!("Error communicating with server: {}", err); } else { log::error!("{}", err); } } Ok(()) }