patching to unity .assets file

Author: ByteZ1337

src/command/patch/assets_patcher.rs

@@ -1,10 +1,35 @@
 use std::ffi::OsStr;
+use std::fs::File;
+use std::io::{BufWriter, Read, Seek, SeekFrom, Write};
 use std::path::{Path, PathBuf};
+
 use anyhow::Context;
+use binrw::__private::write_zeroes;
+use binrw::BinWrite;
+use binrw::io::BufReader;
+use byteorder::WriteBytesExt;
 use walkdir::WalkDir;
+
+use crate::command::pack;
 use crate::command::patch::xml_patcher;
+use crate::command::unpack::RepackInfo;
+use crate::io_ext::WriteExt;
+use crate::unity::{AssetsFile, AssetsFileContent, AssetsFileHeader, ObjectInfo};
+
+/// Length of the header of the Art.dat object.
+/// The header consists of:
+/// - 4 bytes for object name length
+/// - 7 bytes for object name (Art.dat)
+/// - 1 byte for field index
+/// - 4 bytes for data length
+const ART_OBJ_HEADER_LEN: u64 = 4 + 7 + 1 + 4;
 
-pub fn patch_assets(patch: &PathBuf, unpacked: &PathBuf, temp_dir: &PathBuf) -> anyhow::Result<PathBuf> {
+pub fn patch_assets(
+    patch: &PathBuf,
+    unpacked: &PathBuf,
+    temp_dir: &PathBuf,
+    repack_info: RepackInfo,
+) -> anyhow::Result<PathBuf> {
     println!("Patching assets..");
     let patched_assets = temp_dir.join("patched");
     std::fs::create_dir_all(&patched_assets)
@@ -46,11 +71,10 @@ pub fn patch_assets(patch: &PathBuf, unpacked: &PathBuf, temp_dir: &PathBuf) ->
             copy_file(&patch_file.as_path(), rel_path, &patched_assets)?;
         } else if ext == OsStr::new("xml") || ext == OsStr::new("fnt") {
             println!("Patching xml file: {}", rel_path.display());
-            patch_xml(&file, &patch_file, rel_path, &patched_assets)?;
+            patch_xml(&file.path(), &patch_file, rel_path, &patched_assets)?;
         } else {
             anyhow::bail!("Unsupported file type: {}", patch_file.display());
         }
-
     }
 
     // Loop over any files newly added with the patch
@@ -79,7 +103,7 @@ pub fn patch_assets(patch: &PathBuf, unpacked: &PathBuf, temp_dir: &PathBuf) ->
         }
     }
 
-    return Ok(patched_assets);
+    pack_to_assets(temp_dir, &patched_assets, repack_info)
 }
 
 /// Copies a file from one of the input directories to the patched assets directory and makes sure
@@ -94,11 +118,99 @@ fn copy_file(file: &Path, rel_path: &Path, patched_assets: &PathBuf) -> anyhow::
 }
 
 /// Patches an XML file using the given patch file and writes the output to the patched assets directory
