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