Display impls, LevelFile hashing on demand, read-bundle command

Author: yungcomputerchair

examples/extract_bundle.rs

@@ -19,8 +19,8 @@ async fn main() {
         std::fs::remove_dir_all(output_dir).ok();
 
         let time = std::time::Instant::now();
-        let asset = AssetBundle::from_file(&asset_path).unwrap();
-        asset.extract_files(output_dir).unwrap();
+        let (_, bundle) = AssetBundle::from_file(&asset_path).unwrap();
+        bundle.extract_files(output_dir).unwrap();
         info!("Extraction took {:?}", time.elapsed());
 
         let version = Version::from_manifest_file("example_manifest.json").unwrap();

src/bin/ffbuildtool.rs

@@ -24,6 +24,8 @@ enum Commands {
     RepairBuild(RepairBuildArgs),
     ValidateBuild(ValidateBuildArgs),
     #[cfg(feature = "lzma")]
+    ReadBundle(ReadBundleArgs),
+    #[cfg(feature = "lzma")]
     ExtractBundle(ExtractBundleArgs),
 }
 
@@ -95,6 +97,18 @@ struct ValidateBuildArgs {
     uncompressed: bool,
 }
 
+#[cfg(feature = "lzma")]
+#[derive(Args, Debug)]
+struct ReadBundleArgs {
+    /// Path to the compressed asset bundle
+    #[clap(short = 'b', long)]
+    bundle_path: String,
+
+    /// Whether to calculate the hashes of each file in the bundle
+    #[clap(short = 'c', long, action)]
+    calculate_hashes: bool,
+}
+
 #[cfg(feature = "lzma")]
 #[derive(Args, Debug)]
 struct ExtractBundleArgs {
@@ -219,6 +233,8 @@ async fn main() -> Result<(), String> {
         Commands::RepairBuild(args) => repair_build(args).await,
         Commands::ValidateBuild(args) => validate_build(args).await,
         #[cfg(feature = "lzma")]
+        Commands::ReadBundle(args) => read_bundle(args).await,
+        #[cfg(feature = "lzma")]
         Commands::ExtractBundle(args) => extract_bundle(args).await,
     }
 }
@@ -339,11 +355,35 @@ async fn validate_build(args: ValidateBuildArgs) -> Result<(), String> {
     Ok(())
 }
 
+#[cfg(feature = "lzma")]
+async fn read_bundle(args: ReadBundleArgs) -> Result<(), String> {
+    use std::time::Instant;
+
+    let start = Instant::now();
+    let (header, mut bundle) = ffbuildtool::bundle::AssetBundle::from_file(&args.bundle_path)?;
+    println!("Bundle read in {}ms", start.elapsed().as_millis());
+
+    if args.calculate_hashes {
+        let start = Instant::now();
+        bundle.recalculate_all_hashes();
+        println!("Hashes calculated in {}ms", start.elapsed().as_millis());
+    }
+
+    println!(
+        "------------------------\n{}\n------------------------\n{}",
+        header, bundle
+    );
+    Ok(())
+}
+
 #[cfg(feature = "lzma")]
 async fn extract_bundle(args: ExtractBundleArgs) -> Result<(), String> {
-    use std::path::PathBuf;
+    use std::{path::PathBuf, time::Instant};
+
+    let start = Instant::now();
+    let (_, bundle) = ffbuildtool::bundle::AssetBundle::from_file(&args.bundle_path)?;
+    println!("Bundle read in {}ms", start.elapsed().as_millis());
 
-    let asset_bundle = ffbuildtool::bundle::AssetBundle::from_file(&args.bundle_path)?;
     let output_dir = args.output_dir.unwrap_or({
         let bundle_name = ffbuildtool::util::get_file_name_without_parent(&args.bundle_path);
         let bundle_name_url_encoded = ffbuildtool::util::url_encode(bundle_name);
@@ -356,5 +396,10 @@ async fn extract_bundle(args: ExtractBundleArgs) -> Result<(), String> {
             .to_string()
     });
     println!("Extracting bundle {} to {}", args.bundle_path, output_dir);
-    asset_bundle.extract_files(&output_dir)
+
+    let start = Instant::now();
+    bundle.extract_files(&output_dir)?;
+    println!("Bundle extracted in {}ms", start.elapsed().as_millis());
+
+    Ok(())
 }

src/bundle.rs

@@ -2,7 +2,7 @@ use std::{
     collections::HashMap,
     fs::File,
     io::{BufRead, BufReader, BufWriter, Read, Write},
-    path::Path,
+    path::{Path, PathBuf},
 };
 
 use countio::Counter;
@@ -74,6 +74,27 @@ pub struct AssetBundleHeader {
     level_ends: Vec<LevelEnds>,
     bundle_size: u32,
 }
+impl std::fmt::Display for AssetBundleHeader {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        writeln!(f, "Signature: {}", self.signature)?;
+        writeln!(f, "Stream version: {}", self.stream_version)?;
+        writeln!(f, "Player version: {}", self.player_version)?;
+        writeln!(f, "Engine version: {}", self.engine_version)?;
+        writeln!(f, "Min streamed bytes: {}", self.min_streamed_bytes)?;
+        writeln!(f, "Header size: {}", self.header_size)?;
+        writeln!(f, "Number of levels: {}", self.num_levels)?;
+        writeln!(f, "Min levels for load: {}", self.min_levels_for_load)?;
+        writeln!(f, "Level ends:")?;
+        for (i, level) in self.level_ends.iter().enumerate() {
+            writeln!(
+                f,
+                "  Level {}: compressed @ {}, uncompressed @ {}",
+                i, level.compressed_end, level.uncompressed_end
+            )?;
+        }
+        write!(f, "Bundle size: {}", self.bundle_size)
+    }
+}
 impl AssetBundleHeader {
     fn new(level_ends: Vec<LevelEnds>) -> Self {
         let num_levels = level_ends.len() as u32;
@@ -253,15 +274,47 @@ impl LevelHeader {
 struct LevelFile {
     name: String,
     data: Vec<u8>,
+    hash: Option<String>,
 }
 impl std::fmt::Debug for LevelFile {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("LevelFile")
             .field("name", &self.name)
             .field("data", &format_args!("{} bytes", self.data.len()))
+            .field("hash", &self.hash)
             .finish()
     }
 }
+impl std::fmt::Display for LevelFile {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self.hash.as_ref() {
+            Some(hash) => write!(
+                f,
+                "{} - {} ({} bytes) - {}",
+                self.name,
+                util::bytes_to_human_readable(self.data.len() as u64),
+                self.data.len(),
+                hash
+            ),
+            None => write!(
+                f,
+                "{} - {} ({} bytes)",
+                self.name,
+                util::bytes_to_human_readable(self.data.len() as u64),
+                self.data.len()
+            ),
+        }
+    }
+}
+impl LevelFile {
+    fn new(name: String, data: Vec<u8>) -> Self {
+        Self {
+            name,
+            data,
+            hash: None,
+        }
+    }
+}
 
 #[derive(Debug)]
 struct Level {
@@ -278,10 +331,7 @@ impl Level {
             skip_exact(&mut reader, file.offset as usize - offset)?;
             let mut data = vec![0; file.size as usize];
             reader.read_exact(&mut data)?;
-            files.push(LevelFile {
-                name: file.name,
-                data,
-            });
+            files.push(LevelFile::new(file.name, data));
         }
         Ok(Self { files })
     }
@@ -345,8 +395,23 @@ impl Level {
 pub struct AssetBundle {
     levels: Vec<Level>,
 }
+impl std::fmt::Display for AssetBundle {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        for (i, level) in self.levels.iter().enumerate() {
+            writeln!(f, "Level {}", i)?;
+            for file in &level.files {
+                writeln!(f, "  {}", file)?;
+            }
+            write!(f, "End")?;
+        }
+        Ok(())
+    }
+}
 impl AssetBundle {
-    fn read<R: Read + BufRead>(reader: &mut R, expected_size: u32) -> Result<Self, Error> {
+    fn read<R: Read + BufRead>(
+        reader: &mut R,
+        expected_size: u32,
+    ) -> Result<(AssetBundleHeader, Self), Error> {
         let mut reader = Counter::new(reader);
 
         let header = AssetBundleHeader::read(&mut reader)?;
@@ -374,7 +439,7 @@ impl AssetBundle {
             }
         }
 
-        Ok(Self { levels })
+        Ok((header, Self { levels }))
     }
 
     fn write<W: Write>(&self, writer: &mut W, compression: u32) -> Result<(), Error> {
@@ -399,7 +464,27 @@ impl AssetBundle {
         Ok(())
     }
 
-    pub fn from_file(path: &str) -> Result<Self, String> {
+    fn get_level_files_from_dir(dir_path: &Path) -> Result<Vec<LevelFile>, Error> {
+        let mut files = Vec::new();
+        for entry in std::fs::read_dir(dir_path)? {
+            let Ok(entry) = entry else {
+                continue;
+            };
+            let path = entry.path();
+            if path.is_file() {
+                let Some(name) = path.file_name().unwrap().to_str() else {
+                    continue;
+                };
+                let Ok(data) = std::fs::read(&path) else {
+                    continue;
+                };
+                files.push(LevelFile::new(name.to_string(), data));
+            }
+        }
+        Ok(files)
+    }
+
+    pub fn from_file(path: &str) -> Result<(AssetBundleHeader, Self), String> {
         let file = File::open(path).map_err(|e| format!("Couldn't open file {}: {}", path, e))?;
         let metadata = file.metadata().unwrap();
         let mut reader = BufReader::new(file);
@@ -419,24 +504,6 @@ impl AssetBundle {
         Ok(())
     }
 
-    pub fn get_uncompressed_info(&self, level: usize) -> Result<HashMap<String, FileInfo>, Error> {
-        let mut result = HashMap::new();
-        if level >= self.levels.len() {
-            return Err(format!("Level {} does not exist", level).into());
-        }
-
-        for file in &self.levels[level].files {
-            let hash = util::get_buffer_hash(&file.data);
-            let info = FileInfo {
-                hash,
-                size: file.data.len() as u64,
-            };
-            result.insert(file.name.clone(), info);
-        }
-
-        Ok(result)
-    }
-
     pub fn extract_files(&self, output_dir: &str) -> Result<(), String> {
         let make_subdirs = self.levels.len() > 1;
         for (i, level) in self.levels.iter().enumerate() {
@@ -458,4 +525,32 @@ impl AssetBundle {
         }
         Ok(())
     }
+
+    pub fn recalculate_all_hashes(&mut self) {
+        for level in &mut self.levels {
+            for file in &mut level.files {
+                file.hash = Some(util::get_buffer_hash(&file.data));
+            }
+        }
+    }
+
+    pub fn get_uncompressed_info(&self, level: usize) -> Result<HashMap<String, FileInfo>, Error> {
+        let mut result = HashMap::new();
+        if level >= self.levels.len() {
+            return Err(format!("Level {} does not exist", level).into());
+        }
+
+        for file in &self.levels[level].files {
+            let info = FileInfo {
+                hash: file
+                    .hash
+                    .clone()
+                    .unwrap_or_else(|| util::get_buffer_hash(&file.data)),
+                size: file.data.len() as u64,
+            };
+            result.insert(file.name.clone(), info);
+        }
+
+        Ok(result)
+    }
 }

src/lib.rs

@@ -579,7 +579,7 @@ impl BundleInfo {
 
         #[cfg(feature = "lzma")]
         let uncompressed_info = {
-            let bundle = bundle::AssetBundle::from_file(&file_path)?;
+            let (_, bundle) = bundle::AssetBundle::from_file(&file_path)?;
             // ff assets are always in level 0
             bundle.get_uncompressed_info(0)?
         };

src/tests.rs

@@ -89,8 +89,8 @@ async fn test_extract_bundle() {
     let bundle_path = "example_builds/compressed/good/Map_00_00.unity3d";
     let output_dir = TempDir::new();
 
-    let asset_bundle = crate::bundle::AssetBundle::from_file(bundle_path).unwrap();
-    asset_bundle.extract_files(output_dir.path()).unwrap();
+    let (_, bundle) = crate::bundle::AssetBundle::from_file(bundle_path).unwrap();
+    bundle.extract_files(output_dir.path()).unwrap();
 
     let version = Version::from_manifest_file("example_manifest.json").unwrap();
     let bundle_info = version.get_bundle("Map_00_00.unity3d").unwrap();

src/util.rs

@@ -265,3 +265,14 @@ pub fn file_path_to_uri(file_path: &str) -> String {
     // Add file:/// protocol
     format!("file:///{}", path)
 }
+
+pub fn bytes_to_human_readable(bytes: u64) -> String {
+    let units = ["B", "KB", "MB", "GB"];
+    let mut bytes = bytes as f64;
+    let mut unit = 0;
+    while bytes >= 1024.0 && unit < units.len() - 1 {
+        bytes /= 1024.0;
+        unit += 1;
+    }
+    format!("{:.2} {}", bytes, units[unit])
+}