From 14df5b82ce79f7ff20390aff9f96154fc8394ec6 Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sat, 9 Nov 2024 23:33:12 -0800 Subject: [PATCH] Add clone and Pre/Post actions --- src/action.rs | 4 +- src/app.rs | 158 +++++++++++++++++++++++++++++++++++------ src/system/disk.rs | 3 + src/system/diskpart.rs | 1 + 4 files changed, 142 insertions(+), 24 deletions(-) diff --git a/src/action.rs b/src/action.rs index bf8ce58..4285a7c 100644 --- a/src/action.rs +++ b/src/action.rs @@ -29,8 +29,8 @@ use crate::{ #[derive(Debug, Clone, PartialEq, Eq, Display, Serialize, Deserialize)] pub enum Action { // App - Command(String, PathBuf, Vec), // (label, command, args) - Diskpart(String, String), // (label, script_as_string) + Command(PathBuf, Vec), // (command, args) + Diskpart(String), // (script_as_string) InstallDriver, Process, ScanDisks, diff --git a/src/app.rs b/src/app.rs index 68928a8..5dd523d 100644 --- a/src/app.rs +++ b/src/app.rs @@ -14,6 +14,7 @@ // along with Deja-vu. If not, see . // use std::{ + env, iter::zip, path::PathBuf, sync::{Arc, Mutex}, @@ -39,7 +40,7 @@ use crate::{ config::Config, system::{ disk::{get_disks, Disk, PartitionTableType}, - diskpart::refresh_disk_info, + diskpart::{build_dest_format_script, refresh_disk_info}, drivers::{self, Driver}, }, tui::{Event, Tui}, @@ -275,19 +276,10 @@ impl App { Action::Suspend => self.should_suspend = true, Action::Resume => self.should_suspend = false, Action::ClearScreen => tui.terminal.clear()?, - Action::Command(ref label, ref cmd_path, ref cmd_args) => { - self.action_tx.send(Action::DisplayPopup( - popup::Type::Info, - String::from(&label.to_owned()), - ))?; - self.task_handles - .push(lazy_command(&cmd_path, cmd_args.clone())) - } - Action::Diskpart(ref label, ref script) => { - self.action_tx.send(Action::DisplayPopup( - popup::Type::Info, - String::from(&label.to_owned()), - ))?; + Action::Command(ref cmd_path, ref cmd_args) => self + .task_handles + .push(lazy_command(&cmd_path, cmd_args.clone())), + Action::Diskpart(ref script) => { self.task_handles.push(lazy_diskpart(&script.to_owned())) } Action::InstallDriver => { @@ -345,26 +337,146 @@ impl App { .send(Action::UpdateDiskList(disk_list.clone()))?; } Mode::PreClone => { - // TODO: FIXME - self.action_tx.send(Action::Diskpart( + self.action_tx.send(Action::DisplayPopup( + popup::Type::Info, String::from("Formatting destination disk"), - String::from("TODO: FIXME"), ))?; + + // Build Diskpart script to format destination disk + let disk_list = self.disk_list.lock().unwrap(); + if let Some(disk_index) = self.disk_index_dest { + if let Some(disk) = disk_list.get(disk_index) { + let disk_id = disk.id.to_owned(); + let table_type = self.table_type.clone().unwrap(); + let diskpart_script = + build_dest_format_script(&disk_id, &table_type); + self.action_tx.send(Action::Diskpart(diskpart_script))?; + } + } } Mode::Clone => { - self.action_tx.send(Action::Command( + self.action_tx.send(Action::DisplayPopup( + popup::Type::Info, String::from("Running Clone Tool"), + ))?; + self.action_tx.send(Action::Command( self.config.clone_app_path.clone(), Vec::new(), ))?; } Mode::PostClone => { // TODO: FIXME - self.action_tx.send(Action::UpdateDestDisk)?; - self.action_tx.send(Action::Diskpart( - String::from("Creating boot files"), - String::from("TODO: FIXME"), + self.action_tx.send(Action::DisplayPopup( + popup::Type::Info, + String::from("Updating boot configuration"), ))?; + + // Get System32 path + let system32 = if cfg!(windows) { + match env::var("SYSTEMROOT") { + Ok(path) => path, + Err(_) => panic!("Failed to find SYSTEMROOT"), + } + } else { + String::from(".") + }; + + // Add actions + let disk_list = self.disk_list.lock().unwrap(); + if let Some(disk_index) = self.disk_index_dest { + if let Some(disk) = disk_list.get(disk_index) { + let table_type = self.table_type.clone().unwrap(); + let letter_boot = if let Some(s) = disk + .parts + .get(self.part_index_boot.unwrap()) + .clone() + .unwrap() + .letter + .clone() + { + s + } else { + String::from("??") + }; + let letter_os = if let Some(s) = disk + .parts + .get(self.part_index_os.unwrap()) + .unwrap() + .letter + .clone() + { + s + } else { + String::from("??") + }; + + // Create boot files + self.action_tx.send(Action::Command( + PathBuf::from(format!("{system32}/bcdboot.exe")), + vec![ + format!("{letter_os}:\\Windows"), + String::from("/s"), + format!("{letter_boot}:"), + String::from("/f"), + String::from(match table_type { + PartitionTableType::Guid => "UEFI", + PartitionTableType::Legacy => "BIOS", + }), + ], + ))?; + + // Update boot sector (for legacy setups) + if table_type == PartitionTableType::Legacy { + // + self.action_tx.send(Action::Command( + PathBuf::from(format!("{system32}/bootsect.exe")), + vec![ + String::from("/nt60"), + format!("{letter_boot}:"), + String::from("/force"), + String::from("/mbr"), + ], + ))?; + } + + // 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") + } + }; + self.action_tx.send(Action::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"), + ], + ))?; + + // Inject driver(s) (if selected) + if let Some(driver) = &self.driver { + if let Some(os_path) = driver.path.to_str() { + let driver_path_str = String::from(os_path); + self.action_tx.send(Action::Command( + PathBuf::from(format!("{system32}/dism.exe")), + vec![ + format!("/image:{letter_os}:\\"), + String::from("/add-driver"), + format!("/driver:\"{}\"", driver_path_str,), + String::from("/recurse"), + ], + ))?; + } + } + } + } } Mode::Done => { self.action_tx.send(Action::DisplayPopup( @@ -475,6 +587,7 @@ fn get_chunks(r: Rect) -> Vec { } fn lazy_command(cmd_path: &PathBuf, cmd_args: Vec) -> JoinHandle<()> { + info!("Running Command: {cmd_path:?} {cmd_args:?}"); if cfg!(windows) { // TODO: FIXME thread::spawn(|| sleep(Duration::from_secs(1))) @@ -488,6 +601,7 @@ fn lazy_diskpart(script: &str) -> JoinHandle<()> { // TODO: FIXME thread::spawn(|| sleep(Duration::from_secs(1))) } else { + info!("Running (lazy) Diskpart: {:?}", &script); thread::spawn(|| sleep(Duration::from_secs(1))) } } diff --git a/src/system/disk.rs b/src/system/disk.rs index 1c580e8..43e3a38 100644 --- a/src/system/disk.rs +++ b/src/system/disk.rs @@ -20,6 +20,7 @@ use std::{ thread::sleep, time::Duration, }; +use tracing::info; use once_cell::sync::Lazy; use regex::Regex; @@ -137,8 +138,10 @@ impl fmt::Display for PartitionTableType { pub fn get_disks() -> Vec { let disks: Vec; if cfg!(windows) { + info!("Get disks via Diskpart"); disks = diskpart::get_disks(); } else { + info!("Get (fake) disks"); disks = get_fake_disks(); sleep(Duration::from_millis(500)); } diff --git a/src/system/diskpart.rs b/src/system/diskpart.rs index 7e0cfec..74974ee 100644 --- a/src/system/diskpart.rs +++ b/src/system/diskpart.rs @@ -325,6 +325,7 @@ pub fn parse_partition_details(disk: &mut Disk, contents: &str) { pub fn refresh_disk_info(disk: &mut Disk) { // TODO: Needs refactor - assuming add_ functions are replaced with get_ variants + info!("Refresh disk info"); disk.parts.clear(); disk.parts_description.clear(); add_disk_details(disk, None);