diff --git a/config/config.json5 b/config/config.json5 index eb90ed6..90af2f8 100644 --- a/config/config.json5 +++ b/config/config.json5 @@ -190,6 +190,7 @@ }, "ScanWinSources": { "": "Process", + "": "FindWimBackups", "": "FindWimNetwork", "": "Quit", "": "Quit", @@ -200,6 +201,7 @@ "": "Process", "": "KeyUp", "": "KeyDown", + "": "PrevScreen", "": "Quit", "": "Quit", "": "Quit", diff --git a/core/src/action.rs b/core/src/action.rs index 3559c2e..e79620d 100644 --- a/core/src/action.rs +++ b/core/src/action.rs @@ -56,6 +56,7 @@ pub enum Action { Restart, Shutdown, // App (Win-Installer) + FindWimBackups, FindWimNetwork, SetUserName(String), // Screens diff --git a/win_installer/src/app.rs b/win_installer/src/app.rs index 0520d15..c15f3b1 100644 --- a/win_installer/src/app.rs +++ b/win_installer/src/app.rs @@ -47,8 +47,7 @@ use tracing::{debug, info}; use crate::{ components::{set_username::InputUsername, wim_scan::WimScan}, - state::State, - wim::WimImage, + state::{ScanType, State}, }; pub struct App { @@ -110,7 +109,19 @@ impl App { Mode::SelectTableType => Mode::ScanWinSources, Mode::ScanWinSources => Mode::SelectWinSource, Mode::SelectWinSource => Mode::SelectWinImage, - Mode::SelectWinImage => Mode::SetUserName, + Mode::SelectWinImage => { + let mut next_mode = Mode::SetUserName; + // TODO: FIXME - Race condition? + // if let Ok(wim_sources) = self.state.wim_sources.lock() + // && let Some(index) = self.state.wim_image_index + // { + // let image = wim_sources.get_file(index); + // if !image.is_installer { + // next_mode = Mode::Confirm; + // } + // } + next_mode + } Mode::SetUserName => Mode::Confirm, Mode::Confirm => Mode::Process, // i.e. format, apply, etc Mode::Process | Mode::Done => Mode::Done, @@ -150,7 +161,7 @@ impl App { String::from("Scanning Disks..."), ))?; } - Mode::ScanWinSources => self.state.scan_wim_local(), + Mode::ScanWinSources => self.state.scan_wim_local(ScanType::WindowsInstallers), Mode::Done => { self.action_tx .send(Action::DisplayPopup(popup::Type::Success, popup::fortune()))?; @@ -289,6 +300,10 @@ impl App { Action::InstallDriver => { self.action_tx.send(Action::SetMode(Mode::InstallDrivers))?; } + Action::FindWimBackups => { + self.state.reset_local(); + self.state.scan_wim_local(ScanType::GeneralWimFiles); + } Action::FindWimNetwork => { self.state.reset_network(); self.state.scan_wim_network(); @@ -302,7 +317,10 @@ impl App { Mode::SelectTableType => { self.action_tx.send(Action::SetMode(Mode::SelectDisks))?; } - Mode::SetUserName => { + Mode::SelectWinSource => { + self.action_tx.send(Action::SetMode(Mode::ScanWinSources))?; + } + Mode::SelectWinImage | Mode::SetUserName | Mode::Confirm => { self.action_tx .send(Action::SetMode(Mode::SelectWinSource))?; } @@ -413,9 +431,9 @@ fn build_footer_string(cur_mode: Mode) -> String { Mode::SelectWinSource | Mode::SelectWinImage => { String::from("(Enter) to select / (q) to quit") } - Mode::ScanWinSources => { - String::from("(Enter) to continue / (n) to scan network / (q) to quit") - } + Mode::ScanWinSources => String::from( + "(Enter) to continue / (b) to scan for backups / (n) to scan network / (q) to quit", + ), 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"), @@ -633,11 +651,11 @@ fn build_right_items(app: &App) -> Action { } Mode::SelectWinImage | Mode::SetUserName | Mode::Confirm => { info!("Building right items for: {:?}", &app.cur_mode); - let source; + let wim_file; if let Ok(wim_sources) = app.state.wim_sources.lock() && let Some(index) = app.state.wim_file_index { - source = wim_sources.get_file(index); + wim_file = wim_sources.get_file(index); } else { panic!("Failed to get source WIM file"); } @@ -675,7 +693,7 @@ fn build_right_items(app: &App) -> Action { line_colors: vec![Color::Cyan], }, DVLine { - line_parts: vec![source.path.clone()], + line_parts: vec![wim_file.path.clone()], line_colors: vec![Color::Reset], }, DVLine::blank(), @@ -688,7 +706,7 @@ fn build_right_items(app: &App) -> Action { // WIM Info match app.cur_mode { Mode::SelectWinImage => { - source.images.iter().for_each(|image| { + wim_file.images.iter().for_each(|image| { items.push(vec![DVLine { line_parts: vec![format!("{image}")], line_colors: vec![Color::Reset], @@ -697,7 +715,7 @@ fn build_right_items(app: &App) -> Action { } Mode::Confirm => { if let Some(index) = app.state.wim_image_index - && let Some(image) = source.images.get(index) + && let Some(image) = wim_file.images.get(index) { label_dv_lines.append(&mut vec![ DVLine { @@ -707,7 +725,9 @@ fn build_right_items(app: &App) -> Action { DVLine::blank(), ]); } - if let Some(username) = &app.state.username { + if wim_file.is_installer + && let Some(username) = &app.state.username + { label_dv_lines.append(&mut vec![DVLine { line_parts: vec![String::from("Username: "), username.clone()], line_colors: vec![Color::Green, Color::Reset], @@ -719,7 +739,6 @@ fn build_right_items(app: &App) -> Action { } // Done - info!("label_dv_lines: {:?}", &label_dv_lines); labels.push(label_dv_lines); } Mode::SelectTableType => { diff --git a/win_installer/src/state.rs b/win_installer/src/state.rs index 68d9ebd..98428c3 100644 --- a/win_installer/src/state.rs +++ b/win_installer/src/state.rs @@ -20,8 +20,6 @@ use std::{ }; use core::{ - action::Action, - components::popup::Type as PopupType, config::Config, system::{ disk::{Disk, PartitionTableType}, @@ -33,9 +31,14 @@ use tracing::info; use crate::{ net::connect_network_share, - wim::{WimSources, parse_wim_file}, + wim::{WimFile, WimSources, parse_wim_file}, }; +pub enum ScanType { + GeneralWimFiles, // Includes Windows installer WIMs + WindowsInstallers, +} + #[derive(Debug, Default)] pub struct State { pub config: Config, @@ -70,6 +73,12 @@ impl State { } } + pub fn reset_local(&mut self) { + if let Ok(mut sources) = self.wim_sources.lock() { + sources.reset_local(); + } + } + pub fn reset_network(&mut self) { if let Ok(mut sources) = self.wim_sources.lock() { sources.reset_network(); @@ -80,11 +89,11 @@ impl State { self.driver_list = drivers::scan(); } - pub fn scan_wim_local(&mut self) { + pub fn scan_wim_local(&mut self, scan_type: ScanType) { let disk_list_arc = self.disk_list.clone(); let wim_sources_arc = self.wim_sources.clone(); tokio::task::spawn(async move { - scan_local_drives(disk_list_arc, wim_sources_arc); + scan_local_drives(disk_list_arc, wim_sources_arc, scan_type); }); } @@ -97,18 +106,39 @@ impl State { } } +fn get_subfolders(path_str: &str) -> Vec { + if let Ok(read_dir) = read_dir(path_str) { + read_dir + .filter_map(|item| item.ok()) + .map(|item| item.path().to_string_lossy().into_owned()) + .collect() + } else { + // TODO: Use better error handling here? + Vec::new() + } +} + pub fn scan_local_drives( disk_list_arc: Arc>>, wim_sources_arc: Arc>, + scan_type: ScanType, ) { let mut to_check = vec![String::from(".")]; + let mut wim_files: Vec = Vec::new(); // Get drive letters if let Ok(disk_list) = disk_list_arc.lock() { disk_list.iter().for_each(|d| { d.parts.iter().for_each(|p| { if !p.letter.is_empty() { - to_check.push(format!("{}:\\Images", &p.letter)); + match scan_type { + ScanType::GeneralWimFiles => { + to_check.append(&mut get_subfolders(&format!("{}:\\", &p.letter))); + } + ScanType::WindowsInstallers => { + to_check.push(format!("{}:\\Images", &p.letter)); + } + } } }); }) @@ -116,20 +146,28 @@ pub fn scan_local_drives( // Scan drives to_check.iter().for_each(|scan_path| { + let installer = scan_path.ends_with("\\Images"); info!("Scanning: {}", &scan_path); if let Ok(read_dir) = read_dir(scan_path) { read_dir.for_each(|item| { if let Ok(item) = item && item.file_name().to_string_lossy().ends_with(".wim") && let Some(path_str) = item.path().to_str() - && let Ok(new_source) = parse_wim_file(path_str) - && let Ok(mut wim_sources) = wim_sources_arc.lock() + && let Ok(new_source) = parse_wim_file(path_str, installer) { - wim_sources.add_local(new_source); + wim_files.push(new_source); } }); } }); + + // Done + wim_files.sort(); + if let Ok(mut wim_sources) = wim_sources_arc.lock() { + wim_files + .into_iter() + .for_each(|file| wim_sources.add_local(file)); + } } pub fn scan_network_share(config: Config, wim_sources_arc: Arc>) { @@ -139,6 +177,7 @@ pub fn scan_network_share(config: Config, wim_sources_arc: Arc &config.network_user, &config.network_pass, ); + let mut wim_files: Vec = Vec::new(); // Connect to share if result.is_err() { @@ -152,14 +191,19 @@ pub fn scan_network_share(config: Config, wim_sources_arc: Arc if let Ok(item) = item && item.file_name().to_string_lossy().ends_with(".wim") && let Some(path_str) = item.path().to_str() - && let Ok(new_source) = parse_wim_file(path_str) - && let Ok(mut wim_sources) = wim_sources_arc.lock() + && let Ok(new_source) = parse_wim_file(path_str, true) + // Assuming all network sources are installers { - wim_sources.add_network(new_source); + wim_files.push(new_source); } }); } // Done - let _ = 14; + wim_files.sort(); + if let Ok(mut wim_sources) = wim_sources_arc.lock() { + wim_files + .into_iter() + .for_each(|file| wim_sources.add_network(file)); + } } diff --git a/win_installer/src/wim.rs b/win_installer/src/wim.rs index b971569..032ff86 100644 --- a/win_installer/src/wim.rs +++ b/win_installer/src/wim.rs @@ -66,6 +66,7 @@ static WIN_BUILDS: LazyLock> = LazyLock::new(|| { pub struct WimFile { pub path: String, pub images: Vec, + pub is_installer: bool, } impl WimFile { @@ -154,12 +155,10 @@ impl WimSources { pub fn add_local(&mut self, wim_file: WimFile) { self.local.push(wim_file); - self.local.sort(); } pub fn add_network(&mut self, wim_file: WimFile) { self.network.push(wim_file); - self.network.sort(); } pub fn get_file(&self, index: usize) -> WimFile { @@ -183,6 +182,10 @@ impl WimSources { self.network.clear(); } + pub fn reset_local(&mut self) { + self.local.clear(); + } + pub fn reset_network(&mut self) { self.network.clear(); } @@ -203,7 +206,7 @@ fn get_wim_xml(wim_file: &str) -> std::io::Result { Ok(file) } -pub fn parse_wim_file(wim_file: &str) -> std::io::Result { +pub fn parse_wim_file(wim_file: &str, installer: bool) -> std::io::Result { let mut wim_images: Vec = Vec::new(); if !Path::new(wim_file).exists() { return Err(std::io::Error::new( @@ -268,6 +271,7 @@ pub fn parse_wim_file(wim_file: &str) -> std::io::Result { let wim_file = WimFile { path: wim_file.to_string(), images: wim_images, + is_installer: installer, }; Ok(wim_file)