api: cloud-init userdata and ssh keys

This commit is contained in:
snow flurry 2024-08-14 17:33:59 -07:00
parent 997478801c
commit 957499c0a5
5 changed files with 109 additions and 35 deletions

View file

@ -1,13 +1,11 @@
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 nzr_api::{config, NazrinClient};
use std::any::{Any, TypeId};
use std::path::PathBuf;
use std::str::FromStr;
use tarpc::tokio_serde::formats::Bincode;
use tarpc::tokio_util::codec::LengthDelimitedCodec;
use tokio::net::UnixStream;
mod table;
@ -35,11 +33,11 @@ pub struct NewInstanceArgs {
#[arg(short, long, default_value_t = 20)]
primary_size: u32,
/// Secndary HDD size, in GiB
#[arg(short, long)]
secondary_size: Option<u32>,
/// File containing a list of SSH keys to use
#[arg(long)]
sshkey_file: Option<PathBuf>,
secondary_size: Option<u32>,
/// Path to cloud-init userdata, if any
#[arg(long)]
ci_userdata: Option<PathBuf>,
}
#[derive(Debug, Subcommand)]
@ -202,18 +200,12 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
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 framed_io = LengthDelimitedCodec::builder()
.length_field_type::<u32>()
.new_framed(conn);
let transport = tarpc::serde_transport::new(framed_io, Bincode::default());
let client = NazrinClient::new(Default::default(), transport).spawn();
let client = nzr_api::new_client(conn);
match cli.command {
Commands::Instance { command } => match command {
InstanceCmd::Dump { name, quick } => {
let instances = (client
.get_instances(tarpc::context::current(), !quick)
.await?)?;
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)?);
@ -223,6 +215,7 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
}
}
InstanceCmd::New(args) => {
/*
let ssh_keys: Vec<String> = {
let key_file = args.sshkey_file.map_or_else(
|| {
@ -254,6 +247,21 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
}
}
}?;
*/
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,
@ -264,10 +272,10 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
cores: args.cores,
memory: args.mem,
disk_sizes: (args.primary_size, args.secondary_size),
ssh_keys,
ci_userdata,
};
let task_id = (client
.new_instance(tarpc::context::current(), build_args)
.new_instance(nzr_api::default_ctx(), build_args)
.await?)?;
const MAX_RETRIES: i32 = 5;
@ -275,7 +283,7 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
let mut current_pct: f32 = 0.0;
loop {
let status = client
.poll_new_instance(tarpc::context::current(), task_id)
.poll_new_instance(nzr_api::default_ctx(), task_id)
.await;
match status {
Ok(Some(status)) => {
@ -315,21 +323,17 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
}
}
InstanceCmd::Delete { name } => {
(client
.delete_instance(tarpc::context::current(), name)
.await?)?;
(client.delete_instance(nzr_api::default_ctx(), name).await?)?;
}
InstanceCmd::List => {
let instances = client
.get_instances(tarpc::context::current(), true)
.await?;
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(tarpc::context::current()).await?)?,
InstanceCmd::Prune => (client.garbage_collect(nzr_api::default_ctx()).await?)?,
},
Commands::Net { command } => match command {
NetCmd::Add(args) => {
@ -350,12 +354,12 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
},
};
(client
.new_subnet(tarpc::context::current(), build_args)
.new_subnet(nzr_api::default_ctx(), build_args)
.await?)?;
}
NetCmd::Edit(args) => {
let mut net = client
.get_subnets(tarpc::context::current())
.get_subnets(nzr_api::default_ctx())
.await
.map_err(|e| e.to_string())
.and_then(|res| {
@ -391,7 +395,7 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
// run the update
client
.modify_subnet(tarpc::context::current(), net)
.modify_subnet(nzr_api::default_ctx(), net)
.await
.map_err(|err| format!("RPC error: {}", err))
.and_then(|res| {
@ -401,7 +405,7 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
})?;
}
NetCmd::Dump { name } => {
let subnets = (client.get_subnets(tarpc::context::current()).await?)?;
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)?);
@ -411,12 +415,10 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
}
}
NetCmd::Delete { name } => {
(client
.delete_subnet(tarpc::context::current(), name)
.await?)?;
(client.delete_subnet(nzr_api::default_ctx(), name).await?)?;
}
NetCmd::List => {
let subnets = client.get_subnets(tarpc::context::current()).await?;
let subnets = client.get_subnets(nzr_api::default_ctx()).await?;
let tabular: Vec<table::Subnet> =
subnets?.iter().map(table::Subnet::from).collect();
@ -431,7 +433,7 @@ async fn handle_command() -> Result<(), Box<dyn std::error::Error>> {
#[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::<tarpc::client::RpcError>() {
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);

View file

@ -13,7 +13,7 @@ pub struct NewInstance {
pub cores: u8,
pub memory: u32,
pub disk_sizes: (u32, Option<u32>),
pub ssh_keys: Vec<String>,
pub ci_userdata: Option<Vec<u8>>,
}
#[derive(Debug, Serialize, Deserialize)]

View file

@ -0,0 +1 @@
DROP TABLE ssh_keys;

View file

@ -0,0 +1,7 @@
CREATE TABLE ssh_keys (
id INTEGER PRIMARY KEY NOT NULL,
algorithm TEXT NOT NULL,
key_data TEXT NOT NULL,
comment TEXT,
UNIQUE(key_data)
);

View file

@ -54,6 +54,15 @@ diesel::table! {
}
}
diesel::table! {
ssh_keys {
id -> Integer,
algorithm -> Text,
key_data -> Text,
comment -> Nullable<Text>,
}
}
#[derive(
AsChangeset,
Clone,
@ -456,3 +465,58 @@ impl Transactable for Subnet {
self.delete(ctx).await
}
}
#[derive(Clone, Insertable, Identifiable, Selectable, Queryable)]
#[diesel(table_name = ssh_keys)]
pub struct SshPubkey {
pub id: i32,
pub algorithm: String,
pub key_data: String,
pub comment: Option<String>,
}
impl SshPubkey {
pub async fn all(ctx: &Context) -> Result<Vec<Self>, ModelError> {
let res = ctx
.spawn_db(move |mut db| {
Self::table()
.select(Self::as_select())
.load::<Self>(&mut db)
})
.await??;
Ok(res)
}
pub async fn insert(
ctx: &Context,
algorithm: impl AsRef<str>,
key_data: impl AsRef<str>,
comment: Option<impl AsRef<str>>,
) -> Result<Self, ModelError> {
use self::ssh_keys::columns;
let values = (
columns::algorithm.eq(algorithm.as_ref().to_owned()),
columns::key_data.eq(key_data.as_ref().to_owned()),
columns::comment.eq(comment.map(|s| s.as_ref().to_owned())),
);
let ent = ctx
.spawn_db(move |mut db| {
diesel::insert_into(Self::table())
.values(values)
.returning(ssh_keys::table::all_columns())
.get_result::<Self>(&mut db)
})
.await??;
Ok(ent)
}
pub async fn delete(self, ctx: &Context) -> Result<(), ModelError> {
ctx.spawn_db(move |mut db| diesel::delete(&self).execute(&mut db))
.await??;
Ok(())
}
}