deja-vu/win_installer/src/state.rs

// 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 <https://www.gnu.org/licenses/>.
//

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<usize>,
    pub disk_list: Arc<Mutex<Vec<Disk>>>,
    pub driver: Option<drivers::Driver>,
    pub driver_list: Vec<drivers::Driver>,
    pub table_type: Option<PartitionTableType>,
    pub username: Option<String>,
    pub wim_file_index: Option<usize>,
    pub wim_image_index: Option<usize>,
    pub wim_sources: Arc<Mutex<WimSources>>,
}

impl State {
    pub fn new(config: Config, disk_list: Arc<Mutex<Vec<Disk>>>) -> 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<String> {
    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<Mutex<Vec<Disk>>>,
    wim_sources_arc: Arc<Mutex<WimSources>>,
    scan_type: ScanType,
) {
    let mut to_check: Vec<String> = Vec::new();
    let mut wim_files: Vec<WimFile> = 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<Mutex<WimSources>>) {
    let result = connect_network_share(
        &config.network_server,
        &config.network_share,
        &config.network_user,
        &config.network_pass,
    );
    let mut wim_files: Vec<WimFile> = 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));
    }
}