Skip to main content

rkg_utils/
lib.rs

1//! This crate is meant to be a library to completely and coherently access the data in RKGD files,
2//! Mario Kart Wii's native Ghost Data.
3//!
4//! Features:
5//! - [x] Reading and Writing Vanilla Game Data (including embedded Mii data)
6//! - [x] Reading and Writing CTGP Modified Data
7//! - [x] Reading and Writing Pulsar (Retro Rewind) Modified Data
8//! - [x] Reading and Writing MKW-SP Modified Data
9//! - [ ] Implementing `TryFrom<_>` for T where T: `Into<ByteHandler>`, relies on <https://github.com/rust-lang/rust/issues/31844> currently
10//! - [ ] Represent at a Type-system level which types can convert from `T1` (Bytes) to `crate::byte_handler::ByteHandler` to `T2` (Typed Structs)
11//! - [ ] Optimize Little-Endian calculations with `crate::byte_handler::ByteHandler`
12//! - [ ] Figure out whether Big-Endian works with `crate::byte_handler::ByteHandler`
13
14use std::{
15    array::TryFromSliceError,
16    io::{Read, Write},
17};
18
19use chrono::{DateTime, Duration, NaiveDateTime, TimeDelta};
20use sha1::{Digest, Sha1};
21
22use crate::{
23    crc::crc32,
24    ctgp_footer::CTGPFooter,
25    header::{Header, mii::Mii},
26    input_data::InputData,
27    sp_footer::SPFooter,
28};
29
30pub mod byte_handler;
31pub mod crc;
32pub mod ctgp_footer;
33pub mod header;
34pub mod input_data;
35pub mod sp_footer;
36
37#[cfg(test)]
38mod tests;
39
40/// Errors that can occur while parsing or modifying a [`Ghost`].
41#[derive(thiserror::Error, Debug)]
42pub enum GhostError {
43    /// The input data is shorter than the minimum valid ghost size (`0x8E` bytes).
44    #[error("Data length too short for a ghost")]
45    DataLengthTooShort,
46    /// The RKG file header could not be parsed.
47    #[error("Header Error: {0}")]
48    HeaderError(#[from] header::HeaderError),
49    /// The embedded Mii data could not be parsed.
50    #[error("Mii Error: {0}")]
51    MiiError(#[from] header::mii::MiiError),
52    /// The ghost input data could not be parsed.
53    #[error("Input Data Error: {0}")]
54    InputDataError(#[from] input_data::InputDataError),
55    /// The CTGP footer could not be parsed.
56    #[error("CTGP Footer Error: {0}")]
57    CTGPFooterError(#[from] ctgp_footer::CTGPFooterError),
58    /// A `ByteHandler`(byte_handler::ByteHandler) operation failed.
59    #[error("ByteHandler Error: {0}")]
60    ByteHandlerError(#[from] byte_handler::ByteHandlerError),
61    /// A slice-to-array conversion failed (e.g. when extracting a CRC-32 word).
62    #[error("Try From Slice Error: {0}")]
63    TryFromSliceError(#[from] TryFromSliceError),
64    /// A file I/O operation failed.
65    #[error("IO Error: {0}")]
66    IOError(#[from] std::io::Error),
67}
68
69/// A fully parsed Mario Kart Wii RKG ghost file.
70///
71/// Holds the file header, decompressed or compressed input data, optional
72/// external footers (CTGP or SP), and CRC-32 checksums. All setter
73/// operations update the parsed fields; call [`update_raw_data`](Ghost::update_raw_data)
74/// (or [`save_to_file`](Ghost::save_to_file), which calls it implicitly) to
75/// flush all changes back into the raw byte buffer before writing.
76pub struct Ghost {
77    /// The complete raw file bytes, kept in sync by [`update_raw_data`](Ghost::update_raw_data).
78    raw_data: Vec<u8>,
79    /// The parsed 136-byte RKG file header.
80    header: Header,
81    /// The ghost's controller input data (compressed or decompressed).
82    input_data: InputData,
83    /// The CRC-32 of the header and input data only, excluding any external footer.
84    base_crc32: u32,
85    /// The CTGP footer appended to the file, if present.
86    ctgp_footer: Option<CTGPFooter>,
87    /// The SP (MKW Service Pack) footer appended to the file, if present.
88    sp_footer: Option<SPFooter>,
89    /// The CRC-32 of the entire file excluding its final 4 bytes.
90    file_crc32: u32,
91    /// When `true`, any existing external footer is preserved when saving (including footer data from any mods that this crate doesn't support).
92    should_preserve_external_footer: bool,
93}
94
95impl Ghost {
96    /// Parses a [`Ghost`] from an RKG file at the given path.
97    ///
98    /// # Errors
99    ///
100    /// Returns [`GhostError::IOError`] if the file cannot be opened or read,
101    /// and other [`GhostError`] variants if parsing fails.
102    pub fn new_from_file<T: AsRef<std::path::Path>>(path: T) -> Result<Self, GhostError> {
103        let mut buf = Vec::with_capacity(0x100);
104        std::fs::File::open(path)?.read_to_end(&mut buf)?;
105        Self::new(&buf)
106    }
107
108    /// Parses a [`Ghost`] from a byte slice.
109    ///
110    /// Detects and parses an optional CTGP or SP footer if present. The
111    /// base CRC-32 is read from just before the footer when one is found,
112    /// or from the last 4 bytes of the file otherwise.
113    ///
114    /// # Errors
115    ///
116    /// Returns [`GhostError::DataLengthTooShort`] if `bytes` is shorter than
117    /// `0x8E` bytes. Returns other [`GhostError`] variants if any field fails
118    /// to parse.
119    pub fn new(bytes: &[u8]) -> Result<Self, GhostError> {
120        if bytes.len() < 0x8E {
121            return Err(GhostError::DataLengthTooShort);
122        }
123
124        let header = Header::new(&bytes[..0x88])?;
125
126        let file_crc32 = u32::from_be_bytes(bytes[bytes.len() - 0x04..].try_into()?);
127        let mut base_crc32 = file_crc32;
128
129        let ctgp_footer = if let Ok(ctgp_footer) = CTGPFooter::new(bytes) {
130            let input_data_end_offset = bytes.len() - ctgp_footer.len() - 0x08;
131            base_crc32 = u32::from_be_bytes(
132                bytes[input_data_end_offset..input_data_end_offset + 0x04].try_into()?,
133            );
134            Some(ctgp_footer)
135        } else {
136            None
137        };
138
139        let sp_footer = if let Ok(sp_footer) = SPFooter::new(bytes) {
140            let input_data_end_offset = bytes.len() - sp_footer.len() - 0x08;
141            base_crc32 = u32::from_be_bytes(
142                bytes[input_data_end_offset..input_data_end_offset + 0x04].try_into()?,
143            );
144            Some(sp_footer)
145        } else {
146            None
147        };
148
149        let input_data_len = if bytes[0x8C..0x90] == *b"Yaz1" {
150            u32::from_be_bytes(bytes[0x88..0x8C].try_into().unwrap()) as usize + 0x04
151        } else {
152            header.decompressed_input_data_length() as usize
153        };
154
155        let input_data = InputData::new(&bytes[0x88..0x88 + input_data_len])?;
156
157        Ok(Self {
158            raw_data: bytes.to_vec(),
159            header,
160            input_data,
161            base_crc32,
162            ctgp_footer,
163            sp_footer,
164            file_crc32,
165            should_preserve_external_footer: true,
166        })
167    }
168
169    /// Flushes all parsed field modifications back into the raw byte buffer.
170    ///
171    /// This method recomputes the Mii CRC-16, rebuilds the raw buffer from the
172    /// current header and input data, resizes the buffer if the input data
173    /// length has changed, re-inserts any preserved external footer, and
174    /// finally recomputes both the base CRC-32 and the file-level CRC-32. The
175    /// CTGP footer SHA-1 field is also updated if a CTGP footer is present.
176    ///
177    /// # Errors
178    ///
179    /// Returns [`GhostError::MiiError`] if the Mii data is invalid, or
180    /// [`GhostError::CTGPFooterError`] if the SHA-1 field cannot be written.
181    pub fn update_raw_data(&mut self) -> Result<(), GhostError> {
182        let mii_bytes = self.header().mii().raw_data().to_vec();
183        self.header_mut().set_mii(Mii::new(mii_bytes)?);
184        self.header_mut().fix_mii_crc16();
185
186        let mut buf = Vec::from(self.header().raw_data());
187
188        buf.extend_from_slice(self.input_data().raw_data());
189
190        let header_len = 0x88;
191        let new_input_data_end = header_len + self.input_data().raw_data().len();
192
193        // Find input data length of old data
194        let old_input_data_end = if self.raw_data[0x8C..0x90] == *b"Yaz1" {
195            u32::from_be_bytes(self.raw_data[0x88..0x8C].try_into().unwrap()) as usize
196        } else {
197            header_len + 0x2774
198        };
199
200        if new_input_data_end > old_input_data_end {
201            let diff = new_input_data_end - old_input_data_end;
202            let insert_pos = old_input_data_end;
203            self.raw_data
204                .splice(insert_pos..insert_pos, vec![0u8; diff]);
205        } else if new_input_data_end < old_input_data_end {
206            let diff = old_input_data_end - new_input_data_end;
207            let remove_end = old_input_data_end;
208            self.raw_data.drain(remove_end - diff..remove_end);
209        }
210
211        self.raw_data[..new_input_data_end].copy_from_slice(&buf[..new_input_data_end]);
212        let base_crc32 = crc32(&buf);
213
214        if let Some(ctgp_footer) = self.ctgp_footer()
215            && self.should_preserve_external_footer()
216        {
217            buf.extend_from_slice(&base_crc32.to_be_bytes());
218            buf.extend_from_slice(ctgp_footer.raw_data());
219
220            let footer_len = ctgp_footer.len();
221            self.raw_data.drain(new_input_data_end..);
222            self.raw_data
223                .extend_from_slice(&buf[buf.len() - footer_len - 0x04..]);
224            self.raw_data.extend_from_slice(&[0u8; 4]);
225        } else if let Some(sp_footer) = self.sp_footer()
226            && self.should_preserve_external_footer()
227        {
228            buf.extend_from_slice(&base_crc32.to_be_bytes());
229            buf.extend_from_slice(sp_footer.raw_data());
230
231            let footer_len = sp_footer.len();
232            self.raw_data.drain(new_input_data_end..);
233            self.raw_data
234                .extend_from_slice(&buf[buf.len() - footer_len - 0x04..]);
235            self.raw_data.extend_from_slice(&[0u8; 4]);
236        } else if !self.should_preserve_external_footer()
237            && self.raw_data.len() >= new_input_data_end + 0x08
238        {
239            self.raw_data.drain(new_input_data_end + 0x04..);
240        } else if self.should_preserve_external_footer()
241            && self.raw_data.len() >= new_input_data_end + 0x08
242        {
243            self.raw_data[new_input_data_end..new_input_data_end + 0x04]
244                .copy_from_slice(&base_crc32.to_be_bytes());
245        }
246
247        let len = self.raw_data.len();
248        let crc32 = crc32(&self.raw_data[..len - 0x04]);
249        self.raw_data[len - 0x04..].copy_from_slice(&crc32.to_be_bytes());
250
251        let sha1 = compute_sha1_hex(&self.raw_data);
252        if let Some(ctgp_footer) = self.ctgp_footer_mut() {
253            ctgp_footer.set_ghost_sha1(&sha1)?;
254        }
255
256        Ok(())
257    }
258
259    /// Flushes all modifications and writes the ghost to a file at the given path.
260    ///
261    /// # Errors
262    ///
263    /// Returns any error from [`update_raw_data`](Ghost::update_raw_data) or
264    /// from file creation/writing.
265    pub fn save_to_file<T: AsRef<std::path::Path>>(&mut self, path: T) -> Result<(), GhostError> {
266        self.update_raw_data()?;
267        let mut file = std::fs::File::create(path)?;
268        file.write_all(&self.raw_data)?;
269
270        Ok(())
271    }
272
273    /// Compresses the input data using Yaz1 encoding and sets the compression flag in the header.
274    ///
275    /// Does nothing if the input data is already compressed.
276    pub fn compress_input_data(&mut self) {
277        if self.input_data().is_compressed() {
278            return;
279        }
280
281        self.input_data_mut().compress();
282        self.header_mut().set_compressed(true);
283    }
284
285    /// Decompresses the input data and clears the compression flag in the header.
286    ///
287    /// Does nothing if the input data is not compressed.
288    pub fn decompress_input_data(&mut self) {
289        if !self.input_data().is_compressed() {
290            return;
291        }
292
293        self.input_data_mut().decompress();
294        self.header_mut().set_compressed(false);
295    }
296
297    /// Returns the raw file bytes.
298    ///
299    /// May not reflect recent modifications until [`update_raw_data`](Ghost::update_raw_data) is called.
300    pub fn raw_data(&self) -> &[u8] {
301        &self.raw_data
302    }
303
304    /// Returns a mutable reference to the raw file bytes.
305    pub fn raw_data_mut(&mut self) -> &mut [u8] {
306        &mut self.raw_data
307    }
308
309    /// Returns the parsed RKG file header.
310    pub fn header(&self) -> &Header {
311        &self.header
312    }
313
314    /// Returns a mutable reference to the parsed RKG file header.
315    pub fn header_mut(&mut self) -> &mut Header {
316        &mut self.header
317    }
318
319    /// Returns the ghost's controller input data.
320    pub fn input_data(&self) -> &InputData {
321        &self.input_data
322    }
323
324    /// Returns a mutable reference to the ghost's controller input data.
325    pub fn input_data_mut(&mut self) -> &mut InputData {
326        &mut self.input_data
327    }
328
329    /// Returns the CTGP footer, if present.
330    pub fn ctgp_footer(&self) -> Option<&CTGPFooter> {
331        self.ctgp_footer.as_ref()
332    }
333
334    /// Returns a mutable reference to the CTGP footer, if present.
335    pub fn ctgp_footer_mut(&mut self) -> Option<&mut CTGPFooter> {
336        self.ctgp_footer.as_mut()
337    }
338
339    /// Returns the SP footer, if present.
340    pub fn sp_footer(&self) -> Option<&SPFooter> {
341        self.sp_footer.as_ref()
342    }
343
344    /// Returns the CRC-32 of the header and input data, excluding any external footer.
345    pub fn base_crc32(&self) -> u32 {
346        self.base_crc32
347    }
348
349    /// Returns `true` if the stored base CRC-32 matches a freshly computed
350    /// checksum of the current header and input data bytes.
351    pub fn verify_base_crc32(&self) -> bool {
352        let mut data = Vec::from(self.header().raw_data());
353        data.extend_from_slice(self.input_data().raw_data());
354        self.base_crc32 == crc32(&data)
355    }
356
357    /// Returns the CRC-32 of the entire file excluding its final 4 bytes.
358    pub fn file_crc32(&self) -> u32 {
359        self.file_crc32
360    }
361
362    /// Returns `true` if the stored file CRC-32 matches a freshly computed
363    /// checksum of the entire file excluding its final 4 bytes.
364    pub fn verify_file_crc32(&self) -> bool {
365        let len = self.raw_data().len();
366        self.file_crc32 == crc32(&self.raw_data()[..len - 0x04])
367    }
368
369    /// Returns whether an existing external footer will be preserved when saving.
370    pub fn should_preserve_external_footer(&self) -> bool {
371        self.should_preserve_external_footer
372    }
373
374    /// Sets whether an existing external footer should be preserved when saving.
375    pub fn set_should_preserve_external_footer(&mut self, b: bool) {
376        self.should_preserve_external_footer = b;
377    }
378}
379
380/// Used internally for writing bits to a buffer.
381pub(crate) fn write_bits(
382    buf: &mut [u8],
383    byte_offset: usize,
384    bit_offset: usize,
385    bit_width: usize,
386    value: u64,
387) {
388    let bytes_needed = (bit_offset + bit_width).div_ceil(8);
389    let mut chunk: u64 = 0;
390
391    for i in 0..bytes_needed {
392        chunk = (chunk << 8) | buf[byte_offset + i] as u64;
393    }
394
395    let shift = bytes_needed * 8 - bit_offset - bit_width;
396    let mask = ((1u64 << bit_width) - 1) << shift;
397
398    chunk = (chunk & !mask) | ((value << shift) & mask);
399
400    for i in (0..bytes_needed).rev() {
401        buf[byte_offset + i] = (chunk & 0xFF) as u8;
402        chunk >>= 8;
403    }
404}
405
406/// Computes the SHA-1 hash of `input` and returns it as a 20-byte array.
407pub(crate) fn compute_sha1_hex(input: &[u8]) -> [u8; 0x14] {
408    let mut hasher = Sha1::new();
409    hasher.update(input);
410    hasher.finalize().into()
411}
412
413/// Converts a raw Wii tick count into a [`NaiveDateTime`].
414///
415/// Ticks run at 60.75 MHz relative to the Wii epoch of 2000-01-01 00:00:00 UTC.
416pub(crate) fn datetime_from_timestamp(tick_count: u64) -> NaiveDateTime {
417    let clock_rate = 60_750_000.0; // 60.75 MHz tick speed
418    let epoch_shift = 946_684_800; // Shifts epoch from 1970-01-01 to 2000-01-01 (which is what the Wii uses)
419    let total_seconds = tick_count as f64 / clock_rate;
420    let total_nanoseconds = (total_seconds * 1_000_000_000.0) as i64;
421
422    let duration = Duration::nanoseconds(total_nanoseconds);
423    let epoch = DateTime::from_timestamp(epoch_shift, 0).unwrap();
424
425    epoch.naive_utc() + duration
426}
427
428/// Converts a raw Wii tick count into a [`TimeDelta`] (duration from an arbitrary reference).
429///
430/// Ticks run at 60.75 MHz; the result is truncated to millisecond precision.
431pub(crate) fn duration_from_ticks(tick_count: u64) -> TimeDelta {
432    let clock_rate = 60_750_000.0; // 60.75 MHz tick speed
433    let total_seconds = tick_count as f64 / clock_rate;
434    let total_milliseconds = (total_seconds * 1_000.0) as i64;
435
436    Duration::milliseconds(total_milliseconds)
437}