-fn patch_xml(original: &walkdir::DirEntry, patch_file: &PathBuf, rel_path: &Path, patched_assets: &PathBuf) -> anyhow::Result<()> {
+fn patch_xml(original: &Path, patch_file: &PathBuf, rel_path: &Path, patched_assets: &PathBuf) -> anyhow::Result<()> {
     let output = patched_assets.join(rel_path);
     if let Some(parent) = output.parent() {
         std::fs::create_dir_all(parent)
             .context("Failed to create directory")?;
     }
-    xml_patcher::patch(original.path(), patch_file, &output)
+    xml_patcher::patch(original, patch_file, &output)
+}
+
+fn pack_to_assets(temp_dir: &PathBuf, patched: &PathBuf, repack: RepackInfo) -> anyhow::Result<PathBuf> {
+    let output = temp_dir.join("repacked.assets");
+    let temp_art = temp_dir.join("patched-art.dat");
+    pack::pack(&repack.art_key, &Some(patched.clone()), &temp_art)?;
+    let assets = repack.assets;
+    let new_art_len = std::fs::metadata(&temp_art)
+        .context("Failed to get metadata of temp art file")?
+        .len();
+
+    // header
+    let mut header = AssetsFileHeader { file_size: 0, ..assets.header };
+
+    // content
+    let mut objects = Vec::new();
+    let mut current_offset = 0;
+    for obj in &assets.content.objects {
+        if current_offset % 4 != 0 {
+            current_offset += 4 - (current_offset % 4);
+        }
+
+        let mut new_object = ObjectInfo {
+            path_id: obj.path_id,
+            byte_start: current_offset,
+            byte_size: 0,
+            type_id: obj.type_id,
+        };
+        if obj.path_id == repack.art_path_id {
+            new_object.byte_size = (new_art_len + ART_OBJ_HEADER_LEN) as u32;
+        } else {
+            new_object.byte_size = obj.byte_size;
+        }
+        current_offset += new_object.byte_size as u64;
+        objects.push(new_object);
+    }
+    header.file_size = header.offset_first_file + current_offset;
+    let content = AssetsFileContent { objects, ..assets.content };
+    let new_assets = AssetsFile { header, content };
+    let mut writer = BufWriter::new(File::create(&output)
+        .context("Failed to create output file")?);
+    new_assets.write(&mut writer)
+        .context("Failed to write assets file header")?;
+
+    // pad with zeroes until first file offset is reached (yes this is also what Unity does)
+    let pad = assets.header.offset_first_file - writer.seek(SeekFrom::Current(0))
+        .context("Failed to get current position in output file")?;
+    write_zeroes(&mut writer, pad)?;
+
+    // write the actual object data
+    let mut original = BufReader::new(File::open(&repack.original_assets)
+        .context("Failed to open original assets file")?);
+    let original_file_offset = &assets.header.offset_first_file;
+    for (obj, old_obj) in new_assets.content.objects.iter().zip(assets.content.objects) {
+
+        let pos = writer.stream_position()
+            .context("Failed to get current position in output file")?;
+        if pos != obj.byte_start + original_file_offset {
+            // pad with zeroes until the object's start offset is reached
+            let pad = obj.byte_start + original_file_offset - pos;
+            write_zeroes(&mut writer, pad).context("Failed to write padding zeroes")?;
+        }
+
+        if obj.path_id != repack.art_path_id {
+            original.seek(SeekFrom::Start(original_file_offset + old_obj.byte_start))
+                .context("Failed to seek to object in original assets file")?;
+            let mut data = vec![0; obj.byte_size as usize];
+            original.read_exact(&mut data)
+                .context("Failed to read object data from original assets file")?;
+            writer.write_all(&data)?;
+        } else {
+            writer.write_dyn_string("Art.dat", &new_assets.header.endianness)
+                .context("Failed to write object name")?;
+            writer.write_u8(0)
+                .context("Failed to write field index")?;
+            writer.write_u32_order(&new_assets.header.endianness, new_art_len as u32)
+                .context("Failed to write object data length")?;
+            // copy over the new art file
+            let mut art_file = BufReader::new(File::open(&temp_art)
+                .context("Failed to open temp art file")?);
+            std::io::copy(&mut art_file, &mut writer)
+                .context("Failed to copy new art file to assets file")?;
+        }
+
+    }
+
+    println!("Packed objets to: {}", output.display());
+    Ok(output)
 }

src/command/patch/mod.rs

@@ -1,16 +1,16 @@
-pub mod assets_patcher;
-pub mod xml_patcher;
-
 use std::env::temp_dir;
 use std::path::PathBuf;
-use anyhow::Context;
 
+use anyhow::Context;
 use rand::random;
 
 use crate::{I18nCompatMode, NewArgs};
 use crate::command::patch::assets_patcher::patch_assets;
 use crate::command::unpack;
 
+pub mod assets_patcher;
+pub mod xml_patcher;
+
 pub fn patch(args: &NewArgs, patch: &PathBuf, locale_mode: &I18nCompatMode) -> anyhow::Result<()> {
     println!("Patching assets with {:?} with locale mode {:?}", patch, locale_mode);
 
@@ -26,9 +26,9 @@ pub fn patch(args: &NewArgs, patch: &PathBuf, locale_mode: &I18nCompatMode) -> a
     std::fs::create_dir_all(&temp_unpacked)
         .context("Failed to create temp directory")?;
 
-    unpack::unpack_assets(args, &game_files.assets, &temp_unpacked)?;
+    let repack_info = unpack::unpack_assets(args, &game_files.assets, &temp_unpacked)?;
 
-    patch_assets(patch, &temp_unpacked, &temp_dir)?;
+    patch_assets(patch, &temp_unpacked, &temp_dir, repack_info)?;
 
     Ok(())
 }
@@ -82,6 +82,7 @@ fn prepare_file(game_dir: &PathBuf, name: &str) -> anyhow::Result<PathBuf> {
 
 fn create_temp_dir() -> PathBuf {
     let mut temp_dir = temp_dir();
+    temp_dir.push("papers-tools");
     temp_dir.push(format!("papers_please_assets_{}", random::<u64>()));
     temp_dir
 }

src/unity/mod.rs

@@ -94,7 +94,7 @@ pub struct SerializedType {
 #[derive(Debug, PartialEq)]
 pub struct ScriptType {
     local_serialized_file_index: i32,
-    #[br(align_before(4))]
+    #[brw(align_before(4))]
     local_identifier_in_file: i64,
 }
 
@@ -110,7 +110,7 @@ pub struct FileIdentifier {
 #[binrw]
 #[derive(Debug, PartialEq)]
 pub struct ObjectInfo {
-    #[br(align_before(4))]
+    #[brw(align_before(4))]
     pub path_id: i64,
     pub byte_start: u64,
     pub byte_size: u32,