dkl/src/bin/dls.rs

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<ClusterCommand>,
    },
    Hosts,
    Host {
        #[arg(short = 'o', long)]
        out: Option<String>,
        host: String,
        asset: Option<String>,
    },
    #[command(subcommand)]
    DlSet(DlSet),
}

#[derive(Subcommand)]
enum DlSet {
    Sign {
        #[arg(short = 'e', long, default_value = "1d")]
        expiry: String,
        #[arg(value_parser = parse_download_set_item)]
        items: Vec<dls::DownloadSetItem>,
    },
    Show {
        #[arg(env = "DLS_DLSET")]
        signed_set: String,
    },
    Fetch {
        #[arg(long, env = "DLS_DLSET")]
        signed_set: String,
        #[arg(short = 'o', long)]
        out: Option<String>,
        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<String>,
    },
    KubeSign {
        csr: String,
        #[arg(long, default_value = "anonymous", env = "USER")]
        user: String,
        #[arg(long)]
        group: Option<String>,
        #[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 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
                        .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 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 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?;
            }
        },
    };

    Ok(())
}

async fn create_asset_file(
    path: Option<String>,
    kind: &str,
    name: &str,
    asset: &str,
) -> std::io::Result<fs::File> {
    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<Item = reqwest::Result<Bytes>> + 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<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_all(raw).expect("stdout write");
    out.flush().expect("stdout flush");
}

fn parse_download_set_item(s: &str) -> Result<dls::DownloadSetItem, std::io::Error> {
    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)
}