dkl/src/bin/dls.rs

use clap::{Parser, Subcommand};
use eyre::{format_err, Result};
use futures_util::StreamExt;
use tokio::io::AsyncWriteExt;

use dkl::dls;

#[derive(Parser)]
#[command()]
struct Cli {
    #[arg(long, default_value = "http://[::1]:7606")]
    dls: String,

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

#[derive(Subcommand)]
enum Command {
    Clusters,
    Cluster {
        cluster: String,
        #[command(subcommand)]
        command: Option<ClusterCommand>,
    },
    Hosts,
    Host {
        host: String,
        asset: Option<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<String>,
    },
}

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<()> {
    let cli = Cli::parse();

    env_logger::builder()
        .parse_filters("info")
        .parse_default_env()
        .init();

    let token = std::env::var("DLS_TOKEN").map_err(|_| format_err!("DLS_TOKEN should be set"))?;

    let dls = 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 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
                        .sign_ssh_user_pubkey(&dls::SshSignReq {
                            pub_key,
                            principal,
                            validity: Some(validity).filter(|s| s != ""),
                            options,
                        })
                        .await?;
                    write_raw(&cert);
                }
            }
        }
        C::Hosts => write_json(&dls.hosts().await?),
        C::Host { host, asset } => {
            let host_name = host.clone();
            let host = dls.host(host);
            match asset {
                None => write_json(&host.config().await?),
                Some(asset) => {
                    let mut stream = host.asset(&asset).await?;

                    let out_path = format!("{host_name}_{asset}");
                    eprintln!("writing {host_name} asset {asset} to {out_path}");

                    let out = tokio::fs::File::options()
                        .mode(0o600)
                        .write(true)
                        .create(true)
                        .truncate(true)
                        .open(out_path)
                        .await?;
                    let mut out = tokio::io::BufWriter::new(out);

                    let mut n = 0u64;
                    while let Some(chunk) = stream.next().await {
                        let chunk = chunk?;
                        n += chunk.len() as u64;
                        eprint!("wrote {n} bytes\r");
                        out.write_all(&chunk).await?;
                    }
                    eprintln!();

                    out.flush().await?;
                }
            }
        }
    };

    Ok(())
}

fn write_json<T: serde::ser::Serialize>(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(raw).expect("stdout write");
    out.flush().expect("stdout flush");
}