beerpsi_fsdecrypt/src/vhd.rs

use std::{
    fs::{create_dir_all, File, FileTimes, OpenOptions},
    io::{self, BufRead, BufReader, Read, Seek, SeekFrom, Write},
    path::Path,
    time::{Duration, SystemTime},
};

use anyhow::{anyhow, Result};
use indicatif::{ProgressBar, ProgressStyle};
use ntfs::{structured_values::NtfsStandardInformation, Ntfs, NtfsAttributeType, NtfsTime};

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

const SECTOR_SIZE: u64 = 512;
const BUF_SIZE: usize = 256 * 1024;

// VHD format
const VHD_COOKIE: &[u8; 8] = b"conectix";
const VHD_TYPE_FIXED: u32 = 2;
const VHD_TYPE_DYNAMIC: u32 = 3;
const VHD_TYPE_DIFFERENCING: u32 = 4;
const VHD_FOOTER_DISK_TYPE_OFFSET: usize = 0x3C;
const VHD_FOOTER_DATA_OFFSET: usize = 0x10;

// Dynamic/differencing VHD header
const DYNAMIC_HEADER_COOKIE: &[u8; 8] = b"cxsparse";
const DYNAMIC_HEADER_SIZE: usize = 1024;
const DYNAMIC_BAT_OFFSET_FIELD: usize = 0x10;
const DYNAMIC_MAX_ENTRIES_FIELD: usize = 0x18;
const DYNAMIC_BLOCK_SIZE_FIELD: usize = 0x20;
const BAT_UNUSED: u32 = 0xFFFFFFFF;

// MBR
const MBR_SIGNATURE: [u8; 2] = [0x55, 0xAA];
const MBR_PARTITION_TABLE_OFFSET: usize = 0x1BE;
const MBR_PARTITION_ENTRY_SIZE: usize = 16;
const MBR_MAX_PARTITIONS: usize = 4;
const NTFS_PARTITION_TYPE: u8 = 0x07;

// NTFS boot sector magic
const NTFS_MAGIC: [u8; 4] = [0xEB, 0x52, 0x90, 0x4E];
/// Common virtual offsets where NTFS boot sector might start.
const NTFS_PROBE_OFFSETS: [u64; 4] = [0, 32_256, 1_048_576, 512];

// Progress bar
const PROGRESS_STYLE: &str =
    "{prefix} [{bar:20!.bright.yellow/dim.white}] {bytes:>8} [{elapsed}<{eta}, {bytes_per_sec}]";

// Windows epoch -> Unix epoch offset (100ns intervals)
const WINDOWS_EPOCH_OFFSET: u64 = 116_444_736_000_000_000;

// ---------------------------------------------------------------------------
// VHD error type
// ---------------------------------------------------------------------------

#[derive(Debug, thiserror::Error)]
pub enum VhdError {
    #[error(transparent)]
    Io(#[from] io::Error),
    #[error("Not a valid VHD file")]
    InvalidCookie,
    #[error("Unsupported VHD type {0}")]
    UnsupportedType(u32),
    #[error("Invalid dynamic VHD header")]
    InvalidDynamicHeader,
    #[error("No NTFS partition found in VHD")]
    NoNtfsPartition,
}

// ---------------------------------------------------------------------------
// VHD layout: how to map virtual offsets to file offsets
// ---------------------------------------------------------------------------

enum VhdLayout {
    /// Data is contiguous from offset 0 to (file_size - 512).
    Fixed,
    /// Data is in blocks addressed via a Block Allocation Table.
    /// Used for both dynamic (type 3) and differencing (type 4) VHDs.
    Sparse { bat: Vec<u32>, block_size: u64 },
}

impl VhdLayout {
    /// Parse the dynamic/differencing header and BAT.
    fn parse_sparse<R: Read + Seek>(inner: &mut R, footer: &[u8]) -> Result<(Self, u64), VhdError> {
        let header_offset = read_be_u64(footer, VHD_FOOTER_DATA_OFFSET);

        inner.seek(SeekFrom::Start(header_offset))?;
        let mut hdr = [0u8; DYNAMIC_HEADER_SIZE];
        inner.read_exact(&mut hdr)?;
        if &hdr[..8] != DYNAMIC_HEADER_COOKIE {
            return Err(VhdError::InvalidDynamicHeader);
        }

        let bat_offset = read_be_u64(&hdr, DYNAMIC_BAT_OFFSET_FIELD);
        let max_entries = read_be_u32(&hdr, DYNAMIC_MAX_ENTRIES_FIELD) as usize;
        let block_size = read_be_u32(&hdr, DYNAMIC_BLOCK_SIZE_FIELD) as u64;

        inner.seek(SeekFrom::Start(bat_offset))?;
        let mut raw = vec![0u8; max_entries * 4];
        inner.read_exact(&mut raw)?;
        let bat: Vec<u32> = (0..max_entries).map(|i| read_be_u32(&raw, i * 4)).collect();

        Ok((VhdLayout::Sparse { bat, block_size }, max_entries as u64 * block_size))
    }

