machines/rust/program/openbao-helper/src/main.rs

use std::{
    collections::BTreeSet,
    convert::Infallible,
    ffi::{CStr, CString, OsStr, OsString},
};

use anyhow::Context as _;
use clap::{Parser, Subcommand};

mod enventry;

use enventry::*;

fn main() {
    common::entrypoint(program);
}

#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
pub struct Args {
    #[command(subcommand)]
    pub command: Commands,
}

#[derive(Debug, Subcommand)]
pub enum Commands {
    /// Transfers secrets from the current environment into their respective secrets in openbao
    Transfer,
    /// Wraps the program by reading the specified data from bao and setting respective environment variables
    WrapProgram(WrapProgram),
}

#[derive(Debug, Clone, clap::Args)]
pub struct WrapProgram {
    /// The endpoints to fetch,
    #[arg(short = 'e', long = "endpoint", number_of_values = 1)]
    pub endpoint: Vec<Endpoint>,
    /// Command to wrap
    #[arg(allow_hyphen_values = true, last = true)]
    pub cmd: Vec<String>,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, clap::ValueEnum)]
pub enum Endpoint {
    #[value(name = "openstack")]
    Openstack,
    #[value(name = "cloudflare")]
    Cloudflare,
    #[value(name = "aws")]
    Aws,
    #[value(name = "hcloud")]
    Hcloud,
    #[value(name = "unifi")]
    Unifi,
    #[value(name = "vault")]
    Vault,
}

impl Endpoint {
    pub fn try_into_env_data(self) -> anyhow::Result<Vec<(&'static str, String)>> {
        match self {
            Self::Openstack => {
                let data = OpenstackData::read_from_bao()?;
                Ok(data.into())
            }
            Self::Aws => {
                let data = AwsData::read_from_bao()?;
                Ok(data.into())
            }
            Self::Hcloud => {
                let data = HcloudData::read_from_bao()?;
                Ok(data.into())
            }
            Self::Cloudflare => {
                let data = CloudflareData::read_from_bao()?;
                Ok(data.into())
            }
            Self::Unifi => {
                let data = UnifiData::read_from_bao()?;
                Ok(data.into())
            }
            Self::Vault => {
                let data = VaultData::read_from_bao()?;
                Ok(data.into())
            }
        }
    }
}

fn program() -> anyhow::Result<()> {
    let args = Args::parse();
    match args.command {
        Commands::Transfer => transfer(),
        Commands::WrapProgram(w) => wrap_program(w),
    }
}

macro_rules! entry_definition {
    ($config_id:ident, $id: ident, $bao_key: expr, $secrets: expr) => {
        struct $config_id;

        impl EnvEntryConfig for $config_id {
            const SECRETS: &'static [&'static str] = $secrets;
            const BAO_KEY: &'static str = $bao_key;
        }

        type $id = EnvEntry<$config_id>;
    };
}

entry_definition!(
    OpenstackDataConfig,
    OpenstackData,
    "openstack",
    &[
        "TF_VAR_openstack_username",
        "TF_VAR_openstack_password",
        "TF_VAR_openstack_tenant_name",
        "TF_VAR_openstack_auth_url",
        "TF_VAR_openstack_endpoint_type",
        "TF_VAR_openstack_region"
    ]
);
entry_definition!(
    CloudflareDataConfig,
    CloudflareData,
    "cloudflare",
    &["TF_VAR_cloudflare_token", "TF_VAR_cloudflare_email"]
);
entry_definition!(
    AwsDataConfig,
    AwsData,
    "aws",
    &["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"]
);
entry_definition!(
    HcloudDataConfig,
    HcloudData,
    "hcloud",
    &["TF_VAR_hcloud_api_token"]
);
entry_definition!(
    UnifiDataConfig,
    UnifiData,
    "unifi",
    &["UNIFI_USERNAME", "UNIFI_PASSWORD", "UNIFI_API"]
);
entry_definition!(VaultDataConfig, VaultData, "vault", &["VAULT_TOKEN"]);

