stuff

1 files changed, 598 insertions, 0 deletions

diff --git a/src/main.rs b/src/main.rs
new file mode 100644
index 0000000..2184392
--- /dev/null
+++ b/src/main.rs
@@ -0,0 +1,598 @@
+use clap::Clap;
+use merge::Merge;
+use std::collections::HashMap;
+
+use std::borrow::Cow;
+
+use std::process::Stdio;
+use tokio::process::Command;
+
+use std::path::Path;
+
+use std::process;
+
+extern crate pretty_env_logger;
+#[macro_use]
+extern crate log;
+
+#[macro_use]
+extern crate serde_derive;
+
+macro_rules! good_panic {
+    ($($tts:tt)*) => {{
+        error!($($tts)*);
+        process::exit(1);
+    }}
+}
+
+/// Simple Rust rewrite of a simple Nix Flake deployment tool
+#[derive(Clap, Debug)]
+#[clap(version = "1.0", author = "notgne2 <gen2@gen2.space>")]
+struct Opts {
+    /// Log verbosity
+    #[clap(short, long, parse(from_occurrences))]
+    verbose: i32,
+
+    #[clap(subcommand)]
+    subcmd: SubCommand,
+}
+
+/// Deploy profiles
+#[derive(Clap, Debug)]
+struct DeployOpts {
+    /// The flake to deploy
+    #[clap(default_value = ".")]
+    flake: String,
+    /// Prepare server (for first deployments)
+    #[clap(short, long)]
+    prime: bool,
+}
+
+/// Activate a profile on your current machine
+#[derive(Clap, Debug)]
+struct ActivateOpts {
+    profile_path: String,
+    closure: String,
+
+    /// Command for activating the given profile
+    #[clap(short, long)]
+    activate_cmd: Option<String>,
+
+    /// Command for bootstrapping
+    #[clap(short, long)]
+    bootstrap_cmd: Option<String>,
+
+    /// Auto rollback if failure
+    #[clap(short, long)]
+    auto_rollback: bool,
+}
+
+#[derive(Clap, Debug)]
+enum SubCommand {
+    Deploy(DeployOpts),
+    Activate(ActivateOpts),
+}
+
+#[derive(Deserialize, Debug, Clone, Merge)]
+pub struct GenericSettings {
+    #[serde(rename(deserialize = "sshUser"))]
+    pub ssh_user: Option<String>,
+    pub user: Option<String>,
+    #[serde(
+        skip_serializing_if = "Vec::is_empty",
+        default,
+        rename(deserialize = "sshOpts")
+    )]
+    #[merge(strategy = merge::vec::append)]
+    pub ssh_opts: Vec<String>,
+    #[serde(rename(deserialize = "fastConnection"))]
+    pub fast_connection: Option<bool>,
+    #[serde(rename(deserialize = "autoRollback"))]
+    pub auto_rollback: Option<String>,
+}
+
+#[derive(Deserialize, Debug, Clone)]
+pub struct NodeSettings {
+    pub hostname: String,
+}
+
+#[derive(Deserialize, Debug, Clone)]
+pub struct ProfileSettings {
+    pub path: String,
+    pub activate: Option<String>,
+    pub bootstrap: Option<String>,
+}
+
+#[derive(Deserialize, Debug, Clone)]
+pub struct Profile {
+    #[serde(flatten)]
+    pub profile_settings: ProfileSettings,
+    #[serde(flatten)]
+    pub generic_settings: GenericSettings,
+}
+
+#[derive(Deserialize, Debug, Clone)]
+pub struct Node {
+    #[serde(flatten)]
+    pub generic_settings: GenericSettings,
+    #[serde(flatten)]
+    pub node_settings: NodeSettings,
+
+    pub profiles: HashMap<String, Profile>,
+}
+
+#[derive(Deserialize, Debug, Clone)]
+pub struct Data {
+    #[serde(flatten)]
+    pub generic_settings: GenericSettings,
+    pub nodes: HashMap<String, Node>,
+}
+
+async fn deploy_profile(
+    profile: &Profile,
+    profile_name: &str,
+    node: &Node,
+    node_name: &str,
+    top_settings: &GenericSettings,
+    supports_flakes: bool,
+    repo: &str,
+) -> Result<(), Box<dyn std::error::Error>> {
+    info!(
+        "Deploying profile `{}` for node `{}`",
+        profile_name, node_name
+    );
+
+    let mut merged_settings = top_settings.clone();
+    merged_settings.merge(node.generic_settings.clone());
+    merged_settings.merge(profile.generic_settings.clone());
+
+    let ssh_user: Cow<str> = match &merged_settings.ssh_user {
+        Some(u) => u.into(),
+        None => whoami::username().into(),
+    };
+
+    let profile_user: Cow<str> = match &merged_settings.user {
+        Some(x) => x.into(),
+        None => match &merged_settings.ssh_user {
+            Some(x) => x.into(),
+            None => good_panic!(
+                "Neither user nor sshUser set for profile `{}` of node `{}`",
+                profile_name,
+                node_name
+            ),
+        },
+    };
+
+    let sudo: Option<String> = match merged_settings.user {
+        Some(ref user) if user != &ssh_user => Some(format!("sudo -u {}", user).into()),
+        _ => None,
+    };
+
+    let profile_path = match &profile_user[..] {
+        "root" => format!("/nix/var/nix/profiles/{}", profile_name),
+        _ => format!(
+            "/nix/var/nix/profiles/per-user/{}/{}",
+            profile_user, profile_name
+        ),
+    };
+
+    info!(
+        "Building profile `{}` for node `{}`",
+        profile_name, node_name
+    );
+    if supports_flakes {
+        Command::new("nix")
+            .arg("build")
+            .arg("--no-link")
+            .arg(format!(
+                "{}#deploy.nodes.{}.profiles.{}.path",
+                repo, node_name, profile_name
+            ))
+            .stdout(Stdio::null())
+            .stderr(Stdio::null())
+            .spawn()?
+            .await?;
+    } else {
+        Command::new("nix-build")
+            .arg(&repo)
+            .arg("-A")
+            .arg(format!(
+                "deploy.nodes.{}.profiles.{}.path",
+                node_name, profile_name
+            ))
+            .stdout(Stdio::null())
+            .stderr(Stdio::null())
+            .spawn()?
+            .await?;
+    }
+
+    let current_exe = std::env::current_exe().expect("Expected to find current executable path");
+
+    if !current_exe.starts_with("/nix/store/") {
+        good_panic!("The deploy binary must be in the Nix store");
+    }
+
+    if let Ok(local_key) = std::env::var("LOCAL_KEY") {
+        info!(
+            "Signing key present! Signing profile `{}` for node `{}`",
+            profile_name, node_name
+        );
+
+        Command::new("nix")
+            .arg("sign-paths")
+            .arg("-r")
+            .arg("-k")
+            .arg(local_key)
+            .arg(&profile.profile_settings.path)
+            .arg(&current_exe)
+            .stdout(Stdio::null())
+            .stderr(Stdio::null())
+            .spawn()?
+            .await?;
+    }
+
+    info!("Copying profile `{} for node `{}`", profile_name, node_name);
+
+    let mut copy_command_ = Command::new("nix");
+    let mut copy_command = copy_command_.arg("copy");
+
+    if let Some(true) = merged_settings.fast_connection {
+        copy_command = copy_command.arg("--substitute-on-destination");
+    }
+
+    let ssh_opts_str = merged_settings
+        .ssh_opts
+        // This should provide some extra safety, but it also breaks for some reason, oh well
+        // .iter()
+        // .map(|x| format!("'{}'", x))
+        // .collect::<Vec<String>>()
+        .join(" ");
+
+    copy_command
+        .arg("--no-check-sigs")
+        .arg("--to")
+        .arg(format!(
+            "ssh://{}@{}",
+            ssh_user, node.node_settings.hostname
+        ))
+        .arg(&profile.profile_settings.path)
+        .arg(&current_exe)
+        .env("NIX_SSHOPTS", ssh_opts_str)
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .spawn()?
+        .await?;
+
+    info!(
+        "Activating profile `{}` for node `{}`",
+        profile_name, node_name
+    );
+
+    let mut self_activate_command = format!(
+        "{} activate '{}' '{}'",
+        current_exe.as_path().to_str().unwrap(),
+        profile_path,
+        profile.profile_settings.path,
+    );
+
+    if let Some(sudo_cmd) = sudo {
+        self_activate_command = format!("{} {}", sudo_cmd, self_activate_command);
+    }
+
+    if let Some(ref bootstrap_cmd) = profile.profile_settings.bootstrap {
+        self_activate_command = format!(
+            "{} --bootstrap-cmd '{}'",
+            self_activate_command, bootstrap_cmd
+        );
+    }
+
+    if let Some(ref activate_cmd) = profile.profile_settings.activate {
+        self_activate_command = format!(
+            "{} --activate-cmd '{}'",
+            self_activate_command, activate_cmd
+        );
+    }
+
+    let mut c = Command::new("ssh");
+    let mut ssh_command = c.arg(format!(
+        "ssh://{}@{}",
+        ssh_user, node.node_settings.hostname
+    ));
+
+    for ssh_opt in merged_settings.ssh_opts {
+        ssh_command = ssh_command.arg(ssh_opt);
+    }
+
+    ssh_command.arg(self_activate_command).spawn()?.await?;
+
+    Ok(())
+}
+
+#[inline]
+async fn deploy_all_profiles(
+    node: &Node,
+    node_name: &str,
+    supports_flakes: bool,
+    repo: &str,
+    top_settings: &GenericSettings,
+    prime: bool,
+) -> Result<(), Box<dyn std::error::Error>> {
+    info!("Deploying all profiles for `{}`", node_name);
+
+    if prime {
+        info!("Bootstrapping {}", node_name);
+
+        let profile = match node.profiles.get("system") {
+            Some(x) => x,
+            None => good_panic!("No system profile was found, needed for priming"),
+        };
+
+        deploy_profile(
+            &profile,
+            "system",
+            node,
+            node_name,
+            top_settings,
+            supports_flakes,
+            repo,
+        )
+        .await?;
+    }
+
+    for (profile_name, profile) in &node.profiles {
+        // This will have already been deployed
+        if prime && profile_name == "system" {
+            continue;
+        }
+
+        deploy_profile(
+            &profile,
+            profile_name,
+            node,
+            node_name,
+            top_settings,
+            supports_flakes,
+            repo,
+        )
+        .await?;
+    }
+
+    Ok(())
+}
+
+#[tokio::main]
+
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    if let Err(_) = std::env::var("DEPLOY_LOG") {
+        std::env::set_var("DEPLOY_LOG", "info");
+    }
+
+    pretty_env_logger::init_custom_env("DEPLOY_LOG");
+
+    let opts: Opts = Opts::parse();
+
+    match opts.subcmd {
+        SubCommand::Deploy(deploy_opts) => {
+            let flake_fragment_start = deploy_opts.flake.find('#');
+
+            let (repo, maybe_fragment) = match flake_fragment_start {
+                Some(s) => (&deploy_opts.flake[..s], Some(&deploy_opts.flake[s + 1..])),
+                None => (deploy_opts.flake.as_str(), None),
+            };
+
+            let (maybe_node, maybe_profile) = match maybe_fragment {
+                Some(fragment) => {
+                    let fragment_profile_start = fragment.find('.');
+                    match fragment_profile_start {
+                        Some(s) => (Some(&fragment[..s]), Some(&fragment[s + 1..])),
+                        None => (Some(fragment), None),
+                    }
+                }
+                None => (None, None),
+            };
+
+            let test_flake_status = Command::new("nix")
+                .arg("eval")
+                .arg("--expr")
+                .arg("builtins.getFlake")
+                .stdout(Stdio::null())
+                .stderr(Stdio::null())
+                .status()
+                .await?;
+
+            let supports_flakes = test_flake_status.success();
+
+            let data_json = match supports_flakes {
+                true => {
+                    let c = Command::new("nix")
+                        .arg("eval")
+                        .arg("--json")
+                        .arg(format!("{}#deploy", repo))
+                        .stdout(Stdio::null())
+                        .stderr(Stdio::null())
+                        // TODO forward input args?
+                        .output()
+                        .await?;
+
+                    String::from_utf8(c.stdout)?
+                }
+                false => {
+                    let c = Command::new("nix-instanciate")
+                        .arg("--strict")
+                        .arg("--read-write-mode")
+                        .arg("--json")
+                        .arg("--eval")
+                        .arg("--E")
+                        .arg(format!("let r = import {}/.; in if builtins.isFunction r then (r {{}}).deploy else r.deploy", repo))
+                        .stdout(Stdio::null())
+                        .stderr(Stdio::null())
+                        .output()
+                        .await?;
+
+                    String::from_utf8(c.stdout)?
+                }
+            };
+
+            let data: Data = serde_json::from_str(&data_json)?;
+
+            match (maybe_node, maybe_profile) {
+                (Some(node_name), Some(profile_name)) => {
+                    let node = match data.nodes.get(node_name) {
+                        Some(x) => x,
+                        None => good_panic!("No node was found named `{}`", node_name),
+                    };
+                    let profile = match node.profiles.get(profile_name) {
+                        Some(x) => x,
+                        None => good_panic!("No profile was found named `{}`", profile_name),
+                    };
+
+                    deploy_profile(
+                        &profile,
+                        profile_name,
+                        node,
+                        node_name,
+                        &data.generic_settings,
+                        supports_flakes,
+                        repo,
+                    )
+                    .await?;
+                }
+                (Some(node_name), None) => {
+                    let node = match data.nodes.get(node_name) {
+                        Some(x) => x,
+                        None => good_panic!("No node was found named `{}`", node_name),
+                    };
+
+                    deploy_all_profiles(
+                        node,
+                        node_name,
+                        supports_flakes,
+                        repo,
+                        &data.generic_settings,
+                        deploy_opts.prime,
+                    )
+                    .await?;
+                }
+                (None, None) => {
+                    info!("Deploying all profiles on all nodes");
+
+                    for (node_name, node) in &data.nodes {
+                        deploy_all_profiles(
+                            node,
+                            node_name,
+                            supports_flakes,
+                            repo,
+                            &data.generic_settings,
+                            deploy_opts.prime,
+                        )
+                        .await?;
+                    }
+                }
+                (None, Some(_)) => good_panic!(
+                    "Profile provided without a node, this is not (currently) supported"
+                ),
+            };
+        }
+        SubCommand::Activate(activate_opts) => {
+            info!("Activating profile");
+
+            Command::new("nix-env")
+                .arg("-p")
+                .arg(&activate_opts.profile_path)
+                .arg("--set")
+                .arg(&activate_opts.closure)
+                .stdout(Stdio::null())
+                .spawn()?
+                .await?;
+
+            if let (Some(bootstrap_cmd), false) = (
+                activate_opts.bootstrap_cmd,
+                !Path::new(&activate_opts.profile_path).exists(),
+            ) {
+                let bootstrap_status = Command::new("bash")
+                    .arg("-c")
+                    .arg(&bootstrap_cmd)
+                    .env("PROFILE", &activate_opts.profile_path)
+                    .stdout(Stdio::null())
+                    .stderr(Stdio::null())
+                    .status()
+                    .await;
+
+                match bootstrap_status {
+                    Ok(s) if s.success() => (),
+                    _ => {
+                        tokio::fs::remove_file(&activate_opts.profile_path).await?;
+                        good_panic!("Failed to execute bootstrap command");
+                    }
+                }
+            }
+
+            if let Some(activate_cmd) = activate_opts.activate_cmd {
+                let activate_status = Command::new("bash")
+                    .arg("-c")
+                    .arg(&activate_cmd)
+                    .env("PROFILE", &activate_opts.profile_path)
+                    .status()
+                    .await;
+
+                match activate_status {
+                    Ok(s) if s.success() => (),
+                    _ if activate_opts.auto_rollback => {
+                        Command::new("nix-env")
+                            .arg("-p")
+                            .arg(&activate_opts.profile_path)
+                            .arg("--rollback")
+                            .stdout(Stdio::null())
+                            .stderr(Stdio::null())
+                            .spawn()?
+                            .await?;
+
+                        let c = Command::new("nix-env")
+                            .arg("-p")
+                            .arg(&activate_opts.profile_path)
+                            .arg("--list-generations")
+                            .output()
+                            .await?;
+                        let generations_list = String::from_utf8(c.stdout)?;
+
+                        let last_generation_line = generations_list
+                            .lines()
+                            .last()
+                            .expect("Expected to find a generation in list");
+
+                        let last_generation_id = last_generation_line
+                            .split_whitespace()
+                            .next()
+                            .expect("Expected to get ID from generation entry");
+
+                        debug!("Removing generation entry {}", last_generation_line);
+                        warn!("Removing generation by ID {}", last_generation_id);
+
+                        Command::new("nix-env")
+                            .arg("-p")
+                            .arg(&activate_opts.profile_path)
+                            .arg("--delete-generations")
+                            .arg(last_generation_id)
+                            .stdout(Stdio::null())
+                            .stderr(Stdio::null())
+                            .spawn()?
+                            .await?;
+
+                        // TODO: why are we doing this?
+                        // to run the older version as long as the command is the same?
+                        Command::new("bash")
+                            .arg("-c")
+                            .arg(&activate_cmd)
+                            .spawn()?
+                            .await?;
+
+                        good_panic!("Failed to execute activation command");
+                    }
+                    _ => {}
+                }
+            }
+        }
+    }
+
+    Ok(())
+}