    /// Read bytes from a virtual offset according to this layout.
    fn read_at<R: Read + Seek>(
        &self,
        inner: &mut R,
        virt_off: u64,
        virtual_size: u64,
        buf: &mut [u8],
    ) -> io::Result<usize> {
        if virt_off >= virtual_size {
            return Ok(0);
        }
        let cap = std::cmp::min(buf.len() as u64, virtual_size - virt_off) as usize;

        match self {
            VhdLayout::Fixed => {
                inner.seek(SeekFrom::Start(virt_off))?;
                inner.read(&mut buf[..cap])
            }
            VhdLayout::Sparse { bat, block_size } => {
                let bi = (virt_off / block_size) as usize;
                let bo = virt_off % block_size;
                let n = std::cmp::min(cap, (block_size - bo) as usize);

                if bi >= bat.len() || bat[bi] == BAT_UNUSED {
                    buf[..n].fill(0);
                    Ok(n)
                } else {
                    // Each block: bitmap sector + data. Skip bitmap.
                    let file_off = bat[bi] as u64 * SECTOR_SIZE + SECTOR_SIZE + bo;
                    inner.seek(SeekFrom::Start(file_off))?;
                    inner.read(&mut buf[..n])
                }
            }
        }
    }

    /// Read 4 bytes from a virtual offset (for magic-byte probing).
    fn read_magic<R: Read + Seek>(
        &self,
        inner: &mut R,
        offset: u64,
    ) -> io::Result<[u8; 4]> {
        let mut buf = [0u8; 4];
        match self {
            VhdLayout::Fixed => {
                inner.seek(SeekFrom::Start(offset))?;
                inner.read_exact(&mut buf)?;
            }
            VhdLayout::Sparse { bat, block_size } => {
                let bi = (offset / block_size) as usize;
                if bi < bat.len() && bat[bi] != BAT_UNUSED {
                    let file_off = bat[bi] as u64 * SECTOR_SIZE + SECTOR_SIZE + offset % block_size;
                    inner.seek(SeekFrom::Start(file_off))?;
                    inner.read_exact(&mut buf)?;
                }
            }
        }
        Ok(buf)
    }
}

// ---------------------------------------------------------------------------
// VHD reader (single VHD)
// ---------------------------------------------------------------------------

/// Transparently presents the NTFS partition within a fixed or dynamic VHD.
pub struct VhdReader<R> {
    inner: R,
    layout: VhdLayout,
    ntfs_offset: u64,
    virtual_size: u64,
    pos: u64,
}

impl<R: Read + Seek> VhdReader<R> {
    pub fn new(mut inner: R) -> Result<Self, VhdError> {
        let file_size = inner.seek(SeekFrom::End(0))?;
        if file_size < SECTOR_SIZE {
            return Err(VhdError::InvalidCookie);
        }

        inner.seek(SeekFrom::Start(file_size - SECTOR_SIZE))?;
        let mut footer = [0u8; SECTOR_SIZE as usize];
        inner.read_exact(&mut footer)?;
        if &footer[..8] != VHD_COOKIE {
            return Err(VhdError::InvalidCookie);
        }

        let disk_type = read_be_u32(&footer, VHD_FOOTER_DISK_TYPE_OFFSET);
        let (layout, virtual_size) = match disk_type {
            VHD_TYPE_FIXED => (VhdLayout::Fixed, file_size - SECTOR_SIZE),
            VHD_TYPE_DYNAMIC | VHD_TYPE_DIFFERENCING => {
                VhdLayout::parse_sparse(&mut inner, &footer)?
            }
            t => return Err(VhdError::UnsupportedType(t)),
        };

        let ntfs_offset = find_ntfs_offset(&mut inner, &layout, virtual_size)?;
        Ok(Self { inner, layout, ntfs_offset, virtual_size, pos: 0 })
    }

