// 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 std::{ fs::read_dir, sync::{Arc, Mutex}, }; use core::{ config::Config, system::{ disk::{Disk, PartitionTableType}, drivers, }, }; use crate::{ net::connect_network_share, wim::{WimFile, WimSources, parse_wim_file}, }; pub enum ScanType { GeneralWimFiles, // Includes Windows installer WIMs WindowsInstallers, } #[derive(Debug, Default)] pub struct State { pub config: Config, pub disk_index_dest: Option, pub disk_list: Arc>>, pub driver: Option, pub driver_list: Vec, pub table_type: Option, pub username: Option, pub wim_file_index: Option, pub wim_image_index: Option, pub wim_sources: Arc>, } impl State { pub fn new(config: Config, disk_list: Arc>>) -> Self { let wim_sources = Arc::new(Mutex::new(WimSources::new())); State { config, disk_list, wim_sources, ..Default::default() } } pub fn reset_all(&mut self) { self.disk_index_dest = None; self.wim_file_index = None; self.wim_image_index = None; if let Ok(mut sources) = self.wim_sources.lock() { sources.reset_all(); } } 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(); } } pub fn scan_drivers(&mut self) { self.driver_list = drivers::scan(); } 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(); let wim_sources_arc_inner = self.wim_sources.clone(); if let Ok(mut wim_sources) = wim_sources_arc.lock() && wim_sources.thread_local.is_none() { wim_sources.thread_local = Some(tokio::task::spawn(async move { scan_local_drives(disk_list_arc, wim_sources_arc_inner, scan_type); })); } } pub fn scan_wim_network(&mut self) { let wim_sources_arc = self.wim_sources.clone(); let wim_sources_arc_inner = self.wim_sources.clone(); if let Ok(mut wim_sources) = wim_sources_arc.lock() && wim_sources.thread_network.is_none() { let config = self.config.clone(); wim_sources.thread_network = Some(tokio::task::spawn(async move { scan_network_share(config, wim_sources_arc_inner); })); } } } 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 = Vec::new(); 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() { match scan_type { ScanType::GeneralWimFiles => { to_check.append(&mut get_subfolders(&format!("{}:\\", &p.letter))); } ScanType::WindowsInstallers => { to_check.push(format!("{}:\\Images", &p.letter)); } } } }); }) } // Scan drives to_check.iter().for_each(|scan_path| { let is_backup = !scan_path.ends_with("\\Images"); 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, is_backup) { 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>) { let result = connect_network_share( &config.network_server, &config.network_share, &config.network_user, &config.network_pass, ); let mut wim_files: Vec = Vec::new(); // Connect to share if result.is_err() { return; } // Scan share let share_dir = format!("\\\\{}\\{}", &config.network_server, &config.network_share); if let Ok(read_dir) = read_dir(share_dir) { 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, true) // Assuming all network sources are installers { 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_network(file)); } }