diff --git a/src/app.rs b/src/app.rs index d20d934..6576355 100644 --- a/src/app.rs +++ b/src/app.rs @@ -2,8 +2,8 @@ use std::{ io::{self, Write}, path::Path, sync::{ - atomic::{AtomicUsize, Ordering}, Arc, + atomic::{AtomicUsize, Ordering}, }, time::Duration, }; @@ -11,8 +11,7 @@ use std::{ use fs_err as fs; use color_eyre::eyre::bail; -use colored::Colorize; -use dialoguer::{theme::ColorfulTheme, Input, Select}; +use dialoguer::{Input, MultiSelect, Select, theme::ColorfulTheme}; use fs::OpenOptions; use indicatif::{HumanBytes, ProgressBar, ProgressStyle}; use parking_lot::Mutex; @@ -25,30 +24,27 @@ use ree_pak_core::{ write::FileOptions, }; -const FILE_NAME_LIST: &[u8] = include_bytes!("../assets/MHWs_STM_Release.list.zst"); +use crate::{chunk::ChunkName, util::human_bytes}; -#[derive(Debug, Clone)] -pub struct Config { - pub input_path: Option, - pub output_path: Option, - pub use_full_package_mode: bool, - pub use_feature_clone: bool, +const FILE_NAME_LIST: &[u8] = include_bytes!("../assets/MHWs_STM_Release.list.zst"); +const AUTO_CHUNK_SELECTION_SIZE_THRESHOLD: usize = 50 * 1024 * 1024; // 50MB +const FALSE_TRUE_SELECTION: [&str; 2] = ["False", "True"]; + +struct ChunkSelection { + chunk_name: ChunkName, + file_size: u64, } -impl Default for Config { - fn default() -> Self { - Self { - input_path: Default::default(), - output_path: Default::default(), - use_full_package_mode: false, - use_feature_clone: true, - } +impl std::fmt::Display for ChunkSelection { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{} ({})", self.chunk_name, human_bytes(self.file_size))?; + Ok(()) } } #[derive(Default)] pub struct App { - config: Config, + filename_table: Option, } impl App { @@ -57,6 +53,10 @@ impl App { println!("Get updates at https://github.com/eigeen/mhws-tex-decompressor"); println!(); + println!("Loading embedded file name table..."); + let filename_table = FileNameTable::from_bytes(FILE_NAME_LIST)?; + self.filename_table = Some(filename_table); + // Mode selection let mode = Select::with_theme(&ColorfulTheme::default()) .with_prompt("Select mode") @@ -72,47 +72,19 @@ impl App { } } - fn auto_mode(&mut self) -> color_eyre::Result<()> { - todo!() + fn filename_table(&self) -> &FileNameTable { + self.filename_table.as_ref().unwrap() } - fn manual_mode(&mut self) -> color_eyre::Result<()> { - let input: String = Input::with_theme(&ColorfulTheme::default()) - .show_default(true) - .default("re_chunk_000.pak.sub_000.pak".to_string()) - .with_prompt("Input .pak file path") - .interact_text() - .unwrap() - .trim_matches(|c| c == '\"' || c == '\'') - .to_string(); - - let input_path = Path::new(&input); - if !input_path.is_file() { - bail!("input file not exists."); - } - - const FALSE_TRUE_SELECTION: [&str; 2] = ["False", "True"]; - - let use_full_package_mode = Select::with_theme(&ColorfulTheme::default()) - .with_prompt( - "Package all files, including non-tex files (for replacing original files)", - ) - .default(0) - .items(&FALSE_TRUE_SELECTION) - .interact() - .unwrap(); - let use_full_package_mode = use_full_package_mode == 1; - - let use_feature_clone = Select::with_theme(&ColorfulTheme::default()) - .with_prompt("Clone feature flags from original file?") - .default(1) - .items(&FALSE_TRUE_SELECTION) - .interact() - .unwrap(); - let use_feature_clone = use_feature_clone == 1; - - println!("Loading embedded file name table..."); - let filename_table = FileNameTable::from_bytes(FILE_NAME_LIST)?; + fn process_chunk( + &self, + filename_table: &FileNameTable, + input_path: &Path, + output_path: &Path, + use_full_package_mode: bool, + use_feature_clone: bool, + ) -> color_eyre::Result<()> { + println!("Processing chunk: {}", input_path.display()); let file = fs::File::open(input_path)?; let mut reader = io::BufReader::new(file); @@ -135,8 +107,6 @@ impl App { }; // new pak archive - let output_path = input_path.with_extension("uncompressed.pak"); - println!("Output file: {}", output_path.to_string_lossy()); let out_file = OpenOptions::new() .create(true) .truncate(true) @@ -218,13 +188,218 @@ impl App { }; bar.finish(); - println!("{}", "Done!".cyan().bold()); - if !use_full_package_mode { - println!( - "You should rename the output file like `re_chunk_000.pak.sub_000.pak.patch_xxx.pak`." - ); + + Ok(()) + } + + fn auto_mode(&mut self) -> color_eyre::Result<()> { + let current_dir = std::env::current_dir()?; + + wait_for_enter( + r#"Check list: + +1. Your game is already updated to the latest version. +2. Uninstalled all the mods, or the generated files will break mods. + +I'm sure I've checked the list, press Enter to continue"#, + ); + + 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 for pak files + let dir = fs::read_dir(game_dir)?; + let mut all_chunks: Vec = vec![]; + 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(); + if !file_name.ends_with(".pak") || !file_name.starts_with("re_chunk_") { + continue; + } + + let chunk_name = match ChunkName::try_from_str(&file_name) { + Ok(chunk_name) => chunk_name, + Err(e) => { + println!("Invalid chunk name, skipped: {e}"); + continue; + } + }; + all_chunks.push(chunk_name); + } + all_chunks.sort(); + + // show chunks for selection + // only show sub chunks + let chunk_selections = all_chunks + .iter() + .filter_map(|chunk| { + if chunk.sub_id.is_some() { + Some(chunk.to_string()) + } else { + None + } + }) + .map(|file_name| { + let file_path = game_dir.join(&file_name); + let file_size = fs::metadata(file_path)?.len(); + Ok(ChunkSelection { + chunk_name: ChunkName::try_from_str(&file_name)?, + file_size, + }) + }) + .collect::>>()?; + if chunk_selections.is_empty() { + bail!("No available pak files found."); + } + + let selected_chunks: Vec = chunk_selections + .iter() + .map(|chunk_selection| { + Ok(chunk_selection.file_size >= AUTO_CHUNK_SELECTION_SIZE_THRESHOLD as u64) + }) + .collect::>>()?; + + let selected_chunks: Option> = + MultiSelect::with_theme(&ColorfulTheme::default()) + .with_prompt("Select chunks to process (Space to select, Enter to confirm)") + .items(&chunk_selections) + .defaults(&selected_chunks) + .interact_opt()?; + let Some(selected_chunks) = selected_chunks else { + bail!("No chunks selected."); + }; + + let selected_chunks = selected_chunks + .iter() + .map(|i| chunk_selections[*i].chunk_name.clone()) + .collect::>(); + + // replace mode: replace original files with uncompressed files + // patch mode: generate patch files after original patch files + let use_replace_mode = Select::with_theme(&ColorfulTheme::default()) + .with_prompt( + "Replace original files with uncompressed files? (Will automatically backup original files)", + ) + .default(0) + .items(&FALSE_TRUE_SELECTION) + .interact() + .unwrap(); + let use_replace_mode = use_replace_mode == 1; + + // start processing + for chunk_name in selected_chunks { + let chunk_path = game_dir.join(chunk_name.to_string()); + let output_path = if use_replace_mode { + // 如果是替换模式，首先生成一个临时的解压后文件 + let temp_path = chunk_path.with_extension("pak.temp"); + // 备份原始文件 + let backup_path = chunk_path.with_extension("pak.backup"); + if backup_path.exists() { + fs::remove_file(&backup_path)?; + } + fs::rename(&chunk_path, &backup_path)?; + temp_path + } else { + // 如果是补丁模式 + // 查找当前chunk系列的最大patch id + let max_patch_id = all_chunks + .iter() + .filter(|c| { + c.major_id == chunk_name.major_id + && c.patch_id == chunk_name.patch_id + && c.sub_id == chunk_name.sub_id + }) + .filter_map(|c| c.sub_patch_id) + .max() + .unwrap_or(0); + + let new_patch_id = max_patch_id + 1; + + // 创建新的chunk name + let mut output_chunk_name = chunk_name.clone(); + output_chunk_name.sub_patch_id = Some(new_patch_id); + + // 在chunk列表中添加新的补丁，以便后续处理可以找到它 + all_chunks.push(output_chunk_name.clone()); + + game_dir.join(output_chunk_name.to_string()) + }; + + println!("Output patch file: {}", output_path.display()); + self.process_chunk( + self.filename_table(), + &chunk_path, + &output_path, + use_replace_mode, + true, + )?; + + // 如果是替换模式，将临时文件重命名为原始文件名 + if use_replace_mode { + fs::rename(&output_path, &chunk_path)?; + } + println!(); + } + + Ok(()) + } + + fn manual_mode(&mut self) -> color_eyre::Result<()> { + let input: String = Input::with_theme(&ColorfulTheme::default()) + .show_default(true) + .default("re_chunk_000.pak.sub_000.pak".to_string()) + .with_prompt("Input .pak file path") + .interact_text() + .unwrap() + .trim_matches(|c| c == '\"' || c == '\'') + .to_string(); + + let input_path = Path::new(&input); + if !input_path.is_file() { + bail!("input file not exists."); + } + + let use_full_package_mode = Select::with_theme(&ColorfulTheme::default()) + .with_prompt( + "Package all files, including non-tex files (for replacing original files)", + ) + .default(0) + .items(&FALSE_TRUE_SELECTION) + .interact() + .unwrap(); + let use_full_package_mode = use_full_package_mode == 1; + + let use_feature_clone = Select::with_theme(&ColorfulTheme::default()) + .with_prompt("Clone feature flags from original file?") + .default(1) + .items(&FALSE_TRUE_SELECTION) + .interact() + .unwrap(); + let use_feature_clone = use_feature_clone == 1; + + self.process_chunk( + self.filename_table(), + input_path, + &input_path.with_extension("uncompressed.pak"), + use_full_package_mode, + use_feature_clone, + )?; + Ok(()) } } @@ -255,9 +430,9 @@ where Ok(data.len()) } -fn wait_for_exit() { +fn wait_for_enter(msg: &str) { let _: String = Input::with_theme(&ColorfulTheme::default()) - .with_prompt("Press Enter to exit") + .with_prompt(msg) .allow_empty(true) .interact_text() .unwrap(); diff --git a/src/chunk.rs b/src/chunk.rs index 75c3c65..4f05b66 100644 --- a/src/chunk.rs +++ b/src/chunk.rs @@ -1,5 +1,174 @@ //! Chunk file name format +//! +//! File name structure: +//! - Base: re_chunk_XXX.pak +//! - Patch: re_chunk_XXX.pak.patch_XXX.pak +//! - Sub: re_chunk_XXX.pak.sub_XXX.pak +//! - Sub Patch: re_chunk_XXX.pak.sub_XXX.pak.patch_XXX.pak +use color_eyre::eyre; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct ChunkName { - major_id: u32, + /// Major chunk ID (XXX in re_chunk_XXX.pak) + pub major_id: u32, + /// Patch number (XXX in .patch_XXX.pak) + pub patch_id: Option, + /// Sub chunk ID (XXX in .sub_XXX.pak) + pub sub_id: Option, + /// Patch number for sub chunk (YYY in .sub_XXX.pak.patch_YYY.pak) + pub sub_patch_id: Option, +} + +impl ChunkName { + /// Create a new base chunk name (re_chunk_XXX.pak) + pub fn new(major_id: u32) -> Self { + Self { + major_id, + patch_id: None, + sub_id: None, + sub_patch_id: None, + } + } + + /// Create a chunk name from a string + pub fn try_from_str(name: &str) -> color_eyre::Result { + let dot_parts = name.split('.').collect::>(); + if dot_parts.len() < 2 || dot_parts.len() % 2 != 0 { + return Err(eyre::eyre!( + "Invalid chunk name with odd number of parts: {}", + name + )); + } + + // every 2 parts is a component + let components = dot_parts + .chunks_exact(2) + .map(|c| (c[0], c[1])) + .collect::>(); + // check if all parts have the correct extension + if !components.iter().all(|(_, ext)| *ext == "pak") { + return Err(eyre::eyre!( + "Invalid chunk name with invalid extension: {}", + name + )); + } + + let mut this = Self::new(0); + + for (name, _) in components.iter() { + let component = Self::parse_component(name)?; + match component { + Component::Major(id) => this.major_id = id, + Component::Sub(id) => this.sub_id = Some(id), + Component::Patch(id) => { + if this.sub_id.is_some() { + this.sub_patch_id = Some(id); + } else { + this.patch_id = Some(id); + } + } + } + } + + Ok(this) + } + + fn parse_component(name: &str) -> color_eyre::Result { + if name.starts_with("re_chunk_") { + let major_id = name + .strip_prefix("re_chunk_") + .unwrap() + .parse::() + .map_err(|e| eyre::eyre!("Chunk name with invalid major ID: {}", e))?; + Ok(Component::Major(major_id)) + } else if name.starts_with("patch_") { + let patch_id = name + .strip_prefix("patch_") + .unwrap() + .parse::() + .map_err(|e| eyre::eyre!("Chunk name with invalid patch ID: {}", e))?; + Ok(Component::Patch(patch_id)) + } else if name.starts_with("sub_") { + let sub_id = name + .strip_prefix("sub_") + .unwrap() + .parse::() + .map_err(|e| eyre::eyre!("Chunk name with invalid sub ID: {}", e))?; + Ok(Component::Sub(sub_id)) + } else { + Err(eyre::eyre!( + "Invalid chunk name with invalid component: {}", + name + )) + } + } +} + +impl std::fmt::Display for ChunkName { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "re_chunk_{:03}.pak", self.major_id)?; + if let Some(patch_id) = self.patch_id { + write!(f, ".patch_{:03}.pak", patch_id)?; + return Ok(()); + } + if let Some(sub_id) = self.sub_id { + write!(f, ".sub_{:03}.pak", sub_id)?; + } + if let Some(sub_patch_id) = self.sub_patch_id { + write!(f, ".patch_{:03}.pak", sub_patch_id)?; + } + + Ok(()) + } +} + +impl PartialOrd for ChunkName { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for ChunkName { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.major_id + .cmp(&other.major_id) + .then(self.sub_id.cmp(&other.sub_id)) + .then(self.patch_id.cmp(&other.patch_id)) + .then(self.sub_patch_id.cmp(&other.sub_patch_id)) + } +} + +enum Component { + Major(u32), + Patch(u32), + Sub(u32), +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_chunk_name_formats() { + // Test base chunk + let base = ChunkName::new(0); + assert_eq!(base.to_string(), "re_chunk_000.pak"); + + // Test patch chunk + let patch = ChunkName::try_from_str("re_chunk_000.pak.patch_001.pak").unwrap(); + assert_eq!(patch.to_string(), "re_chunk_000.pak.patch_001.pak"); + + // Test sub chunk + let sub = ChunkName::try_from_str("re_chunk_000.pak.sub_000.pak").unwrap(); + assert_eq!(sub.to_string(), "re_chunk_000.pak.sub_000.pak"); + + // Test sub patch chunk + let sub_patch = + ChunkName::try_from_str("re_chunk_000.pak.sub_000.pak.patch_001.pak").unwrap(); + assert_eq!( + sub_patch.to_string(), + "re_chunk_000.pak.sub_000.pak.patch_001.pak" + ); + } } diff --git a/src/main.rs b/src/main.rs index 2030203..66622a2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,32 +1,9 @@ mod app; mod chunk; +mod util; -use std::{ - io::{self, Write}, - path::Path, - sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, - }, - time::Duration, -}; - -use fs_err as fs; - -use color_eyre::eyre::bail; use colored::Colorize; -use dialoguer::{Input, Select, theme::ColorfulTheme}; -use fs::OpenOptions; -use indicatif::{HumanBytes, ProgressBar, ProgressStyle}; -use parking_lot::Mutex; -use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; -use re_tex::tex::Tex; -use ree_pak_core::{ - filename::{FileNameExt, FileNameTable}, - pak::PakEntry, - read::archive::PakArchiveReader, - write::FileOptions, -}; +use dialoguer::{Input, theme::ColorfulTheme}; fn main() { std::panic::set_hook(Box::new(panic_hook)); @@ -38,17 +15,6 @@ fn main() { std::process::exit(1); } wait_for_exit(); - - // println!("Version v{} - Tool by @Eigeen", env!("CARGO_PKG_VERSION")); - // println!("Get updates at https://github.com/eigeen/mhws-tex-decompressor"); - // println!(); - - // if let Err(e) = main_entry() { - // eprintln!("{}: {:#}", "Error".red().bold(), e); - // wait_for_exit(); - // std::process::exit(1); - // } - // wait_for_exit(); } fn panic_hook(info: &std::panic::PanicHookInfo) { @@ -57,181 +23,6 @@ fn panic_hook(info: &std::panic::PanicHookInfo) { std::process::exit(1); } -// fn main_entry() -> color_eyre::Result<()> { -// let input: String = Input::with_theme(&ColorfulTheme::default()) -// .show_default(true) -// .default("re_chunk_000.pak.sub_000.pak".to_string()) -// .with_prompt("Input .pak file path") -// .interact_text() -// .unwrap() -// .trim_matches(|c| c == '\"' || c == '\'') -// .to_string(); - -// let input_path = Path::new(&input); -// if !input_path.is_file() { -// bail!("input file not exists."); -// } - -// const FALSE_TRUE_SELECTION: [&str; 2] = ["False", "True"]; - -// let use_full_package_mode = Select::with_theme(&ColorfulTheme::default()) -// .with_prompt("Package all files, including non-tex files (for replacing original files)") -// .default(0) -// .items(&FALSE_TRUE_SELECTION) -// .interact() -// .unwrap(); -// let use_full_package_mode = use_full_package_mode == 1; - -// let use_feature_clone = Select::with_theme(&ColorfulTheme::default()) -// .with_prompt("Clone feature flags from original file?") -// .default(1) -// .items(&FALSE_TRUE_SELECTION) -// .interact() -// .unwrap(); -// let use_feature_clone = use_feature_clone == 1; - -// println!("Loading embedded file name table..."); -// let filename_table = FileNameTable::from_bytes(FILE_NAME_LIST)?; - -// let file = fs::File::open(input_path)?; -// let mut reader = io::BufReader::new(file); - -// println!("Reading pak archive..."); -// let pak_archive = ree_pak_core::read::read_archive(&mut reader)?; -// let archive_reader = PakArchiveReader::new(reader, &pak_archive); -// let archive_reader_mtx = Mutex::new(archive_reader); - -// // filtered entries -// let entries = if use_full_package_mode { -// pak_archive.entries().iter().collect::>() -// } else { -// println!("Filtering entries..."); -// pak_archive -// .entries() -// .iter() -// .filter(|entry| is_tex_file(entry.hash(), &filename_table)) -// .collect::>() -// }; - -// // new pak archive -// let output_path = input_path.with_extension("uncompressed.pak"); -// println!("Output file: {}", output_path.to_string_lossy()); -// let out_file = OpenOptions::new() -// .create(true) -// .truncate(true) -// .write(true) -// .open(output_path)?; -// let pak_writer = ree_pak_core::write::PakWriter::new(out_file, entries.len() as u64); -// let pak_writer_mtx = Arc::new(Mutex::new(pak_writer)); - -// let bar = ProgressBar::new(entries.len() as u64); -// bar.set_style( -// ProgressStyle::default_bar().template("Bytes written: {msg}\n{pos}/{len} {wide_bar}")?, -// ); -// bar.enable_steady_tick(Duration::from_millis(200)); - -// let pak_writer_mtx1 = Arc::clone(&pak_writer_mtx); -// let bar1 = bar.clone(); -// let bytes_written = AtomicUsize::new(0); -// let err = entries -// .par_iter() -// .try_for_each(move |&entry| -> color_eyre::Result<()> { -// let pak_writer_mtx = &pak_writer_mtx1; -// let bar = &bar1; -// // read raw tex file -// // parse tex file -// let mut entry_reader = { -// let mut archive_reader = archive_reader_mtx.lock(); -// archive_reader.owned_entry_reader(entry.clone())? -// }; - -// if !is_tex_file(entry.hash(), &filename_table) { -// // plain file, just copy -// let mut buf = vec![]; -// std::io::copy(&mut entry_reader, &mut buf)?; -// let mut pak_writer = pak_writer_mtx.lock(); -// let write_bytes = write_to_pak( -// &mut pak_writer, -// entry, -// entry.hash(), -// &buf, -// use_feature_clone, -// )?; -// bytes_written.fetch_add(write_bytes, Ordering::SeqCst); -// } else { -// let mut tex = Tex::from_reader(&mut entry_reader)?; -// // decompress mipmaps -// tex.batch_decompress()?; - -// let tex_bytes = tex.as_bytes()?; -// let mut pak_writer = pak_writer_mtx.lock(); -// let write_bytes = write_to_pak( -// &mut pak_writer, -// entry, -// entry.hash(), -// &tex_bytes, -// use_feature_clone, -// )?; -// bytes_written.fetch_add(write_bytes, Ordering::SeqCst); -// } - -// bar.inc(1); -// if bar.position() % 100 == 0 { -// bar.set_message( -// HumanBytes(bytes_written.load(Ordering::SeqCst) as u64).to_string(), -// ); -// } -// Ok(()) -// }); -// if let Err(e) = err { -// eprintln!("Error occurred when processing tex: {e}"); -// eprintln!( -// "The process terminated early, we'll save the current processed tex files to pak file." -// ); -// } - -// match Arc::try_unwrap(pak_writer_mtx) { -// Ok(pak_writer) => pak_writer.into_inner().finish()?, -// Err(_) => panic!("Arc::try_unwrap failed"), -// }; - -// bar.finish(); -// println!("{}", "Done!".cyan().bold()); -// if !use_full_package_mode { -// println!( -// "You should rename the output file like `re_chunk_000.pak.sub_000.pak.patch_xxx.pak`." -// ); -// } - -// Ok(()) -// } - -fn is_tex_file(hash: u64, file_name_table: &FileNameTable) -> bool { - let Some(file_name) = file_name_table.get_file_name(hash) else { - return false; - }; - file_name.get_name().ends_with(".tex.241106027") -} - -fn write_to_pak( - writer: &mut ree_pak_core::write::PakWriter, - entry: &PakEntry, - file_name: impl FileNameExt, - data: &[u8], - use_feature_clone: bool, -) -> color_eyre::Result -where - W: io::Write + io::Seek, -{ - let mut file_options = FileOptions::default(); - if use_feature_clone { - file_options = file_options.with_unk_attr(*entry.unk_attr()) - } - writer.start_file(file_name, file_options)?; - writer.write_all(data)?; - Ok(data.len()) -} - fn wait_for_exit() { let _: String = Input::with_theme(&ColorfulTheme::default()) .with_prompt("Press Enter to exit") diff --git a/src/util.rs b/src/util.rs new file mode 100644 index 0000000..046de4e --- /dev/null +++ b/src/util.rs @@ -0,0 +1,5 @@ +use indicatif::HumanBytes; + +pub fn human_bytes(bytes: u64) -> String { + HumanBytes(bytes).to_string() +}