diff --git a/src/main.rs b/src/main.rs index 52f9df0..512daba 100644 --- a/src/main.rs +++ b/src/main.rs @@ -25,21 +25,6 @@ mod crypto; mod stream; mod vhd; -/// Info collected from each input file for sorting and post-extraction merge. -struct InputFile { - path: PathBuf, - sequence_number: u8, -} - -/// Info about an extracted VHD, used for delta merge. -struct ExtractedVhd { - vhd_path: PathBuf, - /// Original input .app file path (for resolving sibling files) - input_path: PathBuf, - sequence_number: u8, - game_id: String, -} - fn exfat_timestamp_to_system_time( timestamp: &exfat_fs::timestamp::Timestamp, ) -> Result { @@ -225,34 +210,44 @@ fn extract_internal_vhd(image_path: &Path, sequence_number: u8) -> Result PathBuf { +// --------------------------------------------------------------------------- +// Delta merge support +// --------------------------------------------------------------------------- + +/// Info about an extracted VHD, used for delta merge. +struct ExtractedVhd { + vhd_path: PathBuf, + input_path: PathBuf, + sequence_number: u8, + game_id: String, +} + +/// Strip the `\\?\` prefix that `canonicalize` adds on Windows. +fn strip_unc_prefix(path: PathBuf) -> PathBuf { let s = path.to_string_lossy(); - if let Some(stripped) = s.strip_prefix(r"\\?\") { - PathBuf::from(stripped) - } else { - path - } + s.strip_prefix(r"\\?\").map(PathBuf::from).unwrap_or(path) +} + +/// Search a directory for a file matching `{prefix}*{suffix}`. +fn find_sibling(dir: &Path, prefix: &str, suffix: &str) -> Option { + std::fs::read_dir(dir).ok()? + .filter_map(|e| e.ok()) + .map(|e| e.path()) + .find(|p| { + let name = p.file_name().unwrap_or_default().to_string_lossy(); + name.starts_with(prefix) && name.ends_with(suffix) + }) } /// Merge a differencing VHD into its parent using Hyper-V PowerShell cmdlets. -/// -/// Steps: -/// 1. `Set-VHD` links the delta VHD to its parent -/// 2. `Merge-VHD` merges the delta's changes into the parent (modifies parent in-place) -/// /// Requires elevation (triggers UAC prompt). #[cfg(windows)] fn merge_vhd(base_vhd: &Path, delta_vhd: &Path) -> Result<()> { - let base_abs = strip_extended_path_prefix(std::fs::canonicalize(base_vhd)?); - let delta_abs = strip_extended_path_prefix(std::fs::canonicalize(delta_vhd)?); + let base_abs = strip_unc_prefix(std::fs::canonicalize(base_vhd)?); + let delta_abs = strip_unc_prefix(std::fs::canonicalize(delta_vhd)?); - println!("Merging VHDs..."); - println!(" Base (parent): {}", base_abs.display()); - println!(" Delta (child): {}", delta_abs.display()); + println!("Merging VHDs: {} <- {}", base_abs.display(), delta_abs.display()); - // Set-VHD links the differencing disk to its parent, then - // Merge-VHD (without -DestinationPath) merges the child into its immediate parent. let ps_commands = format!( "Set-VHD -Path '{}' -ParentPath '{}'; Merge-VHD -Path '{}' -Force", delta_abs.display(), @@ -261,9 +256,8 @@ fn merge_vhd(base_vhd: &Path, delta_vhd: &Path) -> Result<()> { ); let error_log = delta_vhd.with_extension("merge_error.txt"); - let error_log_abs = strip_extended_path_prefix(std::path::absolute(&error_log)?); + let error_log_abs = strip_unc_prefix(std::path::absolute(&error_log)?); - // Wrap in try/catch to capture errors from the elevated process let wrapped = format!( "try {{ {} }} catch {{ $_ | Out-File '{}' -Encoding UTF8; throw }}", ps_commands, @@ -285,14 +279,67 @@ fn merge_vhd(base_vhd: &Path, delta_vhd: &Path) -> Result<()> { std::fs::remove_file(&error_log).ok(); println!("WARNING: VHD merge failed: {}", error_text.trim()); } else if status.success() { - println!("Merged delta into base VHD: {}", base_abs.display()); + println!("Merged into: {}", base_abs.display()); } else { - println!("WARNING: PowerShell exited with: {status}. Check UAC was accepted."); + println!("WARNING: Merge failed (exit: {status}). Check UAC was accepted."); } Ok(()) } +/// Resolve the base VHD for a delta: check extracted VHDs, then look for +/// an existing .vhd or .app in the same directory. +#[cfg(windows)] +fn resolve_base_vhd<'a>( + base: Option<&'a ExtractedVhd>, + game_id: &str, + delta_dir: &Path, + out: &'a mut Option, +) -> Option<&'a ExtractedVhd> { + if let Some(b) = base { + return Some(b); + } + + let prefix = format!("{game_id}_"); + + // Look for existing base VHD + if let Some(vhd_path) = find_sibling(delta_dir, &prefix, "_0.vhd") { + println!("Found existing base VHD: {}", vhd_path.display()); + *out = Some(ExtractedVhd { + vhd_path, input_path: PathBuf::new(), + sequence_number: 0, game_id: game_id.into(), + }); + return out.as_ref(); + } + + // Look for base .app and extract + if let Some(app_path) = find_sibling(delta_dir, &prefix, "_0.app") { + println!("Found base APP, extracting: {}", app_path.display()); + if let Ok(f) = File::open(&app_path) + .and_then(|f| FscryptDecryptor::new(f).map_err(|e| std::io::Error::other(e))) + { + let seq = f.bootid.sequence_number; + drop(f); + if let Ok(vhd_path) = extract_internal_vhd(&app_path, seq) { + *out = Some(ExtractedVhd { + vhd_path, input_path: app_path, + sequence_number: 0, game_id: game_id.into(), + }); + return out.as_ref(); + } + } + println!("WARNING: Failed to extract base VHD from {}", app_path.display()); + } + + println!("WARNING: No base (seq=0) found for {game_id}."); + println!(" Place the base .app or .vhd in the same directory."); + None +} + +// --------------------------------------------------------------------------- +// CLI & main +// --------------------------------------------------------------------------- + #[derive(Parser)] #[command(version, about = "decryptor for some SEGA containers", long_about = None)] struct Cli { @@ -306,24 +353,18 @@ struct Cli { fn main() -> Result<()> { let cli = Cli::parse(); - // Pre-read bootids for all files to sort by sequence number (base first) - let mut inputs: Vec = Vec::new(); + // Pre-read bootids to sort by sequence number (base first) + let mut inputs: Vec<(PathBuf, u8)> = Vec::new(); for path in &cli.files { let file = FscryptDecryptor::new(File::open(path)?).map_err(|e| anyhow!(e))?; - inputs.push(InputFile { - path: path.clone(), - sequence_number: file.bootid.sequence_number, - }); + inputs.push((path.clone(), file.bootid.sequence_number)); } - - // Sort so seq=0 (base) is processed before seq>0 (deltas) - inputs.sort_by_key(|f| f.sequence_number); + inputs.sort_by_key(|(_, seq)| *seq); // Track extracted VHDs for post-extraction merge let mut extracted_vhds: Vec = Vec::new(); - for input in &inputs { - let path = &input.path; + for (path, _) in &inputs { let file = FscryptDecryptor::new(File::open(path)?).map_err(|e| anyhow!(e))?; let bootid = file.bootid.clone(); let output_filename = file.filename()?; @@ -396,132 +437,47 @@ fn main() -> Result<()> { // Post-extraction: merge deltas and extract VHD contents (Windows only) #[cfg(windows)] if !cli.no_extract && !extracted_vhds.is_empty() { - // Group by game_id let mut by_game: HashMap> = HashMap::new(); for vhd in &extracted_vhds { by_game.entry(vhd.game_id.clone()).or_default().push(vhd); } - // Track which VHDs to extract contents from at the end let mut vhds_to_extract: Vec = Vec::new(); for (game_id, vhds) in &by_game { - let base = vhds.iter().find(|v| v.sequence_number == 0); + let base = vhds.iter().find(|v| v.sequence_number == 0).copied(); let deltas: Vec<_> = vhds.iter().filter(|v| v.sequence_number > 0).collect(); if deltas.is_empty() { - // Standalone base VHD, just extract its contents if let Some(base) = base { vhds_to_extract.push(base.vhd_path.clone()); } continue; } - // If no base was extracted in this run, search the same directory - let found_base; - let base = match base { - Some(b) => b, - None => { - let delta_dir = deltas[0] - .input_path - .parent() - .unwrap_or_else(|| Path::new(".")); - let pattern_prefix = format!("{game_id}_"); - - // First, look for an existing base VHD ({game_id}_*_0.vhd) - let existing_vhd = std::fs::read_dir(delta_dir)? - .filter_map(|e| e.ok()) - .map(|e| e.path()) - .find(|p| { - let name = p.file_name().unwrap_or_default().to_string_lossy(); - name.starts_with(&pattern_prefix) && name.ends_with("_0.vhd") - }); - - if let Some(vhd_path) = existing_vhd { - println!("Found existing base VHD: {}", vhd_path.display()); - found_base = ExtractedVhd { - vhd_path, - input_path: PathBuf::new(), - sequence_number: 0, - game_id: game_id.clone(), - }; - &found_base - } else { - // Look for a base .app file ({game_id}_*_0.app) and extract it - let base_app = std::fs::read_dir(delta_dir)? - .filter_map(|e| e.ok()) - .map(|e| e.path()) - .find(|p| { - let name = p.file_name().unwrap_or_default().to_string_lossy(); - name.starts_with(&pattern_prefix) && name.ends_with("_0.app") - }); - - match base_app { - Some(app_path) => { - println!( - "Found base APP in same directory, extracting: {}", - app_path.display() - ); - let base_file = FscryptDecryptor::new(File::open(&app_path)?) - .map_err(|e| anyhow!(e))?; - let base_seq = base_file.bootid.sequence_number; - drop(base_file); - - match extract_internal_vhd(&app_path, base_seq) { - Ok(vhd_path) => { - found_base = ExtractedVhd { - vhd_path, - input_path: app_path, - sequence_number: 0, - game_id: game_id.clone(), - }; - &found_base - } - Err(e) => { - println!( - "WARNING: Failed to extract base VHD: {e:#?}" - ); - continue; - } - } - } - None => { - println!( - "WARNING: Delta VHD(s) found for {game_id} but no base (seq=0) found." - ); - println!( - " Place the base .app or .vhd in the same directory." - ); - continue; - } - } - } - } + let delta_dir = deltas[0].input_path.parent().unwrap_or(Path::new(".")); + let mut resolved = None; + let Some(base) = resolve_base_vhd(base, game_id, delta_dir, &mut resolved) else { + continue; }; let mut merge_ok = true; - for delta in deltas { + for delta in &deltas { if let Err(e) = merge_vhd(&base.vhd_path, &delta.vhd_path) { - println!( - "WARNING: Failed to merge {} with {}: {e:#?}", - delta.vhd_path.display(), - base.vhd_path.display() - ); + println!("WARNING: Failed to merge: {e:#}"); merge_ok = false; } } - // Extract the final merged VHD if merge_ok { vhds_to_extract.push(base.vhd_path.clone()); } } - // Extract all VHD contents to folders, then delete the VHDs for vhd_path in &vhds_to_extract { if vhd_path.exists() { if let Err(e) = vhd::extract_vhd(vhd_path) { - println!("WARNING: Failed to extract VHD contents: {e:#?}"); + println!("WARNING: VHD extraction failed: {e:#}"); } } } diff --git a/src/vhd.rs b/src/vhd.rs index 704aed4..f236393 100644 --- a/src/vhd.rs +++ b/src/vhd.rs @@ -10,27 +10,45 @@ use indicatif::{ProgressBar, ProgressStyle}; use ntfs::{structured_values::NtfsStandardInformation, Ntfs, NtfsAttributeType, NtfsTime}; // --------------------------------------------------------------------------- -// VHD constants +// Constants // --------------------------------------------------------------------------- -const VHD_FOOTER_SIZE: u64 = 512; +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_FOOTER_DISK_TYPE_OFFSET: usize = 0x3C; +const VHD_FOOTER_DATA_OFFSET: usize = 0x10; +// Dynamic VHD header const DYNAMIC_HEADER_COOKIE: &[u8; 8] = b"cxsparse"; -const BAT_ENTRY_UNUSED: u32 = 0xFFFFFFFF; - -/// Sectors per bitmap: each data block is preceded by a sector-aligned bitmap. -const BITMAP_SECTORS: u64 = 1; +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; -const NTFS_PROBE_OFFSETS: &[u64] = &[0, 32_256, 1_048_576, 512]; +// 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 @@ -40,16 +58,12 @@ const NTFS_MAGIC: [u8; 4] = [0xEB, 0x52, 0x90, 0x4E]; pub enum VhdError { #[error(transparent)] Io(#[from] io::Error), - - #[error("Not a valid VHD file (bad cookie)")] + #[error("Not a valid VHD file")] InvalidCookie, - - #[error("Unsupported VHD type {0} (only fixed=2 and dynamic=3 are supported)")] + #[error("Unsupported VHD type {0} (only fixed and dynamic are supported)")] UnsupportedType(u32), - #[error("Invalid dynamic VHD header")] InvalidDynamicHeader, - #[error("No NTFS partition found in VHD")] NoNtfsPartition, } @@ -59,234 +73,161 @@ pub enum VhdError { // --------------------------------------------------------------------------- enum VhdLayout { - /// Fixed VHD: data is contiguous from offset 0 to (file_size - 512). Fixed, - /// Dynamic VHD: data is in blocks addressed via a Block Allocation Table. - Dynamic { - bat: Vec, - block_size: u64, - }, + Dynamic { bat: Vec, block_size: u64 }, } -/// A reader that transparently presents the NTFS partition within a VHD file. +/// Transparently presents the NTFS partition within a fixed or dynamic VHD. pub struct VhdReader { inner: R, layout: VhdLayout, - /// Byte offset where the NTFS partition starts within the virtual disk. ntfs_offset: u64, - /// Total virtual disk size. - virtual_disk_size: u64, - /// Current virtual position (relative to NTFS start). + virtual_size: u64, pos: u64, } impl VhdReader { pub fn new(mut inner: R) -> Result { let file_size = inner.seek(SeekFrom::End(0))?; - if file_size < VHD_FOOTER_SIZE { + if file_size < SECTOR_SIZE { return Err(VhdError::InvalidCookie); } - // Read footer (last 512 bytes) - inner.seek(SeekFrom::Start(file_size - VHD_FOOTER_SIZE))?; - let mut footer = [0u8; VHD_FOOTER_SIZE as usize]; + inner.seek(SeekFrom::Start(file_size - SECTOR_SIZE))?; + let mut footer = [0u8; SECTOR_SIZE as usize]; inner.read_exact(&mut footer)?; - - if &footer[0..8] != VHD_COOKIE { + if &footer[..8] != VHD_COOKIE { return Err(VhdError::InvalidCookie); } - let disk_type = - u32::from_be_bytes([footer[0x3C], footer[0x3D], footer[0x3E], footer[0x3F]]); - - let (layout, virtual_disk_size) = match disk_type { - VHD_TYPE_FIXED => { - let vds = file_size - VHD_FOOTER_SIZE; - (VhdLayout::Fixed, vds) - } - VHD_TYPE_DYNAMIC => { - let (layout, vds) = Self::parse_dynamic(&mut inner, &footer)?; - (layout, vds) - } - other => return Err(VhdError::UnsupportedType(other)), + 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 => Self::parse_dynamic(&mut inner, &footer)?, + t => return Err(VhdError::UnsupportedType(t)), }; - // Detect NTFS partition offset within the virtual disk - let ntfs_offset = - Self::detect_ntfs_offset_virtual(&mut inner, &layout, virtual_disk_size)?; - - Ok(Self { - inner, - layout, - ntfs_offset, - virtual_disk_size, - pos: 0, - }) + let ntfs_offset = Self::find_ntfs(&mut inner, &layout, virtual_size)?; + Ok(Self { inner, layout, ntfs_offset, virtual_size, pos: 0 }) } fn parse_dynamic(inner: &mut R, footer: &[u8]) -> Result<(VhdLayout, u64), VhdError> { - // data_offset in footer (big-endian u64 at 0x10) points to dynamic header - let data_offset = u64::from_be_bytes([ - footer[0x10], - footer[0x11], - footer[0x12], - footer[0x13], - footer[0x14], - footer[0x15], - footer[0x16], - footer[0x17], - ]); + let header_offset = read_be_u64(footer, VHD_FOOTER_DATA_OFFSET); - // Read dynamic disk header (1024 bytes) - inner.seek(SeekFrom::Start(data_offset))?; - let mut hdr = [0u8; 1024]; + inner.seek(SeekFrom::Start(header_offset))?; + let mut hdr = [0u8; DYNAMIC_HEADER_SIZE]; inner.read_exact(&mut hdr)?; - - if &hdr[0..8] != DYNAMIC_HEADER_COOKIE { + if &hdr[..8] != DYNAMIC_HEADER_COOKIE { return Err(VhdError::InvalidDynamicHeader); } - // BAT offset (big-endian u64 at 0x10 in header) - let bat_offset = u64::from_be_bytes([ - hdr[0x10], hdr[0x11], hdr[0x12], hdr[0x13], hdr[0x14], hdr[0x15], hdr[0x16], - hdr[0x17], - ]); + 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; - // Max table entries (big-endian u32 at 0x18) - let max_entries = - u32::from_be_bytes([hdr[0x18], hdr[0x19], hdr[0x1A], hdr[0x1B]]) as usize; - - // Block size (big-endian u32 at 0x20) - let block_size = u32::from_be_bytes([hdr[0x20], hdr[0x21], hdr[0x22], hdr[0x23]]) as u64; - - // Read BAT inner.seek(SeekFrom::Start(bat_offset))?; - let mut bat_bytes = vec![0u8; max_entries * 4]; - inner.read_exact(&mut bat_bytes)?; + let mut raw = vec![0u8; max_entries * 4]; + inner.read_exact(&mut raw)?; + let bat: Vec = (0..max_entries).map(|i| read_be_u32(&raw, i * 4)).collect(); - let bat: Vec = (0..max_entries) - .map(|i| { - u32::from_be_bytes([ - bat_bytes[i * 4], - bat_bytes[i * 4 + 1], - bat_bytes[i * 4 + 2], - bat_bytes[i * 4 + 3], - ]) - }) - .collect(); - - let virtual_disk_size = max_entries as u64 * block_size; - - Ok((VhdLayout::Dynamic { bat, block_size }, virtual_disk_size)) + Ok((VhdLayout::Dynamic { bat, block_size }, max_entries as u64 * block_size)) } - /// Read from a virtual disk offset, handling both fixed and dynamic layouts. - fn read_virtual(&mut self, virtual_offset: u64, buf: &mut [u8]) -> io::Result { - if virtual_offset >= self.virtual_disk_size { + /// Translate a virtual disk offset to a file read. + fn read_virtual(&mut self, virt_off: u64, buf: &mut [u8]) -> io::Result { + if virt_off >= self.virtual_size { return Ok(0); } - - let remaining = self.virtual_disk_size - virtual_offset; - let to_read = std::cmp::min(buf.len() as u64, remaining) as usize; + let cap = std::cmp::min(buf.len() as u64, self.virtual_size - virt_off) as usize; match &self.layout { VhdLayout::Fixed => { - self.inner.seek(SeekFrom::Start(virtual_offset))?; - self.inner.read(&mut buf[..to_read]) + self.inner.seek(SeekFrom::Start(virt_off))?; + self.inner.read(&mut buf[..cap]) } VhdLayout::Dynamic { bat, block_size } => { - let block_index = (virtual_offset / block_size) as usize; - let offset_in_block = virtual_offset % block_size; - let max_in_block = (*block_size - offset_in_block) as usize; - let to_read = std::cmp::min(to_read, max_in_block); + 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 block_index >= bat.len() || bat[block_index] == BAT_ENTRY_UNUSED { - // Unallocated block: return zeros - buf[..to_read].fill(0); - Ok(to_read) + if bi >= bat.len() || bat[bi] == BAT_UNUSED { + buf[..n].fill(0); + Ok(n) } else { - // Block starts at sector bat[block_index], skip bitmap sector(s) - let block_file_offset = bat[block_index] as u64 * 512 - + BITMAP_SECTORS * 512 - + offset_in_block; - self.inner.seek(SeekFrom::Start(block_file_offset))?; - self.inner.read(&mut buf[..to_read]) + // Each block: bitmap sector + data. Skip bitmap. + let file_off = bat[bi] as u64 * SECTOR_SIZE + SECTOR_SIZE + bo; + self.inner.seek(SeekFrom::Start(file_off))?; + self.inner.read(&mut buf[..n]) } } } } - /// Detect NTFS partition offset by reading virtual disk data. - fn detect_ntfs_offset_virtual( - inner: &mut R, - layout: &VhdLayout, - virtual_disk_size: u64, - ) -> Result { - // Helper to read 4 bytes from a virtual offset - let read_magic = |inner: &mut R, layout: &VhdLayout, offset: u64| -> io::Result<[u8; 4]> { - let mut magic = [0u8; 4]; - match layout { - VhdLayout::Fixed => { - inner.seek(SeekFrom::Start(offset))?; - inner.read_exact(&mut magic)?; - } - VhdLayout::Dynamic { bat, block_size } => { - let bi = (offset / block_size) as usize; - let bo = offset % block_size; - if bi < bat.len() && bat[bi] != BAT_ENTRY_UNUSED { - let file_off = bat[bi] as u64 * 512 + BITMAP_SECTORS * 512 + bo; - inner.seek(SeekFrom::Start(file_off))?; - inner.read_exact(&mut magic)?; - } + /// Read 4 bytes from a virtual offset (for magic-byte probing). + fn read_virtual_u32(inner: &mut R, layout: &VhdLayout, offset: u64) -> io::Result<[u8; 4]> { + let mut buf = [0u8; 4]; + match layout { + VhdLayout::Fixed => { + inner.seek(SeekFrom::Start(offset))?; + inner.read_exact(&mut buf)?; + } + VhdLayout::Dynamic { 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(magic) - }; + } + Ok(buf) + } - // Stage 1: Try MBR - if virtual_disk_size >= 512 { - let mut mbr = [0u8; 512]; - match layout { - VhdLayout::Fixed => { - inner.seek(SeekFrom::Start(0))?; - inner.read_exact(&mut mbr)?; - } - VhdLayout::Dynamic { bat, block_size: _ } => { - if !bat.is_empty() && bat[0] != BAT_ENTRY_UNUSED { - let file_off = bat[0] as u64 * 512 + BITMAP_SECTORS * 512; - inner.seek(SeekFrom::Start(file_off))?; - inner.read_exact(&mut mbr)?; - } + /// Read a full sector from virtual offset 0. + fn read_first_sector(inner: &mut R, layout: &VhdLayout) -> io::Result<[u8; SECTOR_SIZE as usize]> { + let mut sector = [0u8; SECTOR_SIZE as usize]; + match layout { + VhdLayout::Fixed => { + inner.seek(SeekFrom::Start(0))?; + inner.read_exact(&mut sector)?; + } + VhdLayout::Dynamic { bat, .. } => { + if !bat.is_empty() && bat[0] != BAT_UNUSED { + inner.seek(SeekFrom::Start(bat[0] as u64 * SECTOR_SIZE + SECTOR_SIZE))?; + inner.read_exact(&mut sector)?; } } + } + Ok(sector) + } + + /// Find the byte offset of the NTFS partition within the virtual disk. + fn find_ntfs(inner: &mut R, layout: &VhdLayout, vsize: u64) -> Result { + // Try MBR partition table first + if vsize >= SECTOR_SIZE { + let mbr = Self::read_first_sector(inner, layout)?; if mbr[510..512] == MBR_SIGNATURE { - for i in 0..4 { + 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], - mbr[eo + 9], - mbr[eo + 10], - mbr[eo + 11], - ]); - let offset = lba as u64 * 512; - if offset + 4 <= virtual_disk_size { - if read_magic(inner, layout, offset)? == NTFS_MAGIC { - return Ok(offset); - } + 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 + && Self::read_virtual_u32(inner, layout, offset)? == NTFS_MAGIC + { + return Ok(offset); } } } } } - // Stage 2: Probe known offsets - for &offset in NTFS_PROBE_OFFSETS { - if offset + 4 > virtual_disk_size { - continue; - } - if read_magic(inner, layout, offset)? == NTFS_MAGIC { + // Probe common offsets + for offset in NTFS_PROBE_OFFSETS { + if offset + 4 <= vsize + && Self::read_virtual_u32(inner, layout, offset)? == NTFS_MAGIC + { return Ok(offset); } } @@ -294,20 +235,19 @@ impl VhdReader { Err(VhdError::NoNtfsPartition) } - pub fn ntfs_size(&self) -> u64 { - self.virtual_disk_size - self.ntfs_offset + fn ntfs_size(&self) -> u64 { + self.virtual_size - self.ntfs_offset } } impl Read for VhdReader { fn read(&mut self, buf: &mut [u8]) -> io::Result { - let remaining = self.ntfs_size() - self.pos; + let remaining = self.ntfs_size().saturating_sub(self.pos); if remaining == 0 { return Ok(0); } - let virtual_offset = self.ntfs_offset + self.pos; - let max_read = std::cmp::min(buf.len() as u64, remaining) as usize; - let n = self.read_virtual(virtual_offset, &mut buf[..max_read])?; + let cap = std::cmp::min(buf.len() as u64, remaining) as usize; + let n = self.read_virtual(self.ntfs_offset + self.pos, &mut buf[..cap])?; self.pos += n as u64; Ok(n) } @@ -315,73 +255,57 @@ impl Read for VhdReader { impl Seek for VhdReader { fn seek(&mut self, pos: SeekFrom) -> io::Result { - let new_pos = match pos { - SeekFrom::Start(offset) => offset, - SeekFrom::Current(offset) => { - let target = self.pos as i64 + offset; - if target < 0 { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - "cannot seek before start", - )); - } - target as u64 - } - SeekFrom::End(offset) => { - let target = self.ntfs_size() as i64 + offset; - if target < 0 { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - "cannot seek before start", - )); - } - target as 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, }; - self.pos = new_pos; - Ok(new_pos) + if target < 0 { + return Err(io::Error::new(io::ErrorKind::InvalidInput, "seek before start")); + } + self.pos = target as u64; + Ok(self.pos) } } // --------------------------------------------------------------------------- -// NTFS extraction from VHD +// NTFS extraction // --------------------------------------------------------------------------- -fn ntfs_time_to_system_time(ntfs_time: NtfsTime) -> SystemTime { - let intervals_since_windows_epoch = ntfs_time.nt_timestamp(); - let intervals_since_unix_epoch = intervals_since_windows_epoch - 116_444_736_000_000_000; - let nanos_since_unix_epoch = intervals_since_unix_epoch * 100; - SystemTime::UNIX_EPOCH + Duration::from_nanos(nanos_since_unix_epoch) +fn is_ntfs_system_entry(name: &str) -> bool { + name.starts_with('$') || name == "." || name == ".." || name == "System Volume Information" } -fn set_ntfs_timestamps( - fs: &mut T, - file: &ntfs::NtfsFile, - path: &Path, -) -> Result<()> { +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(fs: &mut T, file: &ntfs::NtfsFile, path: &Path) { let mut attrs = file.attributes(); - while let Some(attr) = attrs.next(fs) { - let attr = attr?; - let attr = attr.to_attribute()?; - if let Ok(NtfsAttributeType::StandardInformation) = attr.ty() { - let info = attr.resident_structured_value::()?; - let handle = OpenOptions::new().write(true).open(path)?; - handle.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; + 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::() { + 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; + } } } - Ok(()) } fn extract_ntfs_dir( ntfs: &Ntfs, fs: &mut T, dir: &ntfs::NtfsFile, - output_dir: &Path, + out: &Path, pb: &ProgressBar, ) -> Result<()> { let index = dir.directory_index(fs)?; @@ -389,52 +313,40 @@ fn extract_ntfs_dir( while let Some(entry) = iter.next(fs) { let entry = entry?; - let file_name = entry - .key() - .ok_or_else(|| anyhow!("missing index entry key"))?; - let file_name = file_name?; - let name = file_name.name().to_string_lossy(); - - if name.starts_with('$') - || name == "." - || name == ".." - || name == "System Volume Information" - { + 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_path = output_dir.join(&*name); + let dest = out.join(&*name); - if file_name.is_directory() { - create_dir_all(&dest_path)?; - extract_ntfs_dir(ntfs, fs, &file, &dest_path, pb)?; - set_ntfs_timestamps(fs, &file, &dest_path).ok(); - } else if let Some(data_item) = file.data(fs, "") { - let data_item = data_item?; - let data_attribute = data_item.to_attribute()?; - let mut data_value = - BufReader::with_capacity(256 * 1024, data_attribute.value(fs)?.attach(fs)); - - let mut output_file = File::create(&dest_path)?; + 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 buffer = data_value.fill_buf()?; - let length = buffer.len(); - if length == 0 { + let buf = reader.fill_buf()?; + if buf.is_empty() { break; } - output_file.write_all(buffer)?; - data_value.consume(length); - pb.inc(length as u64); + out_file.write_all(buf)?; + let n = buf.len(); + reader.consume(n); + pb.inc(n as u64); } - output_file.flush()?; - drop(data_value); - - set_ntfs_timestamps(fs, &file, &dest_path).ok(); + out_file.flush()?; + drop(reader); + set_ntfs_timestamps(fs, &file, &dest); } } - Ok(()) } @@ -443,77 +355,64 @@ fn calculate_ntfs_size( fs: &mut T, dir: &ntfs::NtfsFile, ) -> Result { - let mut total: u64 = 0; + 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 file_name = entry - .key() - .ok_or_else(|| anyhow!("missing index entry key"))?; - let file_name = file_name?; - let name = file_name.name().to_string_lossy(); - - if name.starts_with('$') - || name == "." - || name == ".." - || name == "System Volume Information" - { + 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 file_name.is_directory() { + if key.is_directory() { total += calculate_ntfs_size(ntfs, fs, &file)?; - } else if let Some(data_item) = file.data(fs, "") { - let data_item = data_item?; - let attr = data_item.to_attribute()?; - total += attr.value_length(); + } else if let Some(data) = file.data(fs, "") { + total += data?.to_attribute()?.value_length(); } } - Ok(total) } -/// Extract all files from a VHD's NTFS filesystem into a folder, then delete the VHD. +/// Extract all files from a VHD's NTFS filesystem, then delete the VHD. pub fn extract_vhd(vhd_path: &Path) -> Result<()> { let output_dir = vhd_path.with_extension(""); + println!("Extracting VHD: {}", vhd_path.display()); - println!("Extracting VHD contents..."); - println!(" VHD: {}", vhd_path.display()); - println!(" Output: {}", output_dir.display()); - - let file = File::open(vhd_path)?; - let mut vhd = VhdReader::new(file).map_err(|e| anyhow!(e))?; - + let mut vhd = VhdReader::new(File::open(vhd_path)?).map_err(|e| anyhow!(e))?; let mut ntfs = Ntfs::new(&mut vhd)?; ntfs.read_upcase_table(&mut vhd)?; let root = ntfs.root_directory(&mut vhd)?; - let total_size = calculate_ntfs_size(&ntfs, &mut vhd, &root)?; + let total = calculate_ntfs_size(&ntfs, &mut vhd, &root)?; - let pb = ProgressBar::new(total_size).with_style( - ProgressStyle::default_bar().template( - "{prefix} [{bar:20!.bright.yellow/dim.white}] {bytes:>8} [{elapsed}<{eta}, {bytes_per_sec}]", - )?, - ); - pb.set_prefix(format!( - "Extracting {}", - vhd_path.file_name().unwrap_or_default().to_string_lossy() - )); + let pb = ProgressBar::new(total) + .with_style(ProgressStyle::default_bar().template(PROGRESS_STYLE)?); + pb.set_prefix(vhd_path.file_name().unwrap_or_default().to_string_lossy().to_string()); create_dir_all(&output_dir)?; let root = ntfs.root_directory(&mut vhd)?; extract_ntfs_dir(&ntfs, &mut vhd, &root, &output_dir, &pb)?; - pb.finish(); + 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(()) } + +// --------------------------------------------------------------------------- +// 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()) +}