    fn ntfs_size(&self) -> u64 {
        self.virtual_size - self.ntfs_offset
    }
}

impl<R: Read + Seek> Read for VhdReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let remaining = self.ntfs_size().saturating_sub(self.pos);
        if remaining == 0 {
            return Ok(0);
        }
        let cap = std::cmp::min(buf.len() as u64, remaining) as usize;
        let n = self.layout.read_at(
            &mut self.inner, self.ntfs_offset + self.pos, self.virtual_size, &mut buf[..cap],
        )?;
        self.pos += n as u64;
        Ok(n)
    }
}

impl<R: Read + Seek> Seek for VhdReader<R> {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        let target = match pos {
            SeekFrom::Start(o) => o as i64,
            SeekFrom::Current(o) => self.pos as i64 + o,
            SeekFrom::End(o) => self.ntfs_size() as i64 + o,
        };
        if target < 0 {
            return Err(io::Error::new(io::ErrorKind::InvalidInput, "seek before start"));
        }
        self.pos = target as u64;
        Ok(self.pos)
    }
}

// ---------------------------------------------------------------------------
// Merged VHD reader (base + delta overlay, no disk merge needed)
// ---------------------------------------------------------------------------

/// Reads from a differencing VHD overlaid on a base VHD.
/// For each block, reads from delta if allocated, otherwise from base.
pub struct MergedVhdReader<R> {
    base: R,
    base_layout: VhdLayout,
    delta: R,
    delta_layout: VhdLayout,
    ntfs_offset: u64,
    virtual_size: u64,
    pos: u64,
}

impl<R: Read + Seek> MergedVhdReader<R> {
    pub fn new(mut base: R, mut delta: R) -> Result<Self, VhdError> {
        // Parse base
        let base_file_size = base.seek(SeekFrom::End(0))?;
        if base_file_size < SECTOR_SIZE {
            return Err(VhdError::InvalidCookie);
        }
        base.seek(SeekFrom::Start(base_file_size - SECTOR_SIZE))?;
        let mut base_footer = [0u8; SECTOR_SIZE as usize];
        base.read_exact(&mut base_footer)?;
        if &base_footer[..8] != VHD_COOKIE {
            return Err(VhdError::InvalidCookie);
        }

        let base_type = read_be_u32(&base_footer, VHD_FOOTER_DISK_TYPE_OFFSET);
        let (base_layout, base_vsize) = match base_type {
            VHD_TYPE_FIXED => (VhdLayout::Fixed, base_file_size - SECTOR_SIZE),
            VHD_TYPE_DYNAMIC | VHD_TYPE_DIFFERENCING => {
                VhdLayout::parse_sparse(&mut base, &base_footer)?
            }
            t => return Err(VhdError::UnsupportedType(t)),
        };

        // Parse delta
        let delta_file_size = delta.seek(SeekFrom::End(0))?;
        if delta_file_size < SECTOR_SIZE {
            return Err(VhdError::InvalidCookie);
        }
        delta.seek(SeekFrom::Start(delta_file_size - SECTOR_SIZE))?;
        let mut delta_footer = [0u8; SECTOR_SIZE as usize];
        delta.read_exact(&mut delta_footer)?;
        if &delta_footer[..8] != VHD_COOKIE {
            return Err(VhdError::InvalidCookie);
        }

        let delta_type = read_be_u32(&delta_footer, VHD_FOOTER_DISK_TYPE_OFFSET);
        let (delta_layout, _) = match delta_type {
            VHD_TYPE_DYNAMIC | VHD_TYPE_DIFFERENCING => {
                VhdLayout::parse_sparse(&mut delta, &delta_footer)?
            }
            t => return Err(VhdError::UnsupportedType(t)),
        };

        // Use base's virtual size as the canonical disk size
        let virtual_size = base_vsize;

        // Find NTFS using the merged view
        let ntfs_offset = find_ntfs_offset_merged(
            &mut base, &base_layout, &mut delta, &delta_layout, virtual_size,
        )?;

        Ok(Self {
            base, base_layout,
            delta, delta_layout,
            ntfs_offset, virtual_size, pos: 0,
        })
    }

