From a0144156b972e6e5f78bdfed64fb2470c72fb6aa Mon Sep 17 00:00:00 2001
From: Fernando Ayats <fernando.ayats-llamas@inria.fr>
Date: Tue, 20 Aug 2024 09:54:10 +0200
Subject: [PATCH] Add root check bypass flag, fixes #95

---
 Cargo.lock       |  2 +-
 Cargo.toml       |  2 +-
 src/interface.rs |  4 ++++
 src/nixos.rs     | 34 +++++++++++++++++++++-------------
 4 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 6d438a3..27ceade 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -737,7 +737,7 @@ dependencies = [
 
 [[package]]
 name = "nh"
-version = "3.5.23"
+version = "3.5.24"
 dependencies = [
  "ambassador",
  "anstyle",
diff --git a/Cargo.toml b/Cargo.toml
index 9c7e9bf..203cffe 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "nh"
-version = "3.5.23"
+version = "3.5.24"
 edition = "2021"
 license = "EUPL-1.2"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
diff --git a/src/interface.rs b/src/interface.rs
index 351d7de..ca42a81 100644
--- a/src/interface.rs
+++ b/src/interface.rs
@@ -112,6 +112,10 @@ pub struct OsRebuildArgs {
     /// Extra arguments passed to nix build
     #[arg(last = true)]
     pub extra_args: Vec<String>,
+
+    /// Bypass the check to call nh as root directly.
+    #[arg(short = 'R', long)]
+    pub bypass_root_check: bool,
 }
 
 #[derive(Debug, Args)]
diff --git a/src/nixos.rs b/src/nixos.rs
index 9fe82eb..289c47f 100644
--- a/src/nixos.rs
+++ b/src/nixos.rs
@@ -1,9 +1,10 @@
 use std::ops::Deref;
+use std::{env, vec};
 
 use color_eyre::eyre::{bail, Context};
 use color_eyre::Result;
 
-use tracing::{debug, info};
+use tracing::{debug, info, warn};
 
 use crate::interface::NHRunnable;
 use crate::interface::OsRebuildType::{self, Boot, Build, Switch, Test};
@@ -27,9 +28,19 @@ impl NHRunnable for interface::OsArgs {
 
 impl OsRebuildArgs {
     pub fn rebuild(&self, rebuild_type: &OsRebuildType) -> Result<()> {
-        if nix::unistd::Uid::effective().is_root() {
-            bail!("Don't run nh os as root. I will call sudo internally as needed");
-        }
+        let use_sudo = if self.bypass_root_check {
+            warn!("Bypassing root check, now running nix as root");
+            false
+        } else {
+            if nix::unistd::Uid::effective().is_root() {
+                bail!("Don't run nh os as root. I will call sudo internally as needed");
+            }
+            true
+        };
+
+        // TODO: add a  .maybe_arg to CommandBuilder
+        // so that I can do .maybe_arg( Option<T> )
+        let sudo_args: &[_] = if use_sudo { &["sudo"] } else { &[] };
 
         let hostname = match &self.hostname {
             Some(h) => h.to_owned(),
@@ -134,7 +145,8 @@ impl OsRebuildArgs {
             let switch_to_configuration = switch_to_configuration.to_str().unwrap();
 
             commands::CommandBuilder::default()
-                .args(["sudo", switch_to_configuration, "test"])
+                .args(sudo_args)
+                .args([switch_to_configuration, "test"])
                 .message("Activating configuration")
                 .build()?
                 .exec()?;
@@ -142,13 +154,8 @@ impl OsRebuildArgs {
 
         if let Boot(_) | Switch(_) = rebuild_type {
             commands::CommandBuilder::default()
-                .args([
-                    "sudo",
-                    "nix-env",
-                    "--profile",
-                    SYSTEM_PROFILE,
-                    "--set",
-                ])
+                .args(sudo_args)
+                .args(["nix-env", "--profile", SYSTEM_PROFILE, "--set"])
                 .args([out_path.get_path()])
                 .build()?
                 .exec()?;
@@ -161,7 +168,8 @@ impl OsRebuildArgs {
             let switch_to_configuration = switch_to_configuration.to_str().unwrap();
 
             commands::CommandBuilder::default()
-                .args(["sudo", switch_to_configuration, "boot"])
+                .args(sudo_args)
+                .args([switch_to_configuration, "boot"])
                 .message("Adding configuration to bootloader")
                 .build()?
                 .exec()?;