From cf3d71567e64062882bf7dc8191cb1e6eee80803 Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Fri, 28 Nov 2025 19:25:18 -0800 Subject: [PATCH] Refactor scan_local_drives() --- Cargo.toml | 3 ++ win_installer/src/app.rs | 10 ++---- win_installer/src/state.rs | 68 ++++++++++++++++++++++++++++++++++---- win_installer/src/wim.rs | 56 ++++++++----------------------- 4 files changed, 81 insertions(+), 56 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 0a6700c..e166af2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,3 +17,6 @@ members = ["core", "boot_diags", "deja_vu", "pe_menu", "win_installer"] default-members = ["core", "boot_diags", "deja_vu", "pe_menu", "win_installer"] resolver = "2" + +[profile.release] +lto = true diff --git a/win_installer/src/app.rs b/win_installer/src/app.rs index 7d312f5..71c249d 100644 --- a/win_installer/src/app.rs +++ b/win_installer/src/app.rs @@ -45,7 +45,7 @@ use ratatui::{ use tokio::sync::mpsc; use tracing::{debug, info}; -use crate::{components::wim_scan::WimScan, state::State, wim::scan_local_drives}; +use crate::{components::wim_scan::WimScan, state::State}; pub struct App { // TUI @@ -143,13 +143,7 @@ impl App { String::from("Scanning Disks..."), ))?; } - Mode::ScanWinImages => { - let disk_list_arc = self.state.disk_list.clone(); - let wim_sources_arc = self.state.wim_sources.clone(); - tokio::task::spawn(async move { - scan_local_drives(disk_list_arc, wim_sources_arc); - }); - } + Mode::ScanWinImages => self.state.scan_wim_local(), Mode::Done => { self.action_tx .send(Action::DisplayPopup(popup::Type::Success, popup::fortune()))?; diff --git a/win_installer/src/state.rs b/win_installer/src/state.rs index d85f8ba..1f56254 100644 --- a/win_installer/src/state.rs +++ b/win_installer/src/state.rs @@ -14,14 +14,27 @@ // along with Deja-Vu. If not, see . // -use std::sync::{Arc, Mutex}; - -use core::system::{ - disk::{Disk, PartitionTableType}, - drivers, +use std::{ + fs::read_dir, + sync::{Arc, Mutex}, }; -use crate::wim::WimSources; +use core::{ + action::Action, + components::popup::Type as PopupType, + config::Config, + system::{ + disk::{Disk, PartitionTableType}, + drivers, + }, +}; + +use tracing::info; + +use crate::{ + net::connect_network_share, + wim::{WimSources, parse_wim_file}, +}; #[derive(Debug, Default)] pub struct State { @@ -63,4 +76,47 @@ impl State { pub fn scan_drivers(&mut self) { self.driver_list = drivers::scan(); } + + pub fn scan_wim_local(&mut self) { + 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); + }); + } + +pub fn scan_local_drives( + disk_list_arc: Arc>>, + wim_sources_arc: Arc>, +) { + let mut to_check = vec![String::from(".")]; + + // 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)); + } + }); + }) + } + + // Scan drives + to_check.iter().for_each(|scan_path| { + 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() + { + wim_sources.local.push(new_source); + } + }); + } + }); +} } diff --git a/win_installer/src/wim.rs b/win_installer/src/wim.rs index 7065d9c..b6bc057 100644 --- a/win_installer/src/wim.rs +++ b/win_installer/src/wim.rs @@ -1,4 +1,3 @@ -use core::system::disk::Disk; // This file is part of Deja-Vu. // // Deja-Vu is free software: you can redistribute it and/or modify it @@ -16,18 +15,26 @@ use core::system::disk::Disk; // use std::{ collections::HashMap, - fmt, - fs::{File, read_dir}, + env, fmt, + fs::File, io::BufReader, - path::Path, + path::{Path, PathBuf}, process::Command, - sync::{Arc, LazyLock, Mutex}, + sync::LazyLock, }; use tempfile::NamedTempFile; -use tracing::info; use xml::reader::{EventReader, XmlEvent}; +static WIMINFO_EXE: LazyLock = LazyLock::new(|| { + let program_files = + PathBuf::from(env::var("PROGRAMFILES").expect("Failed to resolve %PROGRAMFILES%")); + program_files + .join("wimlib/wiminfo.cmd") + .to_string_lossy() + .into_owned() +}); + static WIN_BUILDS: LazyLock> = LazyLock::new(|| { HashMap::from([ // Windows 10 @@ -118,7 +125,7 @@ impl WimSources { fn get_wim_xml(wim_file: &str) -> std::io::Result { let tmp_file = NamedTempFile::new()?; - let _ = Command::new("wiminfo") + let _ = Command::new(&*WIMINFO_EXE) .args([ wim_file, "--extract-xml", @@ -200,38 +207,3 @@ pub fn parse_wim_file(wim_file: &str) -> std::io::Result { Ok(wim_file) } - -pub fn scan_local_drives( - disk_list_arc: Arc>>, - wim_sources_arc: Arc>, -) { - let mut to_check = vec![String::from(".")]; - - // 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)); - } - }); - }) - } - - // Scan drives - to_check.iter().for_each(|scan_path| { - 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() - { - wim_sources.local.push(new_source); - } - }); - } - }); -}