    /// Read from the merged view: for each sector, check the delta's bitmap
    /// to decide whether to read from delta or base.
    fn read_merged(&mut self, virt_off: u64, buf: &mut [u8]) -> io::Result<usize> {
        if virt_off >= self.virtual_size {
            return Ok(0);
        }
        let cap = std::cmp::min(buf.len() as u64, self.virtual_size - virt_off) as usize;

        match &self.delta_layout {
            VhdLayout::Fixed => {
                // Shouldn't happen for a delta, but fall through to base
                self.base_layout.read_at(&mut self.base, virt_off, self.virtual_size, &mut buf[..cap])
            }
            VhdLayout::Sparse { bat, block_size } => {
                let bi = (virt_off / block_size) as usize;
                let bo = virt_off % block_size;
                let n = std::cmp::min(cap, (*block_size - bo) as usize);

                if bi >= bat.len() || bat[bi] == BAT_UNUSED {
                    // Block not in delta, read from base
                    return self.base_layout.read_at(
                        &mut self.base, virt_off, self.virtual_size, &mut buf[..n],
                    );
                }

                let block_file_offset = bat[bi] as u64 * SECTOR_SIZE;

                // Read the bitmap sector for this block
                self.delta.seek(SeekFrom::Start(block_file_offset))?;
                let mut bitmap = [0u8; SECTOR_SIZE as usize];
                self.delta.read_exact(&mut bitmap)?;

                // Check if the sector containing our offset has been modified
                let sector_in_block = (bo / SECTOR_SIZE) as usize;
                let bitmap_byte = bitmap[sector_in_block / 8];
                let bitmap_bit = 7 - (sector_in_block % 8); // MSB first
                let sector_modified = (bitmap_byte >> bitmap_bit) & 1 == 1;

                if sector_modified {
                    // Read from delta (skip bitmap sector)
                    let file_off = block_file_offset + SECTOR_SIZE + bo;
                    self.delta.seek(SeekFrom::Start(file_off))?;
                    self.delta.read(&mut buf[..n])
                } else {
                    // Sector not modified in delta, read from base
                    self.base_layout.read_at(
                        &mut self.base, virt_off, self.virtual_size, &mut buf[..n],
                    )
                }
            }
        }
    }

    fn ntfs_size(&self) -> u64 {
        self.virtual_size - self.ntfs_offset
    }
}

impl<R: Read + Seek> Read for MergedVhdReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let remaining = self.ntfs_size().saturating_sub(self.pos);
        if remaining == 0 {
            return Ok(0);
        }
        let cap = std::cmp::min(buf.len() as u64, remaining) as usize;
        let n = self.read_merged(self.ntfs_offset + self.pos, &mut buf[..cap])?;
        self.pos += n as u64;
        Ok(n)
    }
}

impl<R: Read + Seek> Seek for MergedVhdReader<R> {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        let target = match pos {
            SeekFrom::Start(o) => o as i64,
            SeekFrom::Current(o) => self.pos as i64 + o,
            SeekFrom::End(o) => self.ntfs_size() as i64 + o,
        };
        if target < 0 {
            return Err(io::Error::new(io::ErrorKind::InvalidInput, "seek before start"));
        }
        self.pos = target as u64;
        Ok(self.pos)
    }
}

