diff --git a/core/src/system/diskpart.rs b/core/src/system/diskpart.rs index 98fc74d..b1281a8 100644 --- a/core/src/system/diskpart.rs +++ b/core/src/system/diskpart.rs @@ -32,6 +32,12 @@ use crate::system::disk::{ static DEFAULT_MAX_DISKS: usize = 8; +#[derive(Debug, PartialEq)] +pub enum FormatUseCase { + ApplyWimImage, + Clone, +} + pub struct RegexList { detail_all_disks: OnceLock, detail_disk: OnceLock, @@ -205,7 +211,11 @@ pub fn get_partitions( } #[must_use] -pub fn build_dest_format_script(disk_id: usize, part_type: &PartitionTableType) -> String { +pub fn build_dest_format_script( + disk_id: usize, + part_type: &PartitionTableType, + format_use_case: FormatUseCase, +) -> String { let disk_id = format!("{disk_id}"); let mut script = vec!["automount enable noerr", "select disk {disk_id}", "clean"]; match part_type { @@ -221,6 +231,10 @@ pub fn build_dest_format_script(disk_id: usize, part_type: &PartitionTableType) script.push("format fs=ntfs quick label=System"); } } + if format_use_case == FormatUseCase::ApplyWimImage { + script.push("create partition primary"); + script.push("format fs=ntfs quick label=Windows"); + } script.join("\r\n").replace("{disk_id}", &disk_id) } diff --git a/deja_vu/src/app.rs b/deja_vu/src/app.rs index 450f423..dc8ccd6 100644 --- a/deja_vu/src/app.rs +++ b/deja_vu/src/app.rs @@ -28,7 +28,10 @@ use core::{ line::{DVLine, get_disk_description_right, get_part_description}, state::Mode, system::{ - boot, cpu::get_cpu_name, disk::PartitionTableType, diskpart::build_dest_format_script, + boot, + cpu::get_cpu_name, + disk::PartitionTableType, + diskpart::{FormatUseCase, build_dest_format_script}, drivers, }, tasks::{Task, TaskResult, TaskType, Tasks}, @@ -209,7 +212,8 @@ impl App { && let Some(disk) = disk_list.get(disk_index) { let table_type = self.state.table_type.clone().unwrap(); - let diskpart_script = build_dest_format_script(disk.id, &table_type); + let diskpart_script = + build_dest_format_script(disk.id, &table_type, FormatUseCase::Clone); self.tasks.add(TaskType::Diskpart(diskpart_script)); } } diff --git a/win_installer/src/app.rs b/win_installer/src/app.rs index b7f7178..0a56f54 100644 --- a/win_installer/src/app.rs +++ b/win_installer/src/app.rs @@ -26,12 +26,20 @@ use core::{ config::Config, line::{DVLine, get_disk_description_right}, state::Mode, - system::{cpu::get_cpu_name, disk::PartitionTableType, drivers}, - tasks::{TaskType, Tasks}, + system::{ + boot, + cpu::get_cpu_name, + disk::PartitionTableType, + diskpart::{FormatUseCase, build_dest_format_script}, + drivers, + }, + tasks::{Task, TaskResult, TaskType, Tasks}, tui::{Event, Tui}, }; use std::{ + env, iter::zip, + path::PathBuf, sync::{Arc, Mutex}, }; @@ -111,20 +119,20 @@ impl App { Mode::SelectWinSource => Mode::SelectWinImage, Mode::SelectWinImage => Mode::SetUserName, Mode::SetUserName => Mode::Confirm, - Mode::Confirm => Mode::Process, // i.e. format, apply, etc - Mode::Process | Mode::Done => Mode::Done, + Mode::Confirm => Mode::PreClone, + Mode::PreClone => Mode::Clone, + Mode::Clone => Mode::PostClone, + Mode::PostClone | Mode::Done => Mode::Done, Mode::Failed => Mode::Failed, // Invalid States Mode::BootDiags | Mode::BootScan | Mode::BootSetup - | Mode::Clone | Mode::DiagMenu | Mode::InjectDrivers | Mode::LogView | Mode::PEMenu - | Mode::PreClone - | Mode::PostClone + | Mode::Process | Mode::SelectParts | Mode::SetBootMode => panic!("This shouldn't happen?"), } @@ -139,6 +147,126 @@ impl App { // self.action_tx // .send(Action::DisplayPopup(popup::Type::Info, String::from("...")))?; // } + Mode::PreClone => { + self.action_tx.send(Action::DisplayPopup( + popup::Type::Info, + String::from("Formatting destination disk"), + ))?; + + // Get System32 path + let system32 = get_system32_path(&self.action_tx); + + // (Re)Enable volume mounting + self.tasks.add(TaskType::CommandWait( + PathBuf::from(format!("{system32}/mountvol.exe")), + vec![String::from("/e")], + )); + + // Build Diskpart script to format destination disk + let disk_list = self.state.disk_list.lock().unwrap(); + if let Some(disk_index) = self.state.disk_index_dest + && let Some(disk) = disk_list.get(disk_index) + { + let table_type = self.state.table_type.clone().unwrap(); + let diskpart_script = build_dest_format_script( + disk.id, + &table_type, + FormatUseCase::ApplyWimImage, + ); + self.tasks.add(TaskType::Diskpart(diskpart_script)); + } + + // Update drive letters + self.tasks.add(TaskType::Sleep); + if let Some(dest_index) = self.state.disk_index_dest { + self.tasks.add(TaskType::UpdateDestDisk(dest_index)); + } + } + Mode::Clone => { + self.action_tx.send(Action::DisplayPopup( + popup::Type::Info, + String::from("Applying Image"), + ))?; + + // Get wimlib-imagex path + let program_files = PathBuf::from(get_program_files_path(&self.action_tx)); + let wimlib_imagex = program_files.join("wimlib\\wimlib-imagex.exe"); + + // Get image info + let wim_sources = self.state.wim_sources.lock().unwrap(); + let wim_file = wim_sources.get_file(self.state.wim_file_index.unwrap()); + let wim_index = format!("{}", self.state.wim_image_index.unwrap()); + + // Add actions + let disk_list = self.state.disk_list.lock().unwrap(); + if let Some(disk_index) = self.state.disk_index_dest + && let Some(disk) = disk_list.get(disk_index) + { + let num_parts = disk.parts.len(); + let dest_path = format!("{}:\\", disk.get_part_letter(num_parts - 1)); + self.tasks.add(TaskType::CommandWait( + wimlib_imagex, + vec![String::from("apply"), wim_file.path, wim_index, dest_path], + )); + } + } + Mode::PostClone => { + self.action_tx.send(Action::DisplayPopup( + popup::Type::Info, + String::from("Updating boot configuration"), + ))?; + + // Get System32 path + let system32 = get_system32_path(&self.action_tx); + + // Add actions + let disk_list = self.state.disk_list.lock().unwrap(); + if let Some(disk_index) = self.state.disk_index_dest + && let Some(disk) = disk_list.get(disk_index) + { + let table_type = self.state.table_type.clone().unwrap(); + let letter_boot = disk.get_part_letter(0); + let letter_os = + disk.get_part_letter(match self.state.table_type.clone().unwrap() { + PartitionTableType::Guid => 2, + PartitionTableType::Legacy => 1, + }); + info!("PostClone // Disk: {disk:?}"); + info!("\t\tBoot letter: {letter_boot}"); + info!("\t\tOS letter: {letter_os}"); + + // 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, + boot::SafeMode::Enable, + &system32, + &table_type, + ) { + self.tasks.add(task); + } + + // Inject driver(s) (if selected) + if let Some(driver) = &self.state.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::ScanDisks => { self.state.reset_all(); if self.tasks.idle() { @@ -261,7 +389,9 @@ impl App { Action::Tick => { self.last_tick_key_events.drain(..); // Check background task(s) - self.tasks.poll()?; + if let Some(task) = self.tasks.poll()? { + self.handle_task(&task)?; + } if let Ok(mut wim_sources) = self.state.wim_sources.lock() { wim_sources.poll(); } @@ -317,6 +447,11 @@ impl App { self.action_tx.send(Action::DismissPopup)?; self.action_tx.send(Action::SetMode(next_mode))?; } + Action::DisplayPopup(ref popup_type, ref _popup_text) => { + if *popup_type == popup::Type::Error { + self.action_tx.send(Action::SetMode(Mode::Failed))?; + } + } Action::PrevScreen => match self.cur_mode { Mode::SelectTableType => { self.action_tx.send(Action::SetMode(Mode::SelectDisks))?; @@ -405,6 +540,39 @@ impl App { Ok(()) } + fn handle_task(&mut self, task: &Task) -> Result<()> { + match task.task_type { + TaskType::CommandWait(_, _) | TaskType::Diskpart(_) => { + // Check result + if let Some(result) = &task.result { + match result { + TaskResult::Error(msg) => { + self.action_tx + .send(Action::Error(format!("{} Failed: {msg}", task.task_type)))?; + } + TaskResult::Output(stdout, stderr, success) => { + if !success { + let msg = if !stdout.is_empty() { + stdout.clone() + } else if !stderr.is_empty() { + stderr.clone() + } else { + String::from("Unknown Error") + }; + self.action_tx.send(Action::Error(format!( + "{} Failed: {msg}", + task.task_type + )))?; + } + } + } + } + } + _ => {} + } + Ok(()) + } + fn render(&mut self, tui: &mut Tui) -> Result<()> { tui.draw(|frame| { if let [header, _body, footer, center, left, right, username, popup] = @@ -426,7 +594,9 @@ impl App { fn build_footer_string(cur_mode: Mode) -> String { match cur_mode { - Mode::Home | Mode::ScanDisks => String::from("(q) to quit"), + Mode::Home | Mode::ScanDisks | Mode::PreClone | Mode::Clone | Mode::PostClone => { + String::from("(q) to quit") + } Mode::InstallDrivers => String::from("(Enter) to select / (q) to quit"), Mode::SelectDisks => String::from( "(Enter) to select / / (i) to install driver / (r) to rescan / (q) to quit", @@ -440,18 +610,16 @@ fn build_footer_string(cur_mode: Mode) -> String { ), Mode::SetUserName => String::from("(Enter) to continue / (Esc) to go back"), Mode::Confirm => String::from("(Enter) to confirm / (b) to go back / (q) to quit"), - Mode::Done | Mode::Failed | Mode::Process => String::from("(Enter) or (q) to quit"), + Mode::Done | Mode::Failed => String::from("(Enter) or (q) to quit"), // Invalid States Mode::BootDiags | Mode::BootScan | Mode::BootSetup - | Mode::Clone | Mode::DiagMenu | Mode::InjectDrivers | Mode::LogView | Mode::PEMenu - | Mode::PreClone - | Mode::PostClone + | Mode::Process | Mode::SelectParts | Mode::SetBootMode => panic!("This shouldn't happen?"), } @@ -475,15 +643,9 @@ fn build_left_items(app: &App) -> Action { .iter() .for_each(|driver| items.push(driver.to_string())); } - Mode::Process => { - select_type = SelectionType::Loop; - title = String::from("Processing"); - // TODO: FIXME - } Mode::ScanWinSources => { select_type = SelectionType::Loop; title = String::from("Scanning"); - // TODO: FIXME } Mode::SelectWinSource => { select_type = SelectionType::One; @@ -539,6 +701,7 @@ fn build_left_items(app: &App) -> Action { | Mode::BootScan | Mode::BootSetup | Mode::DiagMenu + | Mode::Process | Mode::InjectDrivers | Mode::LogView | Mode::PEMenu @@ -653,7 +816,12 @@ fn build_right_items(app: &App) -> Action { }); } } - Mode::SelectWinImage | Mode::SetUserName | Mode::Confirm => { + Mode::SelectWinImage + | Mode::SetUserName + | Mode::Confirm + | Mode::PreClone + | Mode::Clone + | Mode::PostClone => { info!("Building right items for: {:?}", &app.cur_mode); let wim_file; if let Ok(wim_sources) = app.state.wim_sources.lock() @@ -717,7 +885,7 @@ fn build_right_items(app: &App) -> Action { }]) }); } - Mode::Confirm => { + Mode::Confirm | Mode::PreClone | Mode::Clone | Mode::PostClone => { if let Some(index) = app.state.wim_image_index && let Some(image) = wim_file.images.get(index) { @@ -842,3 +1010,35 @@ fn get_chunks(r: Rect) -> Vec { // Done chunks } + +pub fn get_program_files_path(action_tx: &mpsc::UnboundedSender) -> String { + let mut program_files_path = String::from("."); + if cfg!(windows) { + if let Ok(path) = env::var("PROGRAMFILES") { + program_files_path = path; + } else { + action_tx + .send(Action::Error(String::from( + "ERROR\n\n\nFailed to find PROGRAMFILES", + ))) + .expect("Failed to find PROGRAMFILES and then failed to send action"); + } + } + program_files_path +} + +pub fn get_system32_path(action_tx: &mpsc::UnboundedSender) -> String { + let mut system32_path = String::from("."); + if cfg!(windows) { + if let Ok(path) = env::var("SYSTEMROOT") { + system32_path = format!("{path}/System32"); + } else { + action_tx + .send(Action::Error(String::from( + "ERROR\n\n\nFailed to find SYSTEMROOT", + ))) + .expect("Failed to find SYSTEMROOT and then failed to send action"); + } + } + system32_path +}