From b2b90bb946d2529f74ce190ea1ccd2a56990fa72 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sat, 25 Jan 2025 23:56:35 -0800
Subject: [PATCH] WIP: Add boot mode options to boot-diags

---
 boot_diags/src/app.rs   | 55 +++++++++++++++++++++++++++-------
 core/src/system/boot.rs | 66 +++++++++++++++++++++++++++++------------
 2 files changed, 92 insertions(+), 29 deletions(-)

diff --git a/boot_diags/src/app.rs b/boot_diags/src/app.rs
index ae089df..bdbf32c 100644
--- a/boot_diags/src/app.rs
+++ b/boot_diags/src/app.rs
@@ -22,11 +22,16 @@ use core::{
     config::Config,
     line::{get_disk_description_right, get_part_description, DVLine},
     state::{CloneSettings, Mode},
-    system::{cpu::get_cpu_name, drivers},
+    system::{
+        boot::{self, SafeMode},
+        cpu::get_cpu_name,
+        drivers,
+    },
     tasks::{Task, Tasks},
     tui::{Event, Tui},
 };
 use std::{
+    env,
     iter::zip,
     sync::{Arc, Mutex},
 };
@@ -41,13 +46,6 @@ use ratatui::{
 use tokio::sync::mpsc;
 use tracing::{debug, info};
 
-#[derive(Clone, Debug, Default)]
-pub enum SafeMode {
-    #[default]
-    Disable,
-    Enable,
-}
-
 pub struct App {
     // TUI
     action_rx: mpsc::UnboundedReceiver<Action>,
@@ -65,6 +63,7 @@ pub struct App {
     list: StatefulList<Mode>,
     boot_modes: Vec<SafeMode>,
     selections: Vec<Option<usize>>,
+    system32: String,
     tasks: Tasks,
 }
 
@@ -103,6 +102,7 @@ impl App {
             cur_mode: Mode::Home,
             list,
             boot_modes: vec![SafeMode::Enable, SafeMode::Disable],
+            system32: String::new(),
             selections: vec![None, None],
             tasks,
         })
@@ -124,6 +124,25 @@ impl App {
         }
     }
 
+    pub fn set_boot_mode(&mut self, boot_mode: SafeMode) {
+        info!("Setting boot mode to: {:?}", boot_mode);
+        let disk_list = self.clone.disk_list.lock().unwrap();
+        if let Some(disk_index) = self.clone.disk_index_dest {
+            if let Some(disk) = disk_list.get(disk_index) {
+                if let Some(boot_index) = self.clone.part_index_boot {
+                    if let Ok(task) = boot::set_mode(
+                        disk.get_part_letter(boot_index).as_str(),
+                        boot_mode,
+                        &self.system32,
+                        &disk.part_type,
+                    ) {
+                        self.tasks.add(task);
+                    };
+                }
+            }
+        }
+    }
+
     pub fn set_mode(&mut self, new_mode: Mode) -> Result<()> {
         info!("Setting mode to {new_mode:?}");
         self.cur_mode = new_mode;
@@ -166,7 +185,22 @@ impl App {
         }
 
         let action_tx = self.action_tx.clone();
+
+        // Late init
+        self.system32 = if cfg!(windows) {
+            if let Ok(path) = env::var("SYSTEMROOT") {
+                format!("{path}/System32")
+            } else {
+                self.action_tx.send(Action::Error(String::from(
+                    "ERROR\n\n\nFailed to find SYSTEMROOT",
+                )))?;
+                return Ok(());
+            }
+        } else {
+            String::from(".")
+        };
         action_tx.send(Action::SetMode(Mode::ScanDisks))?;
+
         loop {
             self.handle_events(&mut tui).await?;
             self.handle_actions(&mut tui)?;
@@ -299,8 +333,9 @@ impl App {
                     }
                     Mode::SetBootMode => {
                         if let Some(index) = one {
-                            // TODO: Make real
-                            info!("Setting boot mode to: {:?}", self.boot_modes.get(index));
+                            if let Some(boot_mode) = self.boot_modes.get(index) {
+                                self.set_boot_mode(boot_mode.to_owned());
+                            }
                         }
                     }
                     _ => {}
diff --git a/core/src/system/boot.rs b/core/src/system/boot.rs
index 2bc098f..94edf3e 100644
--- a/core/src/system/boot.rs
+++ b/core/src/system/boot.rs
@@ -18,6 +18,13 @@ use crate::tasks::Task;
 use color_eyre::Result;
 use std::path::PathBuf;
 
+#[derive(Clone, Debug, Default)]
+pub enum SafeMode {
+    #[default]
+    Disable,
+    Enable,
+}
+
 pub fn configure_disk(
     letter_boot: &str,
     letter_os: &str,
@@ -55,25 +62,10 @@ pub fn configure_disk(
     }
 
     // Lock in safe mode
-    let bcd_path = match table_type {
-        PartitionTableType::Guid => {
-            format!("{letter_boot}:\\EFI\\Microsoft\\Boot\\BCD")
-        }
-        PartitionTableType::Legacy => {
-            format!("{letter_boot}:\\Boot\\BCD")
-        }
-    };
-    tasks.push(Task::Command(
-        PathBuf::from(format!("{system32}/bcdedit.exe")),
-        vec![
-            String::from("/store"),
-            bcd_path,
-            String::from("/set"),
-            String::from("{default}"),
-            String::from("safeboot"),
-            String::from("minimal"),
-        ],
-    ));
+    tasks.push(
+        set_mode(letter_boot, SafeMode::Enable, system32, &table_type)
+            .expect("Failed to create set_mode task."),
+    );
 
     // Done
     tasks
@@ -92,3 +84,39 @@ pub fn inject_driver(driver: &Driver, letter_os: &str, system32: &str) -> Result
         ],
     ))
 }
+
+pub fn set_mode(
+    letter_boot: &str,
+    mode: SafeMode,
+    system32: &str,
+    table_type: &PartitionTableType,
+) -> Result<Task> {
+    let bcd_path = match table_type {
+        PartitionTableType::Guid => {
+            format!("{letter_boot}:\\EFI\\Microsoft\\Boot\\BCD")
+        }
+        PartitionTableType::Legacy => {
+            format!("{letter_boot}:\\Boot\\BCD")
+        }
+    };
+
+    // Build Command
+    let mut cmd_args = vec![String::from("/store"), bcd_path];
+    match mode {
+        SafeMode::Disable => {
+            cmd_args.push(String::from("/deletevalue"));
+            cmd_args.push(String::from("{default}"));
+            cmd_args.push(String::from("safeboot"));
+        }
+        SafeMode::Enable => {
+            cmd_args.push(String::from("/set"));
+            cmd_args.push(String::from("{default}"));
+            cmd_args.push(String::from("safeboot"));
+            cmd_args.push(String::from("minimal"));
+        }
+    }
+    Ok(Task::Command(
+        PathBuf::from(format!("{system32}/bcdedit.exe")),
+        cmd_args,
+    ))
+}