letterbomb 5.0.0

use crate::mac::WiiMAC;
use crate::region::Region;
use std::cmp::PartialEq;

use chrono::{DateTime, Datelike, Timelike, Utc};
use hex::encode;
use hmac::{Hmac, Mac};
use sha1::Sha1;

use crate::bundle::bundle;
use std::collections::HashMap;
use std::fs;
use std::path::Path;

/// Packs `blob` with corresponding timestamp and the MAC digest, `digest`.
///
/// # Arguments
///
/// * `digest` - The MAC digest bytes.
/// * `time_stamp` - Unix epoch timestamp.
/// * `blob` - The blob content buffer to modify.
///
/// # Returns
///
/// The modified blob content.
pub fn pack_blob(digest: &[u8], unix_timestamp: u32, mut blob: Vec<u8>) -> Option<Vec<u8>> {
    blob[0x08..0x10].copy_from_slice(&digest[..8]);
    blob[0x7C..0x80].copy_from_slice(&unix_timestamp.to_be_bytes());
    blob[0x80..0x8A].copy_from_slice(format!("{:010}", unix_timestamp).as_bytes());
    let mut hmac = Hmac::<Sha1>::new_from_slice(&digest[8..]).ok()?;
    hmac.update(&blob);
    let hmac_result = hmac.finalize().into_bytes();
    blob[0xB0..0xC4].copy_from_slice(&hmac_result);
    Some(blob)
}

