diff --git a/deja_vu/src/app.rs b/deja_vu/src/app.rs index a5ba031..f68b066 100644 --- a/deja_vu/src/app.rs +++ b/deja_vu/src/app.rs @@ -16,7 +16,6 @@ use std::{ env, iter::zip, - path::PathBuf, sync::{Arc, Mutex}, }; @@ -37,6 +36,7 @@ use crate::{ }, config::Config, system::{ + boot, disk::{Disk, PartitionTableType}, diskpart::build_dest_format_script, drivers::{self, Driver}, @@ -163,6 +163,113 @@ impl App { } } + pub fn set_mode(&mut self, new_mode: Mode) -> Result<()> { + info!("Setting mode to {new_mode:?}"); + self.cur_mode = new_mode; + match new_mode { + Mode::ScanDisks => { + self.prev_mode = self.cur_mode; + if self.tasks.idle() { + self.tasks.add(Task::ScanDisks); + } + } + Mode::PreClone => { + self.action_tx.send(Action::DisplayPopup( + popup::Type::Info, + String::from("Formatting destination disk"), + ))?; + + // 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 table_type = self.table_type.clone().unwrap(); + let diskpart_script = build_dest_format_script(disk.id, &table_type); + self.tasks.add(Task::Diskpart(diskpart_script)); + } + } + } + Mode::Clone => { + self.action_tx.send(Action::DisplayPopup( + popup::Type::Info, + String::from("Running Clone Tool"), + ))?; + self.tasks.add(Task::Command( + self.config.clone_app_path.clone(), + Vec::new(), + )); + if let Some(dest_index) = self.disk_index_dest { + self.tasks.add(Task::UpdateDestDisk(dest_index)); + } + } + Mode::PostClone => { + self.action_tx.send(Action::DisplayPopup( + popup::Type::Info, + String::from("Updating boot configuration"), + ))?; + + // Get System32 path + let 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(".") + }; + + // 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 = disk.get_part_letter(self.part_index_boot.unwrap()); + let letter_os = disk.get_part_letter(self.part_index_os.unwrap()); + + // Safety check + if letter_boot.is_empty() || letter_os.is_empty() { + self.action_tx.send(Action::Error(String::from( + "ERROR\n\n\nFailed to get drive letters for the destination", + )))?; + return Ok(()); + } + + // Create boot files + for task in + boot::configure_disk(&letter_boot, &letter_os, &system32, table_type) + { + self.tasks.add(task); + } + + // Inject driver(s) (if selected) + if let Some(driver) = &self.driver { + if let Ok(task) = boot::inject_driver(&driver, &letter_os, &system32) { + self.tasks.add(task); + } else { + self.action_tx.send(Action::Error(format!( + "Failed to inject driver:\n{}", + driver.name + )))?; + } + } + } + } + } + Mode::Done => { + self.action_tx.send(Action::DisplayPopup( + popup::Type::Success, + String::from("COMPLETE\n\n\nThank you for using this tool!"), + ))?; + } + _ => {} + } + Ok(()) + } + pub async fn run(&mut self) -> Result<()> { let mut tui = Tui::new()? // .mouse(true) // uncomment this line to enable mouse support @@ -282,7 +389,11 @@ impl App { Action::Resize(w, h) => self.handle_resize(tui, w, h)?, Action::Render => self.render(tui)?, Action::PrevScreen => { - self.action_tx.send(Action::SetMode(self.prev_mode))?; + let new_mode = match (self.prev_mode, self.cur_mode) { + (Mode::SelectTableType, Mode::SelectTableType) => Mode::SelectDisks, + (_, _) => self.prev_mode, + }; + self.action_tx.send(Action::SetMode(new_mode))?; } Action::NextScreen => { if let Some(mode) = self.next_mode() { @@ -306,160 +417,7 @@ impl App { self.selections[0] = one; self.selections[1] = two; } - Action::SetMode(new_mode) => { - info!("Setting mode to {new_mode:?}"); - self.cur_mode = new_mode; - match new_mode { - Mode::ScanDisks => { - self.prev_mode = self.cur_mode; - if self.tasks.idle() { - self.tasks.add(Task::ScanDisks); - } - } - Mode::PreClone => { - self.action_tx.send(Action::DisplayPopup( - popup::Type::Info, - String::from("Formatting destination disk"), - ))?; - - // 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 table_type = self.table_type.clone().unwrap(); - let diskpart_script = - build_dest_format_script(disk.id, &table_type); - self.tasks.add(Task::Diskpart(diskpart_script)); - } - } - } - Mode::Clone => { - self.action_tx.send(Action::DisplayPopup( - popup::Type::Info, - String::from("Running Clone Tool"), - ))?; - self.tasks.add(Task::Command( - self.config.clone_app_path.clone(), - Vec::new(), - )); - if let Some(dest_index) = self.disk_index_dest { - self.tasks.add(Task::UpdateDestDisk(dest_index)); - } - } - Mode::PostClone => { - self.action_tx.send(Action::DisplayPopup( - popup::Type::Info, - String::from("Updating boot configuration"), - ))?; - - // Get System32 path - let 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(".") - }; - - // 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 = - disk.get_part_letter(self.part_index_boot.unwrap()); - let letter_os = - disk.get_part_letter(self.part_index_os.unwrap()); - - // Safety check - if letter_boot.is_empty() || letter_os.is_empty() { - self.action_tx.send(Action::Error(String::from( - "ERROR\n\n\nFailed to get drive letters for the destination", - )))?; - return Ok(()); - } - - // Create boot files - self.tasks.add(Task::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.tasks.add(Task::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.tasks.add(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"), - ], - )); - - // 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.tasks.add(Task::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( - popup::Type::Success, - String::from("COMPLETE\n\n\nThank you for using this tool!"), - ))?; - } - _ => {} - } - } + Action::SetMode(new_mode) => self.set_mode(new_mode)?, _ => {} } for component in &mut self.components { diff --git a/deja_vu/src/system.rs b/deja_vu/src/system.rs index a923b3e..1d9692d 100644 --- a/deja_vu/src/system.rs +++ b/deja_vu/src/system.rs @@ -13,6 +13,7 @@ // You should have received a copy of the GNU General Public License // along with Deja-vu. If not, see . // +pub mod boot; pub mod cpu; pub mod disk; pub mod diskpart; diff --git a/deja_vu/src/system/boot.rs b/deja_vu/src/system/boot.rs new file mode 100644 index 0000000..2bc098f --- /dev/null +++ b/deja_vu/src/system/boot.rs @@ -0,0 +1,94 @@ +// This file is part of Deja-vu. +// +// Deja-vu is free software: you can redistribute it and/or modify it +// under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// Deja-vu is distributed in the hope that it will be useful, but +// WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. +// See the GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with Deja-vu. If not, see . +// +use super::{disk::PartitionTableType, drivers::Driver}; +use crate::tasks::Task; +use color_eyre::Result; +use std::path::PathBuf; + +pub fn configure_disk( + letter_boot: &str, + letter_os: &str, + system32: &str, + table_type: PartitionTableType, +) -> Vec { + let mut tasks = Vec::new(); + + // Create + tasks.push(Task::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 { + tasks.push(Task::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") + } + }; + 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"), + ], + )); + + // Done + tasks +} + +pub fn inject_driver(driver: &Driver, letter_os: &str, system32: &str) -> Result { + //if let Some(driver_path_str) = driver.path.to_str() { + let driver_path = driver.path.to_str().unwrap(); + Ok(Task::Command( + PathBuf::from(format!("{system32}/dism.exe")), + vec![ + format!("/image:{letter_os}:\\"), + String::from("/add-driver"), + format!("/driver:\"{}\"", driver_path,), + String::from("/recurse"), + ], + )) +} diff --git a/deja_vu/src/system/cpu.rs b/deja_vu/src/system/cpu.rs index bdca0b2..1187d2c 100644 --- a/deja_vu/src/system/cpu.rs +++ b/deja_vu/src/system/cpu.rs @@ -13,8 +13,6 @@ // You should have received a copy of the GNU General Public License // along with Deja-vu. If not, see . // - -/// CPU Functions #[must_use] pub fn get_cpu_name() -> String { let cpuid = raw_cpuid::CpuId::new();