From e0ff976a116982226756199bdaf03d30d746a3c4 Mon Sep 17 00:00:00 2001 From: Eigeen Date: Mon, 28 Jul 2025 19:31:09 +0800 Subject: [PATCH] add Restore tool, fix metadata write error --- src/app.rs | 207 +++++++++++++++++++++++++++++++++++++++++++++--- src/metadata.rs | 15 +++- 2 files changed, 209 insertions(+), 13 deletions(-) diff --git a/src/app.rs b/src/app.rs index c26badd..25c3f4b 100644 --- a/src/app.rs +++ b/src/app.rs @@ -30,6 +30,24 @@ const FILE_NAME_LIST: &[u8] = include_bytes!("../assets/MHWs_STM_Release.list.zs const AUTO_CHUNK_SELECTION_SIZE_THRESHOLD: usize = 50 * 1024 * 1024; // 50MB const FALSE_TRUE_SELECTION: [&str; 2] = ["False", "True"]; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Mode { + Automatic = 0, + Manual = 1, + Restore = 2, +} + +impl Mode { + fn from_index(index: usize) -> color_eyre::Result { + match index { + 0 => Ok(Mode::Automatic), + 1 => Ok(Mode::Manual), + 2 => Ok(Mode::Restore), + _ => bail!("Invalid mode index: {index}"), + } + } +} + struct ChunkSelection { chunk_name: ChunkName, file_size: u64, @@ -60,15 +78,15 @@ impl App { // Mode selection let mode = Select::with_theme(&ColorfulTheme::default()) .with_prompt("Select mode") - .items(&["Automatic", "Manual"]) + .items(&["Automatic", "Manual", "Restore"]) .default(0) .interact()?; - let auto_mode = mode == 0; + let mode = Mode::from_index(mode)?; - if auto_mode { - self.auto_mode() - } else { - self.manual_mode() + match mode { + Mode::Automatic => self.auto_mode(), + Mode::Manual => self.manual_mode(), + Mode::Restore => self.restore_mode(), } } @@ -111,13 +129,12 @@ impl App { .truncate(true) .write(true) .open(output_path)?; - let mut pak_writer = ree_pak_core::write::PakWriter::new(out_file, entries.len() as u64); + // +1 for metadata + let mut pak_writer = + ree_pak_core::write::PakWriter::new(out_file, (entries.len() as u64) + 1); // write metadata - let metadata = PakMetadata { - version: 1, - is_uncompressed_patch: true, - }; + let metadata = PakMetadata::new(use_full_package_mode); metadata.write_to_pak(&mut pak_writer)?; let pak_writer_mtx = Arc::new(Mutex::new(pak_writer)); @@ -410,6 +427,174 @@ I'm sure I've checked the list, press Enter to continue"#, Ok(()) } + + fn restore_mode(&mut self) -> color_eyre::Result<()> { + let current_dir = std::env::current_dir()?; + + let game_dir: String = Input::::with_theme(&ColorfulTheme::default()) + .show_default(true) + .default(current_dir.to_string_lossy().to_string()) + .with_prompt("Input MonsterHunterWilds directory path") + .interact_text() + .unwrap() + .trim_matches(|c| c == '\"' || c == '\'') + .to_string(); + + let game_dir = Path::new(&game_dir); + if !game_dir.is_dir() { + bail!("game directory not exists."); + } + + // scan all pak files, find files generated by this tool + println!("Scanning tool generated files..."); + let dir = fs::read_dir(game_dir)?; + let mut tool_generated_files = Vec::new(); + let mut backup_files = Vec::new(); + let mut all_chunks = Vec::new(); + + for entry in dir { + let entry = entry?; + if !entry.file_type()?.is_file() { + continue; + } + + let file_name = entry.file_name().to_string_lossy().to_string(); + let file_path = entry.path(); + + // check backup files + if file_name.ends_with(".pak.backup") { + backup_files.push(file_path); + continue; + } + + // check pak files + if !file_name.ends_with(".pak") || !file_name.starts_with("re_chunk_") { + continue; + } + + // collect chunk info + if let Ok(chunk_name) = ChunkName::try_from_str(&file_name) { + all_chunks.push(chunk_name.clone()); + } + + // check if the file is generated by this tool + if let Ok(Some(metadata)) = self.check_tool_generated_file(&file_path) { + tool_generated_files.push((file_path, metadata)); + } + } + + if tool_generated_files.is_empty() && backup_files.is_empty() { + println!("No files found to restore."); + return Ok(()); + } + + println!( + "Found {} tool generated files and {} backup files", + tool_generated_files.len(), + backup_files.len() + ); + + // restore + let mut patch_files_to_remove = Vec::new(); + for (file_path, metadata) in &tool_generated_files { + if metadata.is_full_package() { + // restore full package mode (replace mode) + // this is a replace mode generated file, find the corresponding backup file + let backup_path = file_path.with_extension("pak.backup"); + if backup_path.exists() { + println!("Restore replace mode file: {}", file_path.display()); + + // delete the current file and restore the backup + fs::remove_file(file_path)?; + fs::rename(&backup_path, file_path)?; + + println!(" Restore backup file: {}", backup_path.display()); + } else { + println!("Warning: backup file not found {}", backup_path.display()); + } + } else { + // restore patch mode + // this is a patch mode generated file + if let Ok(chunk_name) = + ChunkName::try_from_str(&file_path.file_name().unwrap().to_string_lossy()) + { + patch_files_to_remove.push((file_path.clone(), chunk_name)); + } + } + } + + // remove patch files + if !patch_files_to_remove.is_empty() { + println!("Remove patch files..."); + + for (file_path, chunk_name) in patch_files_to_remove.iter().rev() { + println!("Remove patch file: {}", file_path.display()); + + // Check if there are any patches with higher numbers + let has_higher_patches = all_chunks.iter().any(|c| { + c.major_id == chunk_name.major_id + && c.sub_id == chunk_name.sub_id + && match (c.sub_id, c.sub_patch_id) { + (Some(_), Some(patch_id)) => { + patch_id > chunk_name.sub_patch_id.unwrap() + } + (None, Some(patch_id)) => patch_id > chunk_name.patch_id.unwrap(), + _ => false, + } + }); + + if has_higher_patches { + // create an empty patch file instead of deleting, to keep the patch sequence continuous + self.create_empty_patch_file(file_path)?; + println!(" Create empty patch file to keep sequence continuous"); + } else { + // no higher patches exist, safe to delete + fs::remove_file(file_path)?; + println!(" Removed patch file"); + } + } + } + + println!("Restore completed!"); + Ok(()) + } + + /// check if the file is generated by this tool, return metadata + fn check_tool_generated_file( + &self, + file_path: &Path, + ) -> color_eyre::Result> { + let file = match fs::File::open(file_path) { + Ok(file) => file, + Err(_) => return Ok(None), + }; + + let mut reader = io::BufReader::new(file); + let pak_archive = match ree_pak_core::read::read_archive(&mut reader) { + Ok(archive) => archive, + Err(_) => return Ok(None), + }; + + PakMetadata::from_pak_archive(reader, &pak_archive) + } + + /// create an empty patch file + fn create_empty_patch_file(&self, file_path: &Path) -> color_eyre::Result<()> { + let out_file = OpenOptions::new() + .create(true) + .truncate(true) + .write(true) + .open(file_path)?; + + let mut pak_writer = ree_pak_core::write::PakWriter::new(out_file, 1); + + // write metadata to mark this is an empty patch file + let metadata = PakMetadata::new(false); + metadata.write_to_pak(&mut pak_writer)?; + + pak_writer.finish()?; + Ok(()) + } } fn is_tex_file(hash: u64, file_name_table: &FileNameTable) -> bool { diff --git a/src/metadata.rs b/src/metadata.rs index 290a6c9..9d0e8cd 100644 --- a/src/metadata.rs +++ b/src/metadata.rs @@ -14,11 +14,22 @@ const METADATA_KEY: &str = "__TEX_DECOMPRESSOR_METADATA__"; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PakMetadata { - pub version: u32, - pub is_uncompressed_patch: bool, + version: u32, + is_full_package: bool, } impl PakMetadata { + pub fn new(is_full_package: bool) -> Self { + Self { + version: 1, + is_full_package, + } + } + + pub fn is_full_package(&self) -> bool { + self.is_full_package + } + pub fn from_pak_archive( reader: R, pak_archive: &PakArchive,