/// Generates the relative SD card path for the LetterBomb based on the provided digest, and time of letter reception.
///
/// # Parameters:
/// - `digest`: MAC digest, typically generated from the `mac_digest` function.
/// - `timestamp`: The time the letter was received.
///
/// # Returns:
/// - A `String` representing the path on the SD card, relative to the root of it.
pub fn sd_path(digest: &[u8], timestamp: DateTime<Utc>) -> String {
    format!(
        "private/wii/title/HAEA/{}/{}/{:04}/{:02}/{:02}/{:02}/{:02}/HABA_#1/txt/{:08X}.000",
        encode(&digest[0..4]).to_uppercase(),
        encode(&digest[4..8]).to_uppercase(),
        timestamp.year(),
        timestamp.month() - 1,
        timestamp.day(),
        timestamp.hour(),
        timestamp.minute(),
        timestamp.timestamp()
    )
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum IncludeBundle {
    IncludeHackMii,
    ExcludeHackMii,
}

/// Constructs a payload for a given WiiMAC and region, optionally including bundled files.
///
/// This function performs the following steps:
/// 1. Computes the SHA-1 digest of the provided `mac`.
/// 2. Captures the current UTC timestamp.
/// 3. Packs the digest, timestamp, and region-specific template data into a single blob.
/// 4. Inserts that blob into a `HashMap` under a path generated from the digest and timestamp.
/// 5. If `include_bundle` is `true`, iterates over `BUNDLED.entries()` and appends each file's
///    contents under its filename as an additional entry in the map.
///
/// # Parameters
///
/// * `mac` - The `WiiMAC` whose SHA-1 digest will be used to generate the payload.
/// * `region` - The `Region` selecting which template data to include in the blob.
/// * `include_bundle` - When `true`, all files in the built-in bundle will also be added
///   to the resulting map, keyed by their filenames.
///
/// # Returns
///
/// On success, returns `Some(HashMap<String, Vec<u8>>>` where:
/// - The primary entry's key is determined by `sd_path(&digest, timestamp)`
///   and its value is the packed blob.
/// - Each bundled file (if so requested) appears under its own filename as the key
///   and its raw contents as the value.
///
/// Returns `None` if the core blob packing (`pack_blob`) fails.
///
/// # Examples
///
/// ```rust
/// use chrono::Utc;
/// use std::collections::HashMap;
/// use letterbomb::generate::{make_payload, IncludeBundle};
/// use letterbomb::mac::WiiMAC;
/// use letterbomb::region::Region;
///
/// // Suppose you have a valid WiiMAC value
/// let mac = WiiMAC::try_from((0x00, 0x17, 0xAB, 0x5A, 0x6E, 0xF5)).expect("invalid mac");
/// let region = Region::E;
/// // Create a payload without bundled files
/// if let Some(map) = make_payload(&mac, &region, IncludeBundle::ExcludeHackMii, Utc::now()) {
///     // The map contains at least the primary SD blob
///     assert!(!map.is_empty());
/// }
/// ```
pub fn make_payload(
    mac: &WiiMAC,
    region: &Region,
    include_bundle: IncludeBundle,
    time: DateTime<Utc>,
) -> Option<HashMap<String, Vec<u8>>> {
    let dig = mac.sha1_digest();
    let mut sd: HashMap<String, Vec<u8>> = HashMap::new();
    let blob = pack_blob(
        &dig,
        time.timestamp() as u32,
        region.template_for().to_vec(),
    )?;
    sd.insert(sd_path(&dig, time), blob);
    if include_bundle == IncludeBundle::IncludeHackMii {
        let bundle = bundle();
        if let Ok(e) = bundle {
            for (name, contents) in e.iter() {
                sd.insert(name.to_owned(), contents.to_vec());
            }
        }
    }
    Some(sd)
}

pub fn write_subpaths(base: &Path, files: &HashMap<String, Vec<u8>>) -> Result<(), std::io::Error> {
    for (subpath, bytes) in files {
        let dest = base.join(subpath);
        if let Some(parent) = dest.parent() {
            fs::create_dir_all(parent)?;
        }
        fs::write(&dest, bytes)?;
        println!(
            "wrote {:>9} bytes to {}",
            &bytes.len(),
            dest.to_string_lossy()
        );
    }
    Ok(())
}

/// Inline-able shortcut that calls `make_payload` with the current UTC timestamp.
/// Returns `None` if payload generation fails.
#[inline]
pub fn make_payload_now(
    mac: &WiiMAC,
    region: &Region,
    include_bundle: IncludeBundle,
) -> Option<HashMap<String, Vec<u8>>> {
    make_payload(mac, region, include_bundle, Utc::now())
}

#[cfg(test)]
mod tests {
    use super::{IncludeBundle, make_payload_now, pack_blob, sd_path};
    use crate::mac::WiiMAC;
    use crate::region::Region;
    use chrono::{TimeZone, Utc};
    use hmac::{Hmac, Mac};
    use sha1::Sha1;

    #[test]
    fn pack_blob_ok() {
        // prepare a 20‐byte digest all set to 1
        let digest = [1u8; 20];
        // 0xC4 is min size
        let template = vec![0u8; 0xC4];
        let timestamp = 0x11223344u32;

        let packed =
            pack_blob(&digest, timestamp, template).expect("pack_blob failed; HMAC digest failed?");

        // first 8 bytes of digest go at 0x08..0x10
        assert_eq!(&packed[0x08..0x10], &digest[..8]);

        // the big‐endian timestamp goes at 0x7C..0x80
        assert_eq!(&packed[0x7C..0x80], &timestamp.to_be_bytes());

        // recompute our expected HMAC, get a 'untainted' copy
        let mut blob = packed.clone();
        blob[0xB0..0xC4].fill(0);

        let mut hmac = Hmac::<Sha1>::new_from_slice(&digest[8..]).unwrap();
        hmac.update(&blob);
        let expected = hmac.finalize().into_bytes();

        assert_eq!(&packed[0xB0..0xC4], &expected[..]);
    }

    #[test]
    fn sd_path_contains_all_components() {
        let digest = [
            0xDE, 0xAD, 0xBE, 0xEF, 0xCA, 0xFE, 0xBA, 0xBE, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        ];
        let timestamp = Utc.with_ymd_and_hms(2020, 5, 20, 12, 30, 45);
        let path = sd_path(&digest, timestamp.unwrap());

        assert!(path.eq(
            "private/wii/title/HAEA/DEADBEEF/CAFEBABE/2020/04/20/12/30/HABA_#1/txt/5EC522F5.000"
        ));
    }

    #[test]
    fn payload_ok() {
        let mac: WiiMAC = "00:17:ab:5a:6e:f5".parse().unwrap();
        let region = Region::U;

        let payload = make_payload_now(&mac, &region, IncludeBundle::ExcludeHackMii)
            .expect("make_payload failed; pack_blob failed?");
        assert_eq!(payload.len(), 1);

        println!("payload files: {} files", payload.len());
        let mut total = 0usize;
        for (name, data) in payload.iter() {
            total += data.len();
            println!("in payload: {:<26}: {:>11} bytes", name, data.len());
        }
        println!("payload size: {} bytes", total);

        let (path, blob) = payload.into_iter().next().unwrap();

        // Must start with the SD-path prefix (cannot match further because of time)
        assert!(path.starts_with("private/wii/title/HAEA/"));

        let filename = path.rsplit('/').next().expect("no filename in path");
        assert!(filename.ends_with(".000"));

        // Get back the 8-hex-digit timestamp
        let hex_ts = &filename[..8];
        let timestamp = u32::from_str_radix(hex_ts, 16).expect("invalid hex timestamp");

        // Recompute digest and expected blob
        let digest = mac.sha1_digest();
        let expected = pack_blob(&digest, timestamp, region.template_for().to_vec())
            .expect("pack_blob failed");

        assert_eq!(blob, expected);
    }

    #[test]
    fn payload_with_bundle_ok() {
        let mac: WiiMAC = "00:17:ab:5a:6e:f5".parse().unwrap();
        let region = Region::U;

        let payload = make_payload_now(&mac, &region, IncludeBundle::IncludeHackMii)
            .expect("make_payload failed; pack_blob failed?");
        // Must have bundle too
        assert!(payload.len() > 1);
        println!("payload files: {} files", payload.len());
        let mut total = 0usize;
        for (name, data) in payload.iter() {
            total += data.len();
            println!("in payload: {:<26}: {:>11} bytes", name, data.len());
        }
        println!("payload size: {} bytes", total);
    }
}