From 00fb900ddb2f9947c72880120745f90a1ae8258b Mon Sep 17 00:00:00 2001
From: devttys0 <heffnercj@gmail.com>
Date: Sun, 17 Nov 2024 20:59:25 -0500
Subject: [PATCH] Added support for compressed CSMAN files

---
 Cargo.lock              |  1 +
 Cargo.toml              |  1 +
 src/extractors/csman.rs | 22 +++++++++++++++++++++-
 src/structures/csman.rs | 28 +++++++++++++++++++++-------
 4 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 74b19d93c..af41cda0a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -115,6 +115,7 @@ dependencies = [
  "env_logger",
  "flate2",
  "log",
+ "miniz_oxide 0.8.0",
  "plotters",
  "serde",
  "serde_json",
diff --git a/Cargo.toml b/Cargo.toml
index a53076486..df8b3a990 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -26,6 +26,7 @@ threadpool = "1.8.1"
 serde_json = "1.0"
 env_logger = "0.11.5"
 flate2 = "1.0.34"
+miniz_oxide = "0.8.0"
 aho-corasick = "1.1.3"
 serde = { version = "1.0", features = ["derive"]}
 clap = { version = "4.5.16", features = ["derive"] }
diff --git a/src/extractors/csman.rs b/src/extractors/csman.rs
index f7b62cf26..c07e3b021 100644
--- a/src/extractors/csman.rs
+++ b/src/extractors/csman.rs
@@ -1,6 +1,7 @@
 use crate::common::is_offset_safe;
 use crate::extractors::common::{Chroot, ExtractionResult, Extractor, ExtractorType};
 use crate::structures::csman::{parse_csman_entry, parse_csman_header, CSManEntry};
+use miniz_oxide::inflate;
 use std::collections::HashMap;
 
 /// Defines the internal extractor function for CSMan DAT files
@@ -38,6 +39,8 @@ pub fn extract_csman_dat(
     offset: usize,
     output_directory: Option<&String>,
 ) -> ExtractionResult {
+    const COMPRESSED_HEADER_SIZE: usize = 2;
+
     // Return value
     let mut result = ExtractionResult {
         ..Default::default()
@@ -47,12 +50,29 @@ pub fn extract_csman_dat(
 
     // Parse the CSMAN header
     if let Ok(csman_header) = parse_csman_header(&file_data[offset..]) {
+        println!("{:?}", csman_header);
         // Calulate the start and end offsets of the CSMAN entries
         let entries_start: usize = offset + csman_header.header_size;
         let entries_end: usize = entries_start + csman_header.data_size;
 
         // Get the CSMAN entry data
-        if let Some(entry_data) = file_data.get(entries_start..entries_end) {
+        if let Some(raw_entry_data) = file_data.get(entries_start..entries_end) {
+            let mut entry_data = raw_entry_data.to_vec();
+
+            // If the entries are compressed, decompress it (zlib compression)
+            if csman_header.compressed {
+                if let Some(compressed_data) = raw_entry_data.get(COMPRESSED_HEADER_SIZE..) {
+                    match inflate::decompress_to_vec(compressed_data) {
+                        Err(_) => {
+                            return result;
+                        }
+                        Ok(decompressed_data) => {
+                            entry_data = decompressed_data.clone();
+                        }
+                    }
+                }
+            }
+
             // Offsets for processing CSMAN entries in entry_data
             let mut next_offset: usize = 0;
             let mut previous_offset = None;
diff --git a/src/structures/csman.rs b/src/structures/csman.rs
index 5a2ae384a..72e23ef9d 100644
--- a/src/structures/csman.rs
+++ b/src/structures/csman.rs
@@ -3,6 +3,7 @@ use crate::structures::common::{self, StructureError};
 /// Struct to store CSMAN header info
 #[derive(Debug, Default, Clone)]
 pub struct CSManHeader {
+    pub compressed: bool,
     pub data_size: usize,
     pub endianness: String,
     pub header_size: usize,
@@ -10,14 +11,15 @@ pub struct CSManHeader {
 
 /// Parses a CSMAN header
 pub fn parse_csman_header(csman_data: &[u8]) -> Result<CSManHeader, StructureError> {
+    const COMPRESSED_MAGIC: &[u8] = b"\x78";
     const LITTLE_ENDIAN_MAGIC: usize = 0x4353;
 
     let csman_header_structure = vec![
         ("magic", "u16"),
         ("unknown1", "u16"),
-        ("data_size_1", "u32"),
+        ("compressed_size", "u32"),
         ("unknown2", "u32"),
-        ("data_size_2", "u32"),
+        ("decompressed_size", "u32"),
     ];
 
     let mut result = CSManHeader {
@@ -39,12 +41,24 @@ pub fn parse_csman_header(csman_data: &[u8]) -> Result<CSManHeader, StructureErr
             result.endianness = "big".to_string();
         }
 
-        // Data size is repeated in both these fields
-        if !result.endianness.is_empty()
-            && csman_header["data_size_1"] == csman_header["data_size_2"]
-        {
-            result.data_size = csman_header["data_size_1"];
+        // Should have been able to determine the endianness
+        if !result.endianness.is_empty() {
+            result.data_size = csman_header["compressed_size"];
             result.header_size = common::size(&csman_header_structure);
+            result.compressed =
+                csman_header["compressed_size"] != csman_header["decompressed_size"];
+
+            // If compressed, check the expected compressed magic bytes
+            if result.compressed {
+                if let Some(compressed_magic) =
+                    csman_data.get(result.header_size..result.header_size + COMPRESSED_MAGIC.len())
+                {
+                    if compressed_magic != COMPRESSED_MAGIC {
+                        return Err(StructureError);
+                    }
+                }
+            }
+
             return Ok(result);
         }
     }