fn transfer() -> anyhow::Result<()> {
    let openstack = OpenstackData::try_new_from_env()?;
    let cloudflare = CloudflareData::try_new_from_env()?;
    let aws = AwsData::try_new_from_env()?;
    let hcloud = HcloudData::try_new_from_env()?;
    let unifi = UnifiData::try_new_from_env()?;
    let vault = VaultData::try_new_from_env()?;

    write_kv_data(openstack)?;
    write_kv_data(cloudflare)?;
    write_kv_data(aws)?;
    write_kv_data(hcloud)?;
    write_kv_data(unifi)?;
    write_kv_data(vault)?;
    Ok(())
}

fn write_kv_data<T: EnvEntryConfig>(entry: EnvEntry<T>) -> anyhow::Result<()> {
    let mut cmd = common::proc::Command::new("bao");
    cmd.args(["kv", "put", "-mount=opentofu"]);
    cmd.arg(T::BAO_KEY);
    for (key, value) in entry {
        cmd.arg(format!("{key}={value}"));
    }
    cmd.try_spawn_to_string()?;
    Ok(())
}

fn wrap_program(wrap_program: WrapProgram) -> anyhow::Result<()> {
    let (args, env) = {
        let WrapProgram { cmd, endpoint } = wrap_program;
        if endpoint.is_empty() {
            return Err(anyhow::format_err!("Must specify at least one endpoint"));
        }
        if cmd.is_empty() {
            return Err(anyhow::format_err!("No command to execute was specified"));
        }
        let unique: BTreeSet<_> = BTreeSet::from_iter(endpoint);
        let mut env = Vec::<(OsString, OsString)>::new();
        for (key, value) in std::env::vars() {
            env.push((OsString::from(key), OsString::from(value)));
        }
        for env_set in unique {
            for (key, value) in env_set.try_into_env_data()? {
                env.push((OsString::from(key), OsString::from(value)));
            }
        }
        let mut args = Vec::new();
        for arg in cmd {
            let arg = CString::new(arg)
                .context("Argument to program to wrap cannot contain null bytes")?;
            args.push(arg);
        }
        (args, env)
    };
    unsafe {
        execvpe(&args[0], args.as_slice(), env.as_slice())?;
    }
    // This will never get executed
    Ok(())
}

#[cfg(not(target_os = "macos"))]
/// Safety: No other threads may read or write environment variables when this function is called.
///         The easiest way to ensure this is using a single threaded program.
/// Note: On Linux specifically this safety requirement is not needed
unsafe fn execvpe<SA: AsRef<CStr>, SEK: AsRef<OsStr>, SEV: AsRef<OsStr>>(
    filename: &CStr,
    args: &[SA],
    environ: &[(SEK, SEV)],
) -> anyhow::Result<Infallible> {
    let environ = environ
        .iter()
        .map(|(k, v)| {
            CString::new(format!(
                "{k}={v}",
                k = k.as_ref().display(),
                v = v.as_ref().display()
            ))
            .with_context(|| {
                format!(
                    "Environment variable {k} contains null bytes",
                    k = k.as_ref().display()
                )
            })
        })
        .collect::<anyhow::Result<Vec<CString>>>()?;
    Ok(nix::unistd::execvpe(filename, args, &environ)?)
}

#[cfg(target_os = "macos")]
/// Safety: No other threads may read or write environment variables when this function is called.
///         The easiest way to ensure this is using a single threaded program.
// Simple "bad" version of execvpe that also works on OSX
unsafe fn execvpe<SA: AsRef<CStr>, SEK: AsRef<OsStr>, SEV: AsRef<OsStr>>(
    filename: &CStr,
    args: &[SA],
    environ: &[(SEK, SEV)],
) -> anyhow::Result<Infallible> {
    let current_env = std::env::vars_os();
    // Safety: Same as this function
    unsafe { nix::env::clearenv()? };
    for (key, val) in environ {
        // Safety: Same as this function
        unsafe { std::env::set_var(key.as_ref(), val.as_ref()) };
    }
    match nix::unistd::execvp(filename, args) {
        Err(err) => {
            unsafe { nix::env::clearenv()? };
            for (key, val) in current_env {
                unsafe { std::env::set_var(key.as_os_str(), val.as_os_str()) };
            }
            Err(err.into())
        }
        _ => unreachable!("execvp doesn't return on success"),
    }
}