use bytes::Bytes; use clap::{CommandFactory, Parser, Subcommand}; use eyre::format_err; use futures_util::Stream; use futures_util::StreamExt; use std::time::{Duration, SystemTime}; use tokio::fs; use tokio::io::{AsyncWrite, AsyncWriteExt}; use dkl::dls; #[derive(Parser)] #[command()] struct Cli { #[arg(long, default_value = "http://[::1]:7606", env = "DLS_URL")] dls: String, #[command(subcommand)] command: Command, } #[derive(Subcommand)] enum Command { Clusters, Cluster { cluster: String, #[command(subcommand)] command: Option, }, Hosts, Host { #[arg(short = 'o', long)] out: Option, host: String, asset: Option, }, #[command(subcommand)] DlSet(DlSet), /// hash a password Hash { salt: String, }, } #[derive(Subcommand)] enum DlSet { Sign { #[arg(short = 'e', long, default_value = "1d")] expiry: String, #[arg(value_parser = parse_download_set_item)] items: Vec, }, Show { #[arg(env = "DLS_DLSET")] signed_set: String, }, Fetch { #[arg(long, env = "DLS_DLSET")] signed_set: String, #[arg(short = 'o', long)] out: Option, kind: String, name: String, asset: String, }, } #[derive(Subcommand)] enum ClusterCommand { CaCert { #[arg(default_value = "cluster")] name: String, }, Token { #[arg(default_value = "admin")] name: String, }, Addons, SshSign { user_public_key: String, #[arg(long, default_value = "root")] principal: String, #[arg(long, default_value = "1d")] validity: String, #[arg(long)] options: Vec, }, KubeSign { csr: String, #[arg(long, default_value = "anonymous", env = "USER")] user: String, #[arg(long)] group: Option, #[arg(long, default_value = "1d")] validity: String, }, } #[tokio::main(flavor = "current_thread")] async fn main() -> eyre::Result<()> { clap_complete::CompleteEnv::with_factory(Cli::command).complete(); let cli = Cli::parse(); env_logger::builder() .parse_filters("info") .parse_default_env() .init(); let dls = || { let token = std::env::var("DLS_TOKEN").expect("DLS_TOKEN should be set"); dls::Client::new(cli.dls, token) }; use Command as C; match cli.command { C::Clusters => write_json(&dls().clusters().await?), C::Cluster { cluster, command } => { let dls = dls(); let cluster = dls.cluster(cluster); use ClusterCommand as CC; match command { None => write_json(&cluster.config().await?), Some(CC::CaCert { name }) => write_raw(&cluster.ca_cert(&name).await?), Some(CC::Token { name }) => println!("{}", &cluster.token(&name).await?), Some(CC::Addons) => write_raw(&cluster.addons().await?), Some(CC::SshSign { user_public_key, principal, validity, options, }) => { let pub_key = tokio::fs::read_to_string(user_public_key).await?; let cert = cluster .ssh_userca_sign(&dls::SshSignReq { pub_key, principal, validity: Some(validity).filter(|s| s != ""), options, }) .await?; write_raw(&cert); } Some(CC::KubeSign { csr, user, group, validity, }) => { let csr = tokio::fs::read_to_string(csr).await?; let cert = cluster .kube_sign(&dls::KubeSignReq { csr, user, group, validity: Some(validity).filter(|s| s != ""), }) .await?; write_raw(&cert); } } } C::Hosts => write_json(&dls().hosts().await?), C::Host { out, host, asset } => { let dls = dls(); let host_name = host.clone(); let host = dls.host(host); match asset { None => write_json(&host.config().await?), Some(asset) => { let stream = host.asset(&asset).await?; let mut out = create_asset_file(out, "host", &host_name, &asset).await?; copy_stream(stream, &mut out).await?; } } } C::DlSet(set) => match set { DlSet::Sign { expiry, items } => { let req = dls::DownloadSetReq { expiry, items }; let signed = dls().sign_dl_set(&req).await?; println!("{signed}"); } DlSet::Show { signed_set } => { let raw = base32::decode(base32::Alphabet::Rfc4648 { padding: false }, &signed_set) .ok_or(format_err!("invalid dlset"))?; let sig_len = raw[0] as usize; let (sig, data) = raw[1..].split_at(sig_len); println!("signature: {}...", hex::encode(&sig[..16])); let data = lz4::Decoder::new(data)?; let data = std::io::read_to_string(data)?; let (expiry, items) = data.split_once('|').ok_or(format_err!("invalid dlset"))?; let expiry = i64::from_str_radix(expiry, 16)?; let expiry = chrono::DateTime::from_timestamp(expiry, 0).unwrap(); println!("expires on {expiry}"); for item in items.split('|') { let mut parts = item.split(':'); let Some(kind) = parts.next() else { continue; }; let Some(name) = parts.next() else { continue; }; for asset in parts { println!("- {kind} {name} {asset}"); } } } DlSet::Fetch { signed_set, out, kind, name, asset, } => { let dls = dls(); let stream = dls.fetch_dl_set(&signed_set, &kind, &name, &asset).await?; let mut out = create_asset_file(out, &kind, &name, &asset).await?; copy_stream(stream, &mut out).await?; } }, C::Hash { salt } => { let salt = dkl::base64_decode(&salt)?; let passphrase = rpassword::prompt_password("password to hash: ")?; let hash = dls::store::hash_password(&salt, &passphrase)?; println!("hash (hex): {}", hex::encode(&hash)); println!("hash (base64): {}", dkl::base64_encode(&hash)); } }; Ok(()) } async fn create_asset_file( path: Option, kind: &str, name: &str, asset: &str, ) -> std::io::Result { let path = &path.unwrap_or(format!("{kind}_{name}_{asset}")); eprintln!("writing {kind} {name} asset {asset} to {path}"); (fs::File::options().write(true).create(true).truncate(true)) .mode(0o600) .open(path) .await } async fn copy_stream( mut stream: impl Stream> + Unpin, out: &mut (impl AsyncWrite + Unpin), ) -> std::io::Result<()> { let mut out = tokio::io::BufWriter::new(out); let info_delay = Duration::from_secs(1); let mut ts = SystemTime::now(); let mut n = 0u64; while let Some(chunk) = stream.next().await { let chunk = chunk.map_err(|e| std::io::Error::other(e))?; n += chunk.len() as u64; out.write_all(&chunk).await?; if ts.elapsed().is_ok_and(|t| t >= info_delay) { eprint!("wrote {n} bytes\r"); ts = SystemTime::now(); } } eprintln!("wrote {n} bytes"); out.flush().await } fn write_json(v: &T) { let data = serde_json::to_string_pretty(v).expect("value should serialize to json"); println!("{data}"); } fn write_raw(raw: &[u8]) { use std::io::Write; let mut out = std::io::stdout(); out.write_all(raw).expect("stdout write"); out.flush().expect("stdout flush"); } fn parse_download_set_item(s: &str) -> Result { let err = |s: &str| std::io::Error::other(s); let mut parts = s.split(':'); let item = dls::DownloadSetItem { kind: parts.next().ok_or(err("no kind"))?.to_string(), name: parts.next().ok_or(err("no name"))?.to_string(), assets: parts.map(|p| p.to_string()).collect(), }; Ok(item) }