deja-vu/win_installer/src/wim.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::{
    collections::HashMap, fmt, fs::File, io::BufReader, path::Path, process::Command,
    sync::LazyLock,
};

use tempfile::NamedTempFile;
use xml::reader::{EventReader, XmlEvent};

static WIN_BUILDS: LazyLock<HashMap<&str, &str>> = LazyLock::new(|| {
    HashMap::from([
        // Windows 10
        ("10240", "1507 \"Threshold 1\""),
        ("10586", "1511 \"Threshold 2\""),
        ("14393", "1607 \"Redstone 1\""),
        ("15063", "1703 \"Redstone 2\""),
        ("16299", "1709 \"Redstone 3\""),
        ("17134", "1803 \"Redstone 4\""),
        ("17763", "1809 \"Redstone 5\""),
        ("18362", "1903 / 19H1"),
        ("18363", "1909 / 19H2"),
        ("19041", "2004 / 20H1"),
        ("19042", "20H2"),
        ("19043", "21H1"),
        ("19044", "21H2"),
        ("19045", "22H2"),
        // Windows 11
        ("22000", "21H2"),
        ("22621", "22H2"),
        ("22631", "23H2"),
        ("26100", "24H2"),
        ("26200", "25H2"),
    ])
});

#[derive(Debug)]
pub struct WimFile<'a> {
    path: &'a Path,
    images: Vec<WimImage>,
}

#[derive(Clone, Debug, Default)]
pub struct WimImage {
    build: String,
    index: String,
    name: String,
    spbuild: String,
    version: String,
}

impl WimImage {
    pub fn new() -> Self {
        Default::default()
    }

    pub fn reset(&mut self) {
        self.build.clear();
        self.index.clear();
        self.name.clear();
        self.spbuild.clear();
        self.version.clear();
    }
}

impl fmt::Display for WimImage {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // Windows 11 Home (24H2, 26100.xxxx)
        let s = if self.version.is_empty() {
            String::new()
        } else {
            format!("{}, ", self.version)
        };
        write!(f, "{} ({}{}.{})", self.name, s, self.build, self.spbuild)
    }
}

#[derive(Debug, Default)]
pub struct WimSources<'a> {
    pub local: Vec<WimFile<'a>>,
    pub network: Vec<WimFile<'a>>,
}

impl WimSources<'_> {
    pub fn new() -> Self {
        Default::default()
    }

    pub fn reset(&mut self) {
        self.local.clear();
        self.network.clear();
    }
}

fn get_wim_xml(wim_file: &str) -> std::io::Result<File> {
    let tmp_file = NamedTempFile::new()?;
    let _ = Command::new("wiminfo")
        .args([
            wim_file,
            "--extract-xml",
            tmp_file.path().as_os_str().to_str().unwrap(),
        ])
        .output()
        .expect("Failed to extract XML data");
    let file = File::open(tmp_file.path())?;

    Ok(file)
}

fn parse_wim_file(wim_file: &str) -> std::io::Result<WimFile<'_>> {
    let mut wim_images: Vec<WimImage> = Vec::new();
    let wim_path = Path::new(wim_file);
    if !wim_path.exists() {
        return Err(std::io::Error::new(
            std::io::ErrorKind::NotFound,
            "Failed to read WIM file",
        ));
    };
    let xml_file = get_wim_xml(wim_file).expect("Failed to open XML file");
    let file = BufReader::new(xml_file);

    let mut current_element = String::new();
    let mut image = WimImage::new();

    let parser = EventReader::new(file);
    let mut depth = 0;
    for e in parser {
        match e {
            Ok(XmlEvent::StartElement {
                name, attributes, ..
            }) => {
                println!("{:spaces$}+{name}", "", spaces = depth * 2);
                depth += 1;
                current_element = name.local_name.to_uppercase();

                if current_element == "IMAGE" {
                    // Update index
                    if let Some(attr) = attributes.first()
                        && attr.name.to_string().to_lowercase() == "index"
                    {
                        image.index = attr.value.clone();
                    }
                }
            }
            Ok(XmlEvent::Characters(char_data)) => {
                if current_element == "BUILD" {
                    let build = char_data.trim();
                    image.build = build.to_string();
                    image.version = WIN_BUILDS.get(build).map_or("", |v| v).to_string();
                    println!("{:spaces$} {char_data}", "", spaces = (depth + 1) * 2);
                }
                if current_element == "NAME" {
                    image.name = char_data.trim().to_string();
                    println!("{:spaces$} {char_data}", "", spaces = (depth + 1) * 2);
                }
                if current_element == "SPBUILD" {
                    image.spbuild = char_data.trim().to_string();
                    println!("{:spaces$} {char_data}", "", spaces = (depth + 1) * 2);
                }
            }
            Ok(XmlEvent::EndElement { name }) => {
                depth -= 1;
                println!("{:spaces$}-{name}", "", spaces = depth * 2);

                if name.local_name.to_uppercase() == "IMAGE" {
                    // Append image to list
                    if !image.build.is_empty() && !image.name.is_empty() && !image.index.is_empty()
                    {
                        wim_images.push(image.clone());
                    }

                    // Reset image
                    image.reset()
                }
            }
            Err(e) => {
                eprintln!("Error: {e}");
                break;
            }
            _ => {}
        }
    }
    let wim_file = WimFile {
        path: wim_path,
        images: wim_images,
    };

    Ok(wim_file)
}

// fn main() -> std::io::Result<()> {
//     let mut sources = WimSources::new();
//     if let Ok(wim_file) = parse_wim_file("./23H2.wim")
//         && !wim_file.images.is_empty()
//     {
//         sources.local.push(wim_file);
//     }
//     if let Ok(wim_file) = parse_wim_file("./24H2.wim") {
//         sources.local.push(wim_file);
//     }
//     dbg!(&sources);
//     sources.local.iter().for_each(|f| {
//         println!("-- {} --", f.path.to_string_lossy());
//         f.images.iter().for_each(|i| {
//             println!("* {i}");
//         });
//         println!();
//     });
//
//     Ok(())
// }