machines/rust/program/infrastructure/src/main.rs

use common::bitwarden::BitwardenSession;
use serde::Deserialize;
use std::{
    borrow::Cow,
    collections::{BTreeMap, BTreeSet},
    path::{Path, PathBuf},
};

use clap::Parser;

mod command;
mod secrets;

use crate::{command::Command, secrets::CliEndpoint};

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

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

fn program() -> anyhow::Result<()> {
    let args = Args::parse();
    args.command.run()
}

#[derive(Deserialize, Debug)]
struct ProvisioningData {
    persistence: OpenTofuConfig,
    compute: OpenTofuConfig,
    compute_with_persistence_attached: OpenTofuConfig,
    configuration: OpenTofuConfig,
    secrets_source: SecretsSource,

    image_username: String,
}

struct EndpointReader {
    bw_session: Option<BitwardenSession>,
    endpoints_read: BTreeSet<CliEndpoint>,
    env: BTreeMap<Cow<'static, str>, String>,
}

impl EndpointReader {
    pub fn new() -> Self {
        Self {
            bw_session: None,
            endpoints_read: BTreeSet::new(),
            env: BTreeMap::new(),
        }
    }
    pub fn env(&self) -> &BTreeMap<Cow<'static, str>, String> {
        &self.env
    }

    pub fn read_endpoints(
        &mut self,
        secrets_source: SecretsSource,
        endpoints: &[CliEndpoint],
    ) -> anyhow::Result<()> {
        let endpoints_to_read: Vec<_> = endpoints
            .iter()
            .copied()
            .filter(|e| !self.endpoints_read.contains(e))
            .collect();

        if endpoints_to_read.is_empty() {
            return Ok(());
        }

        match secrets_source {
            SecretsSource::Bitwarden => {
                if self.bw_session.is_none() {
                    let bw_session = BitwardenSession::unlock()?;
                    let _ = self.bw_session.insert(bw_session);
                }
                let session = self
                    .bw_session
                    .as_mut()
                    .expect("Should have bitwarden session");

                for endpoint in endpoints_to_read {
                    endpoint.read_from_bitwarden(session, &mut self.env)?;
                    self.endpoints_read.insert(endpoint);
                }
            }
            SecretsSource::Vault => {
                for endpoint in endpoints_to_read {
                    endpoint.read_from_openbao(&mut self.env)?;
                    self.endpoints_read.insert(endpoint);
                }
            }
        }

        Ok(())
    }
}

#[derive(Deserialize, Debug, Clone, Copy, PartialEq)]
enum SecretsSource {
    #[serde(rename = "bitwarden")]
    Bitwarden,
    #[serde(rename = "vault")]
    Vault,
}

#[derive(Deserialize, Debug)]
struct OpenTofuConfig {
    config: Option<PathBuf>,
    endpoints: Vec<CliEndpoint>,
}

struct WorkDir {
    path: PathBuf,
}

impl WorkDir {
    pub fn try_new(template: &str) -> anyhow::Result<Self> {
        let mut proc = common::proc::Command::new("mktemp");
        proc.args(["-dt", template]);
        let path: PathBuf = proc.try_spawn_to_string()?.into();
        common::fs::create_dir_recursive(&path)?;
        Ok(Self { path })
    }

    pub fn cleanup(self) -> anyhow::Result<()> {
        common::fs::remove_dir_recursive(&self.path)?;
        Ok(())
    }
}

impl OpenTofuConfig {
    pub fn init<'e, K: AsRef<str>, V: AsRef<str>>(
        &self,
        instance: &str,
        env_map: &'e BTreeMap<K, V>,
    ) -> anyhow::Result<Option<OpenTofuInstance<'e, K, V>>> {
        let Some(config) = self.config.as_ref() else {
            return Ok(None);
        };
        let work_dir = WorkDir::try_new(&format!("{instance}-provision.XXXXXX"))?;
        let config_path: PathBuf = work_dir.path.join("config.tf.json");
        common::fs::create_link(config, &config_path)?;

        let mut init_proc = common::proc::Command::new("tofu");
        let chdir_arg = format!("-chdir={}", work_dir.path.display());
        for (key, value) in env_map {
            init_proc.env_sensitive(key.as_ref(), value.as_ref());
        }
        init_proc.args([&chdir_arg, "init"]);
        init_proc.try_spawn_to_bytes()?;
        Ok(Some(OpenTofuInstance {
            work_dir,
            env_map,
            chdir_arg,
        }))
    }
}

struct OpenTofuInstance<'e, K, V> {
    work_dir: WorkDir,
    chdir_arg: String,
    env_map: &'e BTreeMap<K, V>,
}

impl<'e, K: AsRef<str>, V: AsRef<str>> OpenTofuInstance<'e, K, V> {
    pub fn run(&mut self, action: &str) -> anyhow::Result<()> {
        let mut proc = common::proc::Command::new("tofu");
        for (key, value) in self.env_map {
            proc.env_sensitive(key.as_ref(), value.as_ref());
        }
        proc.args([&self.chdir_arg, action]);
        proc.stdin_inherit();
        proc.stderr_inherit();
        proc.try_spawn_stdout_inherit()?;
        Ok(())
    }

    pub fn output<D: for<'de> serde::Deserialize<'de>>(&mut self) -> anyhow::Result<D> {
        let mut proc = common::proc::Command::new("tofu");
        for (key, value) in self.env_map {
            proc.env_sensitive(key.as_ref(), value.as_ref());
        }
        proc.args([&self.chdir_arg, "output", "-json"]);
        proc.try_spawn_to_json()
    }

    pub fn cleanup(self) -> anyhow::Result<()> {
        self.work_dir.cleanup()
    }
}

impl ProvisioningData {
    pub fn try_new(instance: &str, flake_path: &Path) -> anyhow::Result<Self> {
        let mut proc = common::proc::Command::new("nix");
        let base_attr = format!(
            "{}#nixosConfigurations.\"{instance}\".config.khscodes.infrastructure.provisioning",
            flake_path.display()
        );
        let script = r#"
        let
            data = prov: { inherit (prov) config endpoints; };
        in
        p: {
          persistence = data p.persistence;
          compute = data p.compute;
          compute_with_persistence_attached = data p.combinedPersistenceAttachAndCompute;
          configuration = data p.configuration;
          secrets_source = p.secretsSource;
          image_username = p.imageUsername;
        }
        "#;
        proc.args(["eval", "--json", &base_attr, "--apply", script]);

        log::info!("Building terranix configurations...");
        let result = proc.stderr_inherit().try_spawn_to_json()?;
        log::info!("Terranix configurations built.");
        Ok(result)
    }
}