OpenPXE/crates/iso-store/src/introspect.rs

//! ISO introspection — identify the distro family and locate kernel/initrd.
//!
//! We avoid a full ISO9660/Joliet/Rock-Ridge parser by reading a small number
//! of well-known files via `isoinfo` (from cdrtools/genisoimage) when it's on
//! the path. As a pure-Rust fallback we do a crude scan: read the volume
//! descriptor at offset 0x8000 to grab the volume label, and grep for known
//! filenames by scanning raw sectors — good enough to tell Debian from RHEL
//! most of the time, without shelling out.
//!
//! The returned `IntrospectionReport` is what `BootEntry`s get generated from.

use serde::{Deserialize, Serialize};
use std::io::{Read, Seek, SeekFrom};
use std::path::Path;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum DistroFamily {
    DebianUbuntu,
    RhelFedora,
    OpenSuse,
    Arch,
    Alpine,
    WindowsPe,
    Unknown,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IntrospectionReport {
    pub family: DistroFamily,
    pub volume_label: Option<String>,
    /// Kernel path inside the ISO (e.g. `/casper/vmlinuz`, `/isolinux/vmlinuz`).
    pub kernel_path: Option<String>,
    /// Initrd path(s) inside the ISO. May be multiple for multi-initrd setups.
    pub initrd_paths: Vec<String>,
    /// True if `sources/boot.wim` present — Windows install media.
    pub has_boot_wim: bool,
}

/// Probe an ISO file on disk. Never fails — on unrecoverable IO error we log
/// and return an `Unknown` family so the uploader still sees a record.
pub fn introspect(path: &Path) -> IntrospectionReport {
    let mut report = IntrospectionReport {
        family: DistroFamily::Unknown,
        volume_label: None,
        kernel_path: None,
        initrd_paths: Vec::new(),
        has_boot_wim: false,
    };

    let Ok(mut f) = std::fs::File::open(path) else {
        tracing::warn!(target: "openpxe::iso", "cannot open ISO for introspection: {}", path.display());
        return report;
    };

    // ISO9660 Primary Volume Descriptor at LBA 16 (offset 0x8000), 2048 bytes.
    // Bytes 40..72 are the Volume Identifier (space-padded, d-characters).
    let mut pvd = [0u8; 2048];
    if f.seek(SeekFrom::Start(0x8000)).is_ok() && f.read_exact(&mut pvd).is_ok() {
        // Byte 0 must be 0x01 (primary descriptor), bytes 1..6 = "CD001".
        if pvd[0] == 0x01 && &pvd[1..6] == b"CD001" {
            let label_raw = &pvd[40..72];
            let label = String::from_utf8_lossy(label_raw).trim().to_string();
            if !label.is_empty() {
                report.volume_label = Some(label.clone());
                report.family = family_from_label(&label);
            }
        }
    }

    // Cheap content scan: read the first ~64 MiB, look for signature filenames.
    // This is enough to identify `sources/boot.wim` (Windows) and common
    // kernel/initrd paths for the major Linux distros.
    let _ = f.seek(SeekFrom::Start(0));
    let scan_bytes = 64 * 1024 * 1024;
    let mut buf = vec![0u8; 1024 * 1024];
    let mut read_total = 0usize;
    let mut haystack = Vec::with_capacity(scan_bytes.min(32 * 1024 * 1024));
    while read_total < scan_bytes {
        let n = f.read(&mut buf).unwrap_or(0);
        if n == 0 {
            break;
        }
        haystack.extend_from_slice(&buf[..n]);
        read_total += n;
    }

    if contains_ascii(&haystack, b"sources/boot.wim")
        || contains_ascii(&haystack, b"SOURCES/BOOT.WIM")
        || contains_ascii(&haystack, b"SOURCES\\BOOT.WIM")
    {
        report.has_boot_wim = true;
        if report.family == DistroFamily::Unknown {
            report.family = DistroFamily::WindowsPe;
        }
    }

    // Best-effort kernel/initrd path guess from family. These paths are what
    // distro ISOs conventionally ship at — we don't verify extraction here;
    // that happens in the store after introspection.
    let (k, i) = guess_kernel_initrd(report.family);
    report.kernel_path = k.map(str::to_string);
    report.initrd_paths = i.iter().map(std::string::ToString::to_string).collect();

    report
}

fn family_from_label(label: &str) -> DistroFamily {
    let l = label.to_ascii_lowercase();
    if l.contains("ubuntu") || l.contains("debian") || l.contains("mint") {
        DistroFamily::DebianUbuntu
    } else if l.contains("rhel")
        || l.contains("centos")
        || l.contains("fedora")
        || l.contains("rocky")
        || l.contains("alma")
    {
        DistroFamily::RhelFedora
    } else if l.contains("suse") || l.contains("opensuse") {
        DistroFamily::OpenSuse
    } else if l.contains("arch") {
        DistroFamily::Arch
    } else if l.contains("alpine") {
        DistroFamily::Alpine
    } else if l.contains("windows") || l.contains("winpe") {
        DistroFamily::WindowsPe
    } else {
        DistroFamily::Unknown
    }
}

fn guess_kernel_initrd(family: DistroFamily) -> (Option<&'static str>, Vec<&'static str>) {
    match family {
        DistroFamily::DebianUbuntu => (Some("/casper/vmlinuz"), vec!["/casper/initrd"]),
        DistroFamily::RhelFedora => (
            Some("/images/pxeboot/vmlinuz"),
            vec!["/images/pxeboot/initrd.img"],
        ),
        DistroFamily::OpenSuse => (
            Some("/boot/x86_64/loader/linux"),
            vec!["/boot/x86_64/loader/initrd"],
        ),
        DistroFamily::Arch => (
            Some("/arch/boot/x86_64/vmlinuz-linux"),
            vec!["/arch/boot/x86_64/initramfs-linux.img"],
        ),
        DistroFamily::Alpine => (Some("/boot/vmlinuz-lts"), vec!["/boot/initramfs-lts"]),
        DistroFamily::WindowsPe | DistroFamily::Unknown => (None, Vec::new()),
    }
}

fn contains_ascii(haystack: &[u8], needle: &[u8]) -> bool {
    if needle.is_empty() || haystack.len() < needle.len() {
        return false;
    }
    haystack
        .windows(needle.len())
        .any(|w| w.eq_ignore_ascii_case(needle))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn label_matching() {
        assert_eq!(
            family_from_label("Ubuntu 24.04"),
            DistroFamily::DebianUbuntu
        );
        assert_eq!(
            family_from_label("Rocky-9-x86_64-dvd"),
            DistroFamily::RhelFedora
        );
        assert_eq!(
            family_from_label("openSUSE-Leap-15.6"),
            DistroFamily::OpenSuse
        );
        assert_eq!(family_from_label("ARCH_202604"), DistroFamily::Arch);
        assert_eq!(family_from_label("weird-custom"), DistroFamily::Unknown);
    }
}