// ---------------------------------------------------------------------------
// NTFS partition detection (shared between single and merged readers)
// ---------------------------------------------------------------------------

/// Find NTFS offset in a single VHD.
fn find_ntfs_offset<R: Read + Seek>(
    inner: &mut R,
    layout: &VhdLayout,
    vsize: u64,
) -> Result<u64, VhdError> {
    // Try MBR
    if vsize >= SECTOR_SIZE {
        let mut mbr = [0u8; SECTOR_SIZE as usize];
        let _ = layout.read_at(inner, 0, vsize, &mut mbr);

        if mbr[510..512] == MBR_SIGNATURE {
            for i in 0..MBR_MAX_PARTITIONS {
                let eo = MBR_PARTITION_TABLE_OFFSET + i * MBR_PARTITION_ENTRY_SIZE;
                if mbr[eo + 4] == NTFS_PARTITION_TYPE {
                    let lba = u32::from_le_bytes(mbr[eo + 8..eo + 12].try_into().unwrap());
                    let offset = lba as u64 * SECTOR_SIZE;
                    if offset + 4 <= vsize && layout.read_magic(inner, offset)? == NTFS_MAGIC {
                        return Ok(offset);
                    }
                }
            }
        }
    }

    // Probe common offsets
    for offset in NTFS_PROBE_OFFSETS {
        if offset + 4 <= vsize && layout.read_magic(inner, offset)? == NTFS_MAGIC {
            return Ok(offset);
        }
    }

    Err(VhdError::NoNtfsPartition)
}

/// Read from merged view with bitmap awareness: if the delta has the block
/// allocated but the specific sector's bitmap bit is 0, read from base instead.
fn read_merged_sector<R: Read + Seek>(
    base: &mut R,
    base_layout: &VhdLayout,
    delta: &mut R,
    delta_layout: &VhdLayout,
    vsize: u64,
    virt_off: u64,
    buf: &mut [u8],
) -> io::Result<usize> {
    match delta_layout {
        VhdLayout::Fixed => base_layout.read_at(base, virt_off, vsize, buf),
        VhdLayout::Sparse { bat, block_size } => {
            let bi = (virt_off / block_size) as usize;
            if bi >= bat.len() || bat[bi] == BAT_UNUSED {
                return base_layout.read_at(base, virt_off, vsize, buf);
            }
            let bo = virt_off % block_size;
            let block_file_offset = bat[bi] as u64 * SECTOR_SIZE;

            // Read bitmap
            delta.seek(SeekFrom::Start(block_file_offset))?;
            let mut bitmap = [0u8; SECTOR_SIZE as usize];
            delta.read_exact(&mut bitmap)?;

            let sector_in_block = (bo / SECTOR_SIZE) as usize;
            let bitmap_byte = bitmap[sector_in_block / 8];
            let bitmap_bit = 7 - (sector_in_block % 8);
            let modified = (bitmap_byte >> bitmap_bit) & 1 == 1;

            if modified {
                let file_off = block_file_offset + SECTOR_SIZE + bo;
                delta.seek(SeekFrom::Start(file_off))?;
                delta.read(buf)
            } else {
                base_layout.read_at(base, virt_off, vsize, buf)
            }
        }
    }
}

