From 1cc38ad621df7556a9713fc431098a8264a54960 Mon Sep 17 00:00:00 2001 From: Luke Street Date: Thu, 7 Nov 2024 08:43:20 -0700 Subject: [PATCH] Add WAD support to VFS & `wad` commands --- Cargo.lock | 4 +- Cargo.toml | 4 +- README.md | 32 +++++ src/cmd/mod.rs | 1 + src/cmd/wad.rs | 108 ++++++++++++++ src/main.rs | 2 + src/util/mod.rs | 2 + src/util/read.rs | 26 ++++ src/util/wad.rs | 220 ++++++++++++++++++++++++++++ src/vfs/common.rs | 2 +- src/vfs/mod.rs | 19 ++- src/vfs/wad.rs | 359 ++++++++++++++++++++++++++++++++++++++++++++++ 12 files changed, 770 insertions(+), 9 deletions(-) create mode 100644 src/cmd/wad.rs create mode 100644 src/util/read.rs create mode 100644 src/util/wad.rs create mode 100644 src/vfs/wad.rs diff --git a/Cargo.lock b/Cargo.lock index b7af474..01c33dc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -348,14 +348,16 @@ dependencies = [ [[package]] name = "decomp-toolkit" -version = "1.2.0" +version = "1.3.0" dependencies = [ + "aes", "anyhow", "ar", "argp", "base16ct", "base64", "byteorder", + "cbc", "crossterm", "cwdemangle", "cwextab", diff --git a/Cargo.toml b/Cargo.toml index 3e05e3a..c741642 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,7 +3,7 @@ name = "decomp-toolkit" description = "Yet another GameCube/Wii decompilation toolkit." authors = ["Luke Street "] license = "MIT OR Apache-2.0" -version = "1.2.0" +version = "1.3.0" edition = "2021" publish = false repository = "https://github.com/encounter/decomp-toolkit" @@ -25,6 +25,7 @@ strip = "debuginfo" codegen-units = 1 [dependencies] +aes = "0.8" anyhow = { version = "1.0", features = ["backtrace"] } ar = { git = "https://github.com/bjorn3/rust-ar.git", branch = "write_symbol_table" } argp = "0.3" @@ -32,6 +33,7 @@ base16ct = "0.2" base64 = "0.22" byteorder = "1.5" typed-path = "0.9" +cbc = "0.1" crossterm = "0.28" cwdemangle = "1.0" cwextab = "1.0" diff --git a/README.md b/README.md index 00d7649..0922424 100644 --- a/README.md +++ b/README.md @@ -52,6 +52,9 @@ project structure and build system that uses decomp-toolkit under the hood. - [yay0 compress](#yay0-compress) - [yaz0 decompress](#yaz0-decompress) - [yaz0 compress](#yaz0-compress) + - [wad info](#wad-info) + - [wad extract](#wad-extract) + - [wad verify](#wad-verify) ## Goals @@ -474,6 +477,7 @@ Supported containers: - Disc images (see [disc info](#disc-info) for supported formats) - RARC archives (older .arc) - U8 archives (newer .arc) +- WAD files (Wii VC) Supported compression formats are handled transparently: - Yay0 (SZP) / Yaz0 (SZS) @@ -562,3 +566,31 @@ $ dtk yaz0 compress input.bin -o output.bin.yaz0 # or, for batch processing $ dtk yaz0 compress rels/* -o rels ``` + +### wad info + +Prints information about a WAD file. + +```shell +$ dtk wad info input.wad +``` + +### wad extract + +> [!NOTE] +> [vfs cp](#vfs-cp) is more flexible and supports WAD files. +> This command is now equivalent to `dtk vfs cp input.wad: output_dir` + +Extracts the contents of a WAD file. + +```shell +$ dtk wad extract input.wad -o output_dir +``` + +### wad verify + +Verifies the contents of a WAD file. + +```shell +$ dtk wad verify input.wad +``` diff --git a/src/cmd/mod.rs b/src/cmd/mod.rs index 7059750..e9a3e72 100644 --- a/src/cmd/mod.rs +++ b/src/cmd/mod.rs @@ -14,5 +14,6 @@ pub mod rso; pub mod shasum; pub mod u8_arc; pub mod vfs; +pub mod wad; pub mod yay0; pub mod yaz0; diff --git a/src/cmd/wad.rs b/src/cmd/wad.rs new file mode 100644 index 0000000..17eb8c2 --- /dev/null +++ b/src/cmd/wad.rs @@ -0,0 +1,108 @@ +use anyhow::Result; +use argp::FromArgs; +use size::Size; +use typed_path::Utf8NativePathBuf; + +use crate::{ + cmd::vfs, + util::{ + path::native_path, + wad::{process_wad, verify_wad}, + }, + vfs::open_file, +}; + +#[derive(FromArgs, PartialEq, Debug)] +/// Commands for processing Wii WAD files. +#[argp(subcommand, name = "wad")] +pub struct Args { + #[argp(subcommand)] + command: SubCommand, +} + +#[derive(FromArgs, PartialEq, Debug)] +#[argp(subcommand)] +enum SubCommand { + Extract(ExtractArgs), + Info(InfoArgs), + Verify(VerifyArgs), +} + +#[derive(FromArgs, PartialEq, Eq, Debug)] +/// Extracts WAD file contents. +#[argp(subcommand, name = "extract")] +pub struct ExtractArgs { + #[argp(positional, from_str_fn(native_path))] + /// WAD file + file: Utf8NativePathBuf, + #[argp(option, short = 'o', from_str_fn(native_path))] + /// output directory + output: Option, + #[argp(switch)] + /// Do not decompress files when copying. + no_decompress: bool, + #[argp(switch, short = 'q')] + /// Quiet output. Don't print anything except errors. + quiet: bool, +} + +#[derive(FromArgs, PartialEq, Eq, Debug)] +/// Views WAD file information. +#[argp(subcommand, name = "info")] +pub struct InfoArgs { + #[argp(positional, from_str_fn(native_path))] + /// WAD file + file: Utf8NativePathBuf, +} + +#[derive(FromArgs, PartialEq, Eq, Debug)] +/// Verifies WAD file integrity. +#[argp(subcommand, name = "verify")] +pub struct VerifyArgs { + #[argp(positional, from_str_fn(native_path))] + /// WAD file + file: Utf8NativePathBuf, +} + +pub fn run(args: Args) -> Result<()> { + match args.command { + SubCommand::Info(c_args) => info(c_args), + SubCommand::Verify(c_args) => verify(c_args), + SubCommand::Extract(c_args) => extract(c_args), + } +} + +fn info(args: InfoArgs) -> Result<()> { + let mut file = open_file(&args.file, true)?; + let wad = process_wad(file.as_mut())?; + println!("Title ID: {}", hex::encode(wad.ticket().title_id)); + println!("Title key: {}", hex::encode(wad.title_key)); + println!("Fake signed: {}", wad.fake_signed); + for content in wad.contents() { + println!( + "Content {:08x}: Offset {:#X}, size {}", + content.content_index.get(), + wad.content_offset(content.content_index.get()), + Size::from_bytes(content.size.get()) + ); + } + Ok(()) +} + +fn verify(args: VerifyArgs) -> Result<()> { + let mut file = open_file(&args.file, true)?; + let wad = process_wad(file.as_mut())?; + verify_wad(&wad, file.as_mut())?; + println!("Verification successful"); + Ok(()) +} + +fn extract(args: ExtractArgs) -> Result<()> { + let path = Utf8NativePathBuf::from(format!("{}:", args.file)); + let output = args.output.unwrap_or_else(|| Utf8NativePathBuf::from(".")); + vfs::cp(vfs::CpArgs { + paths: vec![path, output], + no_decompress: args.no_decompress, + quiet: args.quiet, + }) +} diff --git a/src/main.rs b/src/main.rs index 5d4a490..67a64b8 100644 --- a/src/main.rs +++ b/src/main.rs @@ -106,6 +106,7 @@ enum SubCommand { Vfs(cmd::vfs::Args), Yay0(cmd::yay0::Args), Yaz0(cmd::yaz0::Args), + Wad(cmd::wad::Args), } // Duplicated from supports-color so we can check early. @@ -181,6 +182,7 @@ fn main() { SubCommand::Vfs(c_args) => cmd::vfs::run(c_args), SubCommand::Yay0(c_args) => cmd::yay0::run(c_args), SubCommand::Yaz0(c_args) => cmd::yaz0::run(c_args), + SubCommand::Wad(c_args) => cmd::wad::run(c_args), }); if let Err(e) = result { eprintln!("Failed: {e:?}"); diff --git a/src/util/mod.rs b/src/util/mod.rs index 81d2398..b81fc92 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -18,6 +18,7 @@ pub mod nested; pub mod nlzss; pub mod path; pub mod rarc; +pub mod read; pub mod reader; pub mod rel; pub mod rso; @@ -25,6 +26,7 @@ pub mod signatures; pub mod split; pub mod take_seek; pub mod u8_arc; +pub mod wad; #[inline] pub const fn align_up(value: u32, align: u32) -> u32 { (value + (align - 1)) & !(align - 1) } diff --git a/src/util/read.rs b/src/util/read.rs new file mode 100644 index 0000000..419a699 --- /dev/null +++ b/src/util/read.rs @@ -0,0 +1,26 @@ +use std::{io, io::Read}; + +use zerocopy::{FromBytes, FromZeros, IntoBytes}; + +#[inline(always)] +pub fn read_from(reader: &mut R) -> io::Result +where + T: FromBytes + IntoBytes, + R: Read + ?Sized, +{ + let mut ret = ::new_zeroed(); + reader.read_exact(ret.as_mut_bytes())?; + Ok(ret) +} + +#[inline(always)] +pub fn read_box_slice(reader: &mut R, count: usize) -> io::Result> +where + T: FromBytes + IntoBytes, + R: Read + ?Sized, +{ + let mut ret = <[T]>::new_box_zeroed_with_elems(count) + .map_err(|_| io::Error::from(io::ErrorKind::OutOfMemory))?; + reader.read_exact(ret.as_mut().as_mut_bytes())?; + Ok(ret) +} diff --git a/src/util/wad.rs b/src/util/wad.rs new file mode 100644 index 0000000..4b6e8fa --- /dev/null +++ b/src/util/wad.rs @@ -0,0 +1,220 @@ +use std::{ + io, + io::{BufRead, Read, Seek}, +}; + +use aes::cipher::{BlockDecryptMut, KeyIvInit}; +use anyhow::{bail, Result}; +use nodtool::nod::{Ticket, TmdHeader}; +use sha1::{Digest, Sha1}; +use size::Size; +use zerocopy::{big_endian::*, FromBytes, FromZeros, Immutable, IntoBytes, KnownLayout}; + +use crate::{ + array_ref_mut, static_assert, + util::read::{read_box_slice, read_from}, +}; + +// TODO: other WAD types? +pub const WAD_MAGIC: [u8; 8] = [0x00, 0x00, 0x00, 0x20, 0x49, 0x73, 0x00, 0x00]; + +#[derive(Clone, Debug, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[repr(C, align(4))] +pub struct WadHeader { + pub header_size: U32, + pub wad_type: [u8; 0x2], + pub wad_version: U16, + pub cert_chain_size: U32, + pub _reserved1: [u8; 0x4], + pub ticket_size: U32, + pub tmd_size: U32, + pub data_size: U32, + pub footer_size: U32, +} + +static_assert!(size_of::() == 0x20); + +#[derive(Clone, Debug, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)] +#[repr(C, align(4))] +pub struct ContentMetadata { + pub content_id: U32, + pub content_index: U16, + pub content_type: U16, + pub size: U64, + pub hash: HashBytes, +} + +static_assert!(size_of::() == 0x24); + +impl ContentMetadata { + #[inline] + pub fn iv(&self) -> [u8; 0x10] { + let mut iv = [0u8; 0x10]; + *array_ref_mut!(iv, 0, 2) = self.content_index.get().to_be_bytes(); + iv + } +} + +const ALIGNMENT: usize = 0x40; + +#[inline(always)] +pub fn align_up(value: u64, alignment: u64) -> u64 { (value + alignment - 1) & !(alignment - 1) } + +pub type HashBytes = [u8; 20]; +pub type KeyBytes = [u8; 16]; + +type Aes128Cbc = cbc::Decryptor; + +#[derive(Debug, Clone)] +pub struct WadFile { + pub header: WadHeader, + pub title_key: KeyBytes, + pub fake_signed: bool, + pub raw_cert_chain: Box<[u8]>, + pub raw_ticket: Box<[u8]>, + pub raw_tmd: Box<[u8]>, + pub content_offset: u64, +} + +impl WadFile { + pub fn ticket(&self) -> &Ticket { + Ticket::ref_from_bytes(&self.raw_ticket).expect("Invalid ticket alignment") + } + + pub fn tmd(&self) -> &TmdHeader { + TmdHeader::ref_from_prefix(&self.raw_tmd).expect("Invalid TMD alignment").0 + } + + pub fn contents(&self) -> &[ContentMetadata] { + let (_, cmd_data) = + TmdHeader::ref_from_prefix(&self.raw_tmd).expect("Invalid TMD alignment"); + <[ContentMetadata]>::ref_from_bytes(cmd_data).expect("Invalid CMD alignment") + } + + pub fn content_offset(&self, content_index: u16) -> u64 { + let contents = self.contents(); + let mut offset = self.content_offset; + for content in contents.iter().take(content_index as usize) { + offset = align_up(offset + content.size.get(), ALIGNMENT as u64); + } + offset + } + + pub fn trailer_offset(&self) -> u64 { + let contents = self.contents(); + let mut offset = self.content_offset; + for content in contents.iter() { + offset = align_up(offset + content.size.get(), ALIGNMENT as u64); + } + offset + } +} + +pub fn process_wad(reader: &mut R) -> Result +where R: BufRead + Seek + ?Sized { + let header: WadHeader = read_from(reader)?; + let mut offset = align_up(header.header_size.get() as u64, ALIGNMENT as u64); + + reader.seek(io::SeekFrom::Start(offset))?; + let raw_cert_chain: Box<[u8]> = read_box_slice(reader, header.cert_chain_size.get() as usize)?; + offset = align_up(offset + header.cert_chain_size.get() as u64, ALIGNMENT as u64); + + reader.seek(io::SeekFrom::Start(offset))?; + let raw_ticket: Box<[u8]> = read_box_slice(reader, header.ticket_size.get() as usize)?; + offset = align_up(offset + header.ticket_size.get() as u64, ALIGNMENT as u64); + + reader.seek(io::SeekFrom::Start(offset))?; + let raw_tmd: Box<[u8]> = read_box_slice(reader, header.tmd_size.get() as usize)?; + offset = align_up(offset + header.tmd_size.get() as u64, ALIGNMENT as u64); + + let content_offset = offset; + let mut file = WadFile { + header, + title_key: [0; 16], + fake_signed: false, + raw_cert_chain, + raw_ticket, + raw_tmd, + content_offset, + }; + + let mut title_key_found = false; + if file.ticket().header.sig.iter().all(|&x| x == 0) { + // Fake signed, try to determine common key index + file.fake_signed = true; + let contents = file.contents(); + if let Some(smallest_content) = contents.iter().min_by_key(|x| x.size.get()) { + let mut ticket = file.ticket().clone(); + for i in 0..2 { + ticket.common_key_idx = i; + let title_key = ticket.decrypt_title_key()?; + let offset = file.content_offset(smallest_content.content_index.get()); + reader.seek(io::SeekFrom::Start(offset))?; + if verify_content(reader, smallest_content, &title_key)? { + file.title_key = title_key; + title_key_found = true; + break; + } + } + } + if !title_key_found { + bail!("Failed to determine title key for fake signed WAD"); + } + } + if !title_key_found { + let title_key = file.ticket().decrypt_title_key()?; + file.title_key = title_key; + } + + Ok(file) +} + +pub fn verify_wad(file: &WadFile, reader: &mut R) -> Result<()> +where R: Read + Seek + ?Sized { + for content in file.contents() { + let content_index = content.content_index.get(); + println!( + "Verifying content {:08x} (size {})", + content_index, + Size::from_bytes(content.size.get()) + ); + let offset = file.content_offset(content_index); + reader.seek(io::SeekFrom::Start(offset))?; + if !verify_content(reader, content, &file.title_key)? { + bail!("Content {:08x} hash mismatch", content_index); + } + } + Ok(()) +} + +fn verify_content( + reader: &mut R, + content: &ContentMetadata, + title_key: &KeyBytes, +) -> Result +where + R: Read + ?Sized, +{ + let mut buf = <[[u8; 0x10]]>::new_box_zeroed_with_elems(0x200) + .map_err(|_| io::Error::from(io::ErrorKind::OutOfMemory))?; + // Read full padded size for decryption + let read_size = align_up(content.size.get(), 0x40); + let mut decryptor = Aes128Cbc::new(title_key.into(), (&content.iv()).into()); + let mut digest = Sha1::default(); + let mut read = 0; + while read < read_size { + let len = buf.len().min(usize::try_from(read_size - read).unwrap_or(usize::MAX)); + debug_assert_eq!(len % 0x10, 0); + reader.read_exact(&mut buf.as_mut_bytes()[..len])?; + for block in buf.iter_mut().take(len / 0x10) { + decryptor.decrypt_block_mut(block.into()); + } + // Only hash up to content size + let hash_len = (read + len as u64).min(content.size.get()).saturating_sub(read) as usize; + if hash_len > 0 { + digest.update(&buf.as_bytes()[..hash_len]); + } + read += len as u64; + } + Ok(HashBytes::from(digest.finalize()) == content.hash) +} diff --git a/src/vfs/common.rs b/src/vfs/common.rs index 249cad0..e5ef0fa 100644 --- a/src/vfs/common.rs +++ b/src/vfs/common.rs @@ -112,7 +112,7 @@ impl Seek for WindowedFile { } #[inline] - fn stream_position(&mut self) -> io::Result { Ok(self.pos) } + fn stream_position(&mut self) -> io::Result { Ok(self.pos - self.begin) } } impl VfsFile for WindowedFile { diff --git a/src/vfs/mod.rs b/src/vfs/mod.rs index 5a1c2da..55c85c8 100644 --- a/src/vfs/mod.rs +++ b/src/vfs/mod.rs @@ -3,6 +3,7 @@ mod disc; mod rarc; mod std_fs; mod u8_arc; +mod wad; use std::{ error::Error, @@ -22,12 +23,14 @@ use rarc::RarcFs; pub use std_fs::StdFs; use typed_path::{Utf8NativePath, Utf8UnixPath, Utf8UnixPathBuf}; use u8_arc::U8Fs; +use wad::WadFs; use crate::util::{ ncompress::{YAY0_MAGIC, YAZ0_MAGIC}, nlzss, rarc::RARC_MAGIC, u8_arc::U8_MAGIC, + wad::WAD_MAGIC, }; pub trait Vfs: DynClone + Send + Sync { @@ -154,6 +157,7 @@ pub enum ArchiveKind { Rarc, U8, Disc(nod::Format), + Wad, } impl Display for ArchiveKind { @@ -162,6 +166,7 @@ impl Display for ArchiveKind { ArchiveKind::Rarc => write!(f, "RARC"), ArchiveKind::U8 => write!(f, "U8"), ArchiveKind::Disc(format) => write!(f, "Disc ({})", format), + ArchiveKind::Wad => write!(f, "WAD"), } } } @@ -169,18 +174,19 @@ impl Display for ArchiveKind { pub fn detect(file: &mut R) -> io::Result where R: Read + Seek + ?Sized { file.seek(SeekFrom::Start(0))?; - let mut magic = [0u8; 4]; + let mut magic = [0u8; 8]; match file.read_exact(&mut magic) { Ok(_) => {} Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(FileFormat::Regular), Err(e) => return Err(e), } - file.seek_relative(-4)?; + file.seek_relative(-8)?; match magic { - YAY0_MAGIC => Ok(FileFormat::Compressed(CompressionKind::Yay0)), - YAZ0_MAGIC => Ok(FileFormat::Compressed(CompressionKind::Yaz0)), - RARC_MAGIC => Ok(FileFormat::Archive(ArchiveKind::Rarc)), - U8_MAGIC => Ok(FileFormat::Archive(ArchiveKind::U8)), + _ if magic.starts_with(&YAY0_MAGIC) => Ok(FileFormat::Compressed(CompressionKind::Yay0)), + _ if magic.starts_with(&YAZ0_MAGIC) => Ok(FileFormat::Compressed(CompressionKind::Yaz0)), + _ if magic.starts_with(&RARC_MAGIC) => Ok(FileFormat::Archive(ArchiveKind::Rarc)), + _ if magic.starts_with(&U8_MAGIC) => Ok(FileFormat::Archive(ArchiveKind::U8)), + WAD_MAGIC => Ok(FileFormat::Archive(ArchiveKind::Wad)), _ => { let format = nod::Disc::detect(file)?; file.seek(SeekFrom::Start(0))?; @@ -332,6 +338,7 @@ pub fn open_fs(mut file: Box, kind: ArchiveKind) -> io::Result Ok(Box::new(WadFs::new(file)?)), } } diff --git a/src/vfs/wad.rs b/src/vfs/wad.rs new file mode 100644 index 0000000..0bed0ec --- /dev/null +++ b/src/vfs/wad.rs @@ -0,0 +1,359 @@ +use std::{ + io, + io::{BufRead, Cursor, Read, Seek, SeekFrom}, + sync::Arc, +}; + +use aes::cipher::{BlockDecryptMut, KeyIvInit}; +use filetime::FileTime; +use nodtool::nod::DiscStream; +use typed_path::Utf8UnixPath; +use zerocopy::FromZeros; + +use crate::{ + array_ref, + util::wad::{align_up, process_wad, ContentMetadata, WadFile}, + vfs::{ + common::{StaticFile, WindowedFile}, + Vfs, VfsError, VfsFile, VfsFileType, VfsMetadata, VfsResult, + }, +}; + +#[derive(Clone)] +pub struct WadFs { + file: Box, + wad: WadFile, +} + +enum WadFindResult<'a> { + Root, + Static(&'a [u8]), + Content(u16, &'a ContentMetadata), + Window(u64, u64), +} + +impl WadFs { + pub fn new(mut file: Box) -> io::Result { + let wad = process_wad(file.as_mut()) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + Ok(Self { file, wad }) + } + + fn find(&self, path: &str) -> Option { + let filename = path.trim_start_matches('/'); + if filename.contains('/') { + return None; + } + if filename.is_empty() { + return Some(WadFindResult::Root); + } + let filename = filename.to_ascii_lowercase(); + if let Some(id) = filename.strip_suffix(".app") { + if let Ok(content_index) = u16::from_str_radix(id, 16) { + if let Some(content) = self.wad.contents().get(content_index as usize) { + return Some(WadFindResult::Content(content_index, content)); + } + } + return None; + } + let title_id = hex::encode(self.wad.ticket().title_id); + match filename.strip_prefix(&title_id) { + Some(".tik") => Some(WadFindResult::Static(&self.wad.raw_ticket)), + Some(".tmd") => Some(WadFindResult::Static(&self.wad.raw_tmd)), + Some(".cert") => Some(WadFindResult::Static(&self.wad.raw_cert_chain)), + Some(".trailer") => { + if self.wad.header.footer_size.get() == 0 { + return None; + } + Some(WadFindResult::Window( + self.wad.trailer_offset(), + self.wad.header.footer_size.get() as u64, + )) + } + _ => None, + } + } +} + +impl Vfs for WadFs { + fn open(&mut self, path: &Utf8UnixPath) -> VfsResult> { + if let Some(result) = self.find(path.as_str()) { + match result { + WadFindResult::Root => Err(VfsError::IsADirectory), + WadFindResult::Static(data) => { + Ok(Box::new(StaticFile::new(Arc::from(data), self.file.metadata()?.mtime))) + } + WadFindResult::Content(content_index, content) => { + let offset = self.wad.content_offset(content_index); + Ok(Box::new(WadContent::new( + AesCbcStream::new( + self.file.clone(), + offset, + content.size.get(), + &self.wad.title_key, + &content.iv(), + ), + self.file.metadata()?.mtime, + ))) + } + WadFindResult::Window(offset, len) => { + Ok(Box::new(WindowedFile::new(self.file.clone(), offset, len)?)) + } + } + } else { + Err(VfsError::NotFound) + } + } + + fn exists(&mut self, path: &Utf8UnixPath) -> VfsResult { + Ok(self.find(path.as_str()).is_some()) + } + + fn read_dir(&mut self, path: &Utf8UnixPath) -> VfsResult> { + let path = path.as_str().trim_start_matches('/'); + if !path.is_empty() { + return Err(VfsError::NotFound); + } + let title_id = hex::encode(self.wad.ticket().title_id); + let mut entries = Vec::new(); + entries.push(format!("{}.tik", title_id)); + entries.push(format!("{}.tmd", title_id)); + entries.push(format!("{}.cert", title_id)); + if self.wad.header.footer_size.get() > 0 { + entries.push(format!("{}.trailer", title_id)); + } + for content in self.wad.contents() { + entries.push(format!("{:08x}.app", content.content_index.get())); + } + Ok(entries) + } + + fn metadata(&mut self, path: &Utf8UnixPath) -> VfsResult { + let mtime = self.file.metadata()?.mtime; + if let Some(result) = self.find(path.as_str()) { + match result { + WadFindResult::Root => { + Ok(VfsMetadata { file_type: VfsFileType::Directory, len: 0, mtime }) + } + WadFindResult::Static(data) => { + Ok(VfsMetadata { file_type: VfsFileType::File, len: data.len() as u64, mtime }) + } + WadFindResult::Content(_, content) => { + Ok(VfsMetadata { file_type: VfsFileType::File, len: content.size.get(), mtime }) + } + WadFindResult::Window(_, len) => { + Ok(VfsMetadata { file_type: VfsFileType::File, len, mtime }) + } + } + } else { + Err(VfsError::NotFound) + } + } +} + +#[derive(Clone)] +enum WadContentInner { + Stream(AesCbcStream), + Mapped(Cursor>), +} + +#[derive(Clone)] +struct WadContent { + inner: WadContentInner, + mtime: Option, +} + +impl WadContent { + fn new(inner: AesCbcStream, mtime: Option) -> Self { + Self { inner: WadContentInner::Stream(inner), mtime } + } + + fn convert_to_mapped(&mut self) { + match &mut self.inner { + WadContentInner::Stream(stream) => { + let pos = stream.stream_position().unwrap(); + stream.seek(SeekFrom::Start(0)).unwrap(); + let mut data = vec![0u8; stream.len() as usize]; + stream.read_exact(&mut data).unwrap(); + let mut cursor = Cursor::new(Arc::from(data.as_slice())); + cursor.set_position(pos); + self.inner = WadContentInner::Mapped(cursor); + } + WadContentInner::Mapped(_) => {} + }; + } +} + +impl BufRead for WadContent { + fn fill_buf(&mut self) -> io::Result<&[u8]> { + match &mut self.inner { + WadContentInner::Stream(stream) => stream.fill_buf(), + WadContentInner::Mapped(data) => data.fill_buf(), + } + } + + fn consume(&mut self, amt: usize) { + match &mut self.inner { + WadContentInner::Stream(stream) => stream.consume(amt), + WadContentInner::Mapped(data) => data.consume(amt), + } + } +} + +impl Read for WadContent { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + match &mut self.inner { + WadContentInner::Stream(stream) => stream.read(buf), + WadContentInner::Mapped(data) => data.read(buf), + } + } +} + +impl Seek for WadContent { + fn seek(&mut self, pos: SeekFrom) -> io::Result { + match &mut self.inner { + WadContentInner::Stream(stream) => stream.seek(pos), + WadContentInner::Mapped(data) => data.seek(pos), + } + } +} + +impl VfsFile for WadContent { + fn map(&mut self) -> io::Result<&[u8]> { + self.convert_to_mapped(); + match &mut self.inner { + WadContentInner::Stream(_) => unreachable!(), + WadContentInner::Mapped(data) => Ok(data.get_ref()), + } + } + + fn metadata(&mut self) -> io::Result { + match &mut self.inner { + WadContentInner::Stream(stream) => Ok(VfsMetadata { + file_type: VfsFileType::File, + len: stream.len(), + mtime: self.mtime, + }), + WadContentInner::Mapped(data) => Ok(VfsMetadata { + file_type: VfsFileType::File, + len: data.get_ref().len() as u64, + mtime: self.mtime, + }), + } + } + + fn into_disc_stream(self: Box) -> Box { self } +} + +#[derive(Clone)] +struct AesCbcStream { + inner: Box, + position: u64, + content_offset: u64, + content_size: u64, + key: [u8; 0x10], + init_iv: [u8; 0x10], + last_iv: [u8; 0x10], + block_idx: u64, + block: Box<[u8; 0x200]>, +} + +impl AesCbcStream { + fn new( + inner: Box, + content_offset: u64, + content_size: u64, + key: &[u8; 0x10], + iv: &[u8; 0x10], + ) -> Self { + let block = <[u8; 0x200]>::new_box_zeroed().unwrap(); + Self { + inner, + position: 0, + content_offset, + content_size, + key: *key, + init_iv: *iv, + last_iv: [0u8; 0x10], + block_idx: u64::MAX, + block, + } + } + + #[inline] + fn len(&self) -> u64 { self.content_size } + + #[inline] + fn remaining(&self) -> u64 { self.content_size.saturating_sub(self.position) } +} + +impl Read for AesCbcStream { + fn read(&mut self, mut buf: &mut [u8]) -> io::Result { + let mut total = 0; + while !buf.is_empty() { + let block = self.fill_buf()?; + if block.is_empty() { + break; + } + let len = buf.len().min(block.len()); + buf[..len].copy_from_slice(&block[..len]); + buf = &mut buf[len..]; + self.consume(len); + total += len; + } + Ok(total) + } +} + +impl BufRead for AesCbcStream { + fn fill_buf(&mut self) -> io::Result<&[u8]> { + if self.position >= self.content_size { + return Ok(&[]); + } + let block_size = self.block.len(); + let current_block = self.position / block_size as u64; + if current_block != self.block_idx { + let block_offset = current_block * block_size as u64; + let mut iv = [0u8; 0x10]; + if current_block == 0 { + // Use the initial IV for the first block + self.inner.seek(SeekFrom::Start(self.content_offset))?; + iv = self.init_iv; + } else if self.block_idx.checked_add(1) == Some(current_block) { + // Shortcut to avoid seeking when reading sequentially + iv = self.last_iv; + } else { + // Read the IV from the previous block + self.inner.seek(SeekFrom::Start(self.content_offset + block_offset - 0x10))?; + self.inner.read_exact(&mut iv)?; + } + let aligned_size = align_up(self.content_size, 0x10); + let remaining = aligned_size.saturating_sub(block_offset); + let read = remaining.min(block_size as u64) as usize; + self.inner.read_exact(&mut self.block[..read])?; + self.last_iv = *array_ref!(self.block, read - 0x10, 0x10); + let mut decryptor = + cbc::Decryptor::::new((&self.key).into(), (&iv).into()); + for aes_block in self.block[..read].chunks_exact_mut(0x10) { + decryptor.decrypt_block_mut(aes_block.into()); + } + self.block_idx = current_block; + } + let offset = (self.position % block_size as u64) as usize; + let len = self.remaining().min((block_size - offset) as u64) as usize; + Ok(&self.block[offset..offset + len]) + } + + fn consume(&mut self, amt: usize) { self.position = self.position.saturating_add(amt as u64); } +} + +impl Seek for AesCbcStream { + fn seek(&mut self, pos: SeekFrom) -> io::Result { + self.position = match pos { + SeekFrom::Start(p) => p, + SeekFrom::End(p) => self.content_size.saturating_add_signed(p), + SeekFrom::Current(p) => self.position.saturating_add_signed(p), + }; + Ok(self.position) + } +}