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    /// Error with data length being too short.
42    #[error("Data length too short")]
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    /// `0x90` bytes OR input data is shorter than the expected input data size.
114    /// Returns other [`GhostError`] variants if any field fails
115    /// to parse.
116    pub fn new(bytes: &[u8]) -> Result<Self, GhostError> {
117        if bytes.len() < 0x90 {
118            return Err(GhostError::DataLengthTooShort);
119        }
120
121        let header = Header::new(&bytes[..0x88])?;
122
123        let file_crc32 = u32::from_be_bytes(bytes[bytes.len() - 0x04..].try_into()?);
124        let mut base_crc32 = file_crc32;
125
126        let footer = if let Ok(ctgp_footer) = CTGPFooter::new(bytes) {
127            let input_data_end_offset = bytes.len() - ctgp_footer.len() - 0x08;
128            base_crc32 = u32::from_be_bytes(
129                bytes[input_data_end_offset..input_data_end_offset + 0x04].try_into()?,
130            );
131            Some(FooterType::CTGPFooter(ctgp_footer))
132        } else if let Ok(sp_footer) = SPFooter::new(bytes) {
133            let input_data_end_offset = bytes.len() - sp_footer.len() - 0x08;
134            base_crc32 = u32::from_be_bytes(
135                bytes[input_data_end_offset..input_data_end_offset + 0x04].try_into()?,
136            );
137            Some(FooterType::SPFooter(sp_footer))
138        } else {
139            None
140        };
141
142        let input_data_len = if bytes[0x8C..0x90] == *b"Yaz1" {
143            u32::from_be_bytes(bytes[0x88..0x8C].try_into().unwrap()) as usize + 0x04
144        } else {
145            header.decompressed_input_data_length() as usize
146        };
147
148        if bytes.len() < 0x88 + input_data_len + 0x04 {
149            return Err(GhostError::DataLengthTooShort);
150        }
151
152        let input_data = InputData::new(&bytes[0x88..0x88 + input_data_len])?;
153
154        Ok(Self {
155            raw_data: bytes.to_vec(),
156            header,
157            input_data,
158            base_crc32,
159            footer,
160            file_crc32,
161            should_preserve_external_footer: true,
162        })
163    }
164
165    /// Flushes all parsed field modifications back into the raw byte buffer.
166    ///
167    /// This method recomputes the Mii CRC-16, rebuilds the raw buffer from the
168    /// current header and input data, resizes the buffer if the input data
169    /// length has changed, re-inserts any preserved external footer, and
170    /// finally recomputes both the base CRC-32 and the file-level CRC-32. The
171    /// CTGP footer SHA-1 field is also updated if a CTGP footer is present.
172    ///
173    /// # Errors
174    ///
175    /// Returns [`GhostError::MiiError`] if the Mii data is invalid, or
176    /// [`GhostError::CTGPFooterError`] if the SHA-1 field cannot be written.
177    pub fn update_raw_data(&mut self) -> Result<(), GhostError> {
178        let mii_bytes = self.header().mii().raw_data().to_vec();
179        self.header_mut().set_mii(Mii::new(mii_bytes)?);
180        self.header_mut().fix_mii_crc16();
181
182        let mut buf = Vec::from(self.header().raw_data());
183
184        buf.extend_from_slice(self.input_data().raw_data());
185
186        let header_len = 0x88;
187        let new_input_data_end = header_len + self.input_data().raw_data().len();
188
189        // Find input data length of old data
190        let old_input_data_end = if self.raw_data[0x8C..0x90] == *b"Yaz1" {
191            u32::from_be_bytes(self.raw_data[0x88..0x8C].try_into().unwrap()) as usize
192        } else {
193            header_len + 0x2774
194        };
195
196        if new_input_data_end > old_input_data_end {
197            let diff = new_input_data_end - old_input_data_end;
198            let insert_pos = old_input_data_end;
199            self.raw_data
200                .splice(insert_pos..insert_pos, vec![0u8; diff]);
201        } else if new_input_data_end < old_input_data_end {
202            let diff = old_input_data_end - new_input_data_end;
203            let remove_end = old_input_data_end;
204            self.raw_data.drain(remove_end - diff..remove_end);
205        }
206
207        self.raw_data[..new_input_data_end].copy_from_slice(&buf[..new_input_data_end]);
208        let base_crc32 = crc32(&buf);
209
210        match (self.footer(), self.should_preserve_external_footer()) {
211            (Some(FooterType::CTGPFooter(ctgp_footer)), true) => {
212                buf.extend_from_slice(&base_crc32.to_be_bytes());
213                buf.extend_from_slice(ctgp_footer.raw_data());
214
215                let footer_len = ctgp_footer.len();
216                self.raw_data.drain(new_input_data_end..);
217                self.raw_data
218                    .extend_from_slice(&buf[buf.len() - footer_len - 0x04..]);
219                self.raw_data.extend_from_slice(&[0u8; 4]);
220            }
221            (Some(FooterType::SPFooter(sp_footer)), true) => {
222                buf.extend_from_slice(&base_crc32.to_be_bytes());
223                buf.extend_from_slice(sp_footer.raw_data());
224
225                let footer_len = sp_footer.len();
226                self.raw_data.drain(new_input_data_end..);
227                self.raw_data
228                    .extend_from_slice(&buf[buf.len() - footer_len - 0x04..]);
229                self.raw_data.extend_from_slice(&[0u8; 4]);
230            }
231            (_, true) if self.raw_data.len() >= new_input_data_end + 0x08 => {
232                self.raw_data.drain(new_input_data_end + 0x04..);
233            }
234            (_, false) if self.raw_data.len() >= new_input_data_end + 0x08 => self.raw_data
235                [new_input_data_end..new_input_data_end + 0x04]
236                .copy_from_slice(&base_crc32.to_be_bytes()),
237            _ => (),
238        }
239
240        let len = self.raw_data.len();
241        let crc32 = crc32(&self.raw_data[..len - 0x04]);
242        self.raw_data[len - 0x04..].copy_from_slice(&crc32.to_be_bytes());
243
244        let sha1 = compute_sha1_hex(&self.raw_data);
245        if let Some(FooterType::CTGPFooter(ctgp_footer)) = self.footer_mut() {
246            ctgp_footer.set_ghost_sha1(&sha1)?;
247        }
248
249        Ok(())
250    }
251
252    /// Flushes all modifications and writes the ghost to a file at the given path.
253    ///
254    /// # Errors
255    ///
256    /// Returns any error from [`update_raw_data`](Ghost::update_raw_data) or
257    /// from file creation/writing.
258    pub fn save_to_file<T: AsRef<std::path::Path>>(&mut self, path: T) -> Result<(), GhostError> {
259        self.update_raw_data()?;
260        let mut file = std::fs::File::create(path)?;
261        file.write_all(&self.raw_data)?;
262
263        Ok(())
264    }
265
266    /// Compresses the input data using Yaz1 encoding and sets the compression flag in the header.
267    ///
268    /// Does nothing if the input data is already compressed.
269    pub fn compress_input_data(&mut self) {
270        if self.input_data().is_compressed() {
271            return;
272        }
273
274        self.input_data_mut().compress();
275        self.header_mut().set_compressed(true);
276    }
277
278    /// Decompresses the input data and clears the compression flag in the header.
279    ///
280    /// Does nothing if the input data is not compressed.
281    pub fn decompress_input_data(&mut self) {
282        if !self.input_data().is_compressed() {
283            return;
284        }
285
286        self.input_data_mut().decompress();
287        self.header_mut().set_compressed(false);
288    }
289
290    /// Returns the raw file bytes.
291    ///
292    /// May not reflect recent modifications until [`update_raw_data`](Ghost::update_raw_data) is called.
293    pub fn raw_data(&self) -> &[u8] {
294        &self.raw_data
295    }
296
297    /// Returns a mutable reference to the raw file bytes.
298    pub fn raw_data_mut(&mut self) -> &mut [u8] {
299        &mut self.raw_data
300    }
301
302    /// Returns the parsed RKG file header.
303    pub fn header(&self) -> &Header {
304        &self.header
305    }
306
307    /// Returns a mutable reference to the parsed RKG file header.
308    pub fn header_mut(&mut self) -> &mut Header {
309        &mut self.header
310    }
311
312    /// Returns the ghost's controller input data.
313    pub fn input_data(&self) -> &InputData {
314        &self.input_data
315    }
316
317    /// Returns a mutable reference to the ghost's controller input data.
318    pub fn input_data_mut(&mut self) -> &mut InputData {
319        &mut self.input_data
320    }
321
322    /// Returns the footer, if present.
323    pub fn footer(&self) -> Option<&FooterType> {
324        self.footer.as_ref()
325    }
326
327    /// Returns a mutable reference to the footer, if present.
328    pub fn footer_mut(&mut self) -> Option<&mut FooterType> {
329        self.footer.as_mut()
330    }
331
332    /// Returns the CRC-32 of the header and input data, excluding any external footer.
333    pub fn base_crc32(&self) -> u32 {
334        self.base_crc32
335    }
336
337    /// Returns `true` if the stored base CRC-32 matches a freshly computed
338    /// checksum of the current header and input data bytes.
339    pub fn verify_base_crc32(&self) -> bool {
340        let mut data = Vec::from(self.header().raw_data());
341        data.extend_from_slice(self.input_data().raw_data());
342        self.base_crc32 == crc32(&data)
343    }
344
345    /// Returns the CRC-32 of the entire file excluding its final 4 bytes.
346    pub fn file_crc32(&self) -> u32 {
347        self.file_crc32
348    }
349
350    /// Returns `true` if the stored file CRC-32 matches a freshly computed
351    /// checksum of the entire file excluding its final 4 bytes.
352    pub fn verify_file_crc32(&self) -> bool {
353        let len = self.raw_data().len();
354        self.file_crc32 == crc32(&self.raw_data()[..len - 0x04])
355    }
356
357    /// Returns whether an existing external footer will be preserved when saving.
358    pub fn should_preserve_external_footer(&self) -> bool {
359        self.should_preserve_external_footer
360    }
361
362    /// Sets whether an existing external footer should be preserved when saving.
363    pub fn set_should_preserve_external_footer(&mut self, b: bool) {
364        self.should_preserve_external_footer = b;
365    }
366}
367
368/// Used internally for writing bits to a buffer.
369pub(crate) fn write_bits(
370    buf: &mut [u8],
371    byte_offset: usize,
372    bit_offset: usize,
373    bit_width: usize,
374    value: u64,
375) {
376    let bytes_needed = (bit_offset + bit_width).div_ceil(8);
377    let mut chunk: u64 = 0;
378
379    for i in 0..bytes_needed {
380        chunk = (chunk << 8) | buf[byte_offset + i] as u64;
381    }
382
383    let shift = bytes_needed * 8 - bit_offset - bit_width;
384    let mask = ((1u64 << bit_width) - 1) << shift;
385
386    chunk = (chunk & !mask) | ((value << shift) & mask);
387
388    for i in (0..bytes_needed).rev() {
389        buf[byte_offset + i] = (chunk & 0xFF) as u8;
390        chunk >>= 8;
391    }
392}
393
394/// Computes the SHA-1 hash of `input` and returns it as a 20-byte array.
395pub(crate) fn compute_sha1_hex(input: &[u8]) -> [u8; 0x14] {
396    let mut hasher = Sha1::new();
397    hasher.update(input);
398    hasher.finalize().into()
399}
400
401/// Converts a raw Wii tick count into a [`NaiveDateTime`].
402///
403/// Ticks run at 60.75 MHz relative to the Wii epoch of 2000-01-01 00:00:00 UTC.
404pub(crate) fn datetime_from_timestamp(tick_count: u64) -> NaiveDateTime {
405    let clock_rate = 60_750_000.0; // 60.75 MHz tick speed
406    let epoch_shift = 946_684_800; // Shifts epoch from 1970-01-01 to 2000-01-01 (which is what the Wii uses)
407    let total_seconds = tick_count as f64 / clock_rate;
408    let total_nanoseconds = (total_seconds * 1_000_000_000.0) as i64;
409
410    let duration = Duration::nanoseconds(total_nanoseconds);
411    let epoch = DateTime::from_timestamp(epoch_shift, 0).unwrap();
412
413    epoch.naive_utc() + duration
414}
415
416/// Converts a raw Wii tick count into a [`TimeDelta`] (duration from an arbitrary reference).
417///
418/// Ticks run at 60.75 MHz; the result is truncated to millisecond precision.
419pub(crate) fn duration_from_ticks(tick_count: u64) -> TimeDelta {
420    let clock_rate = 60_750_000.0; // 60.75 MHz tick speed
421    let total_seconds = tick_count as f64 / clock_rate;
422    let total_milliseconds = (total_seconds * 1_000.0) as i64;
423
424    Duration::milliseconds(total_milliseconds)
425}