/// Find NTFS offset in a merged (base + delta) view.
fn find_ntfs_offset_merged<R: Read + Seek>(
    base: &mut R,
    base_layout: &VhdLayout,
    delta: &mut R,
    delta_layout: &VhdLayout,
    vsize: u64,
) -> Result<u64, VhdError> {
    let read_magic = |base: &mut R, delta: &mut R, offset: u64| -> io::Result<[u8; 4]> {
        let mut buf = [0u8; 4];
        read_merged_sector(base, base_layout, delta, delta_layout, vsize, offset, &mut buf)?;
        Ok(buf)
    };

    // Try MBR from merged view
    if vsize >= SECTOR_SIZE {
        let mut mbr = [0u8; SECTOR_SIZE as usize];
        read_merged_sector(base, base_layout, delta, delta_layout, vsize, 0, &mut mbr)?;

        if mbr[510..512] == MBR_SIGNATURE {
            for i in 0..MBR_MAX_PARTITIONS {
                let eo = MBR_PARTITION_TABLE_OFFSET + i * MBR_PARTITION_ENTRY_SIZE;
                if mbr[eo + 4] == NTFS_PARTITION_TYPE {
                    let lba = u32::from_le_bytes(mbr[eo + 8..eo + 12].try_into().unwrap());
                    let offset = lba as u64 * SECTOR_SIZE;
                    if offset + 4 <= vsize && read_magic(base, delta, offset)? == NTFS_MAGIC {
                        return Ok(offset);
                    }
                }
            }
        }
    }

    for offset in NTFS_PROBE_OFFSETS {
        if offset + 4 <= vsize && read_magic(base, delta, offset)? == NTFS_MAGIC {
            return Ok(offset);
        }
    }

    Err(VhdError::NoNtfsPartition)
}

// ---------------------------------------------------------------------------
// NTFS extraction (shared logic)
// ---------------------------------------------------------------------------

fn is_ntfs_system_entry(name: &str) -> bool {
    name.starts_with('$') || name == "." || name == ".." || name == "System Volume Information"
}

fn ntfs_time_to_system_time(t: NtfsTime) -> SystemTime {
    let nanos = (t.nt_timestamp() - WINDOWS_EPOCH_OFFSET) * 100;
    SystemTime::UNIX_EPOCH + Duration::from_nanos(nanos)
}

fn set_ntfs_timestamps<T: Read + Seek>(fs: &mut T, file: &ntfs::NtfsFile, path: &Path) {
    let mut attrs = file.attributes();
    while let Some(Ok(attr)) = attrs.next(fs) {
        if let Ok(attr) = attr.to_attribute() {
            if let Ok(NtfsAttributeType::StandardInformation) = attr.ty() {
                if let Ok(info) = attr.resident_structured_value::<NtfsStandardInformation>() {
                    let _ = OpenOptions::new().write(true).open(path).and_then(|h| {
                        h.set_times(
                            FileTimes::new()
                                .set_accessed(ntfs_time_to_system_time(info.access_time()))
                                .set_modified(ntfs_time_to_system_time(info.modification_time())),
                        )
                    });
                }
                break;
            }
        }
    }
}

fn extract_ntfs_dir<T: Read + Seek>(
    ntfs: &Ntfs,
    fs: &mut T,
    dir: &ntfs::NtfsFile,
    out: &Path,
    pb: &ProgressBar,
) -> Result<()> {
    let index = dir.directory_index(fs)?;
    let mut iter = index.entries();

    while let Some(entry) = iter.next(fs) {
        let entry = entry?;
        let key = entry.key().ok_or_else(|| anyhow!("missing key"))??;
        let name = key.name().to_string_lossy();
        if is_ntfs_system_entry(&name) {
            continue;
        }

        let file = entry.to_file(ntfs, fs)?;
        let dest = out.join(&*name);

        if key.is_directory() {
            create_dir_all(&dest)?;
            extract_ntfs_dir(ntfs, fs, &file, &dest, pb)?;
            set_ntfs_timestamps(fs, &file, &dest);
        } else if let Some(data) = file.data(fs, "") {
            let data_item = data?;
            let attr = data_item.to_attribute()?;
            let mut reader = BufReader::with_capacity(BUF_SIZE, attr.value(fs)?.attach(fs));
            let mut out_file = File::create(&dest)?;

            loop {
                let buf = reader.fill_buf()?;
                if buf.is_empty() {
                    break;
                }
                out_file.write_all(buf)?;
                let n = buf.len();
                reader.consume(n);
                pb.inc(n as u64);
            }
            out_file.flush()?;
            drop(reader);
            set_ntfs_timestamps(fs, &file, &dest);
        }
    }
    Ok(())
}

fn calculate_ntfs_size<T: Read + Seek>(
    ntfs: &Ntfs,
    fs: &mut T,
    dir: &ntfs::NtfsFile,
) -> Result<u64> {
    let mut total = 0u64;
    let index = dir.directory_index(fs)?;
    let mut iter = index.entries();

    while let Some(entry) = iter.next(fs) {
        let entry = entry?;
        let key = entry.key().ok_or_else(|| anyhow!("missing key"))??;
        if is_ntfs_system_entry(&key.name().to_string_lossy().as_ref()) {
            continue;
        }
        let file = entry.to_file(ntfs, fs)?;
        if key.is_directory() {
            total += calculate_ntfs_size(ntfs, fs, &file)?;
        } else if let Some(data) = file.data(fs, "") {
            total += data?.to_attribute()?.value_length();
        }
    }
    Ok(total)
}

/// Shared extraction logic: given an NTFS-bearing Read+Seek, extract to output_dir.
fn extract_ntfs_to_dir<T: Read + Seek>(fs: &mut T, output_dir: &Path, prefix: &str) -> Result<()> {
    let mut ntfs = Ntfs::new(fs)?;
    ntfs.read_upcase_table(fs)?;

    let root = ntfs.root_directory(fs)?;
    let total = calculate_ntfs_size(&ntfs, fs, &root)?;

    let pb = ProgressBar::new(total)
        .with_style(ProgressStyle::default_bar().template(PROGRESS_STYLE)?);
    pb.set_prefix(prefix.to_string());

    create_dir_all(output_dir)?;
    let root = ntfs.root_directory(fs)?;
    extract_ntfs_dir(&ntfs, fs, &root, output_dir, &pb)?;
    pb.finish();

    Ok(())
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/// Extract all files from a single VHD's NTFS filesystem, then delete the VHD.
/// The `output_dir` is where files are extracted to.
pub fn extract_vhd(vhd_path: &Path, output_dir: &Path) -> Result<()> {
    println!("Extracting VHD: {}", vhd_path.display());

    let mut vhd = VhdReader::new(File::open(vhd_path)?).map_err(|e| anyhow!(e))?;
    let prefix = output_dir.file_name().unwrap_or_default().to_string_lossy().to_string();
    extract_ntfs_to_dir(&mut vhd, output_dir, &prefix)?;

    println!("Extracted to: {}", output_dir.display());

    drop(vhd);
    if let Err(e) = std::fs::remove_file(vhd_path) {
        println!("WARNING: Could not delete VHD: {e}");
    }
    Ok(())
}

/// Extract files from a merged view of base + delta VHDs (pure Rust, no admin needed).
/// Reads delta blocks where available, falls back to base. Deletes both VHDs after.
/// The `output_dir` is where files are extracted to.
pub fn extract_merged_vhd(base_path: &Path, delta_path: &Path, output_dir: &Path) -> Result<()> {
    println!("Extracting merged VHD: {} + {}", base_path.display(), delta_path.display());

    let base = File::open(base_path)?;
    let delta = File::open(delta_path)?;
    let mut merged = MergedVhdReader::new(base, delta).map_err(|e| anyhow!(e))?;

    let prefix = output_dir.file_name().unwrap_or_default().to_string_lossy().to_string();
    extract_ntfs_to_dir(&mut merged, output_dir, &prefix)?;

    println!("Extracted to: {}", output_dir.display());

    drop(merged);
    for path in [base_path, delta_path] {
        if let Err(e) = std::fs::remove_file(path) {
            println!("WARNING: Could not delete {}: {e}", path.display());
        }
    }
    Ok(())
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

fn read_be_u32(buf: &[u8], offset: usize) -> u32 {
    u32::from_be_bytes(buf[offset..offset + 4].try_into().unwrap())
}

fn read_be_u64(buf: &[u8], offset: usize) -> u64 {
    u64::from_be_bytes(buf[offset..offset + 8].try_into().unwrap())
}