wzlib-rs 0.1.1

//! WZ key management — lazy key generation matching `WzMutableKey`.
//!
//! Keys are generated by AES-ECB encrypting an IV repeatedly.
//! The resulting byte stream is used for XOR-based string and data decryption.

use crate::crypto::aes_encryption::generate_wz_key;

pub struct WzKey {
    iv: [u8; 4],
    user_key: Option<[u8; 128]>,
    keys: Vec<u8>,
}

impl WzKey {
    pub fn new(iv: [u8; 4]) -> Self {
        WzKey {
            iv,
            user_key: None,
            keys: Vec::new(),
        }
    }

    /// Create a `WzKey` with a custom 128-byte user key for AES key generation.
    pub fn with_user_key(iv: [u8; 4], user_key: [u8; 128]) -> Self {
        WzKey {
            iv,
            user_key: Some(user_key),
            keys: Vec::new(),
        }
    }

    pub fn ensure_size(&mut self, size: usize) {
        if self.keys.len() < size {
            self.keys = generate_wz_key(&self.iv, size, self.user_key.as_ref());
        }
    }

    pub fn get(&mut self, index: usize) -> u8 {
        self.ensure_size(index + 1);
        self.keys[index]
    }

    pub fn get_slice(&mut self, start: usize, len: usize) -> &[u8] {
        self.ensure_size(start + len);
        &self.keys[start..start + len]
    }

    pub fn iv(&self) -> [u8; 4] {
        self.iv
    }

}

impl std::ops::Index<usize> for WzKey {
    type Output = u8;

    fn index(&self, index: usize) -> &Self::Output {
        &self.keys[index]
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // ── ensure_size ───────────────────────────────────────────────

    #[test]
    fn test_ensure_size_generates_keys() {
        let mut key = WzKey::new([0; 4]);
        assert_eq!(key.keys.len(), 0);
        key.ensure_size(256);
        assert!(key.keys.len() >= 256);
    }

    #[test]
    fn test_ensure_size_no_shrink() {
        let mut key = WzKey::new([0; 4]);
        key.ensure_size(512);
        let len1 = key.keys.len();
        key.ensure_size(100); // smaller request
        assert_eq!(key.keys.len(), len1); // should not shrink
    }

    #[test]
    fn test_ensure_size_expands() {
        let mut key = WzKey::new([0; 4]);
        key.ensure_size(16);
        let first = key.keys.clone();
        key.ensure_size(1024);
        assert!(key.keys.len() >= 1024);
        // First 16 bytes of expanded key should still match (regenerated from same IV)
        assert_eq!(&key.keys[..16], &first[..16]);
    }

    // ── get ───────────────────────────────────────────────────────

    #[test]
    fn test_get_auto_expands() {
        let mut key = WzKey::new([0; 4]);
        // keys buffer is empty, get(0) should auto-expand
        let _byte = key.get(0);
        assert!(!key.keys.is_empty());
    }

    #[test]
    fn test_get_deterministic() {
        let mut key1 = WzKey::new([0x4D, 0x23, 0xC7, 0x2B]);
        let mut key2 = WzKey::new([0x4D, 0x23, 0xC7, 0x2B]);
        for i in 0..32 {
            assert_eq!(key1.get(i), key2.get(i), "mismatch at index {}", i);
        }
    }

    #[test]
    fn test_get_different_iv_different_keys() {
        let mut gms = WzKey::new([0x4D, 0x23, 0xC7, 0x2B]);
        let mut ems = WzKey::new([0xB9, 0x7D, 0x63, 0xE9]);
        // Different IVs should produce different keys
        let differs = (0..32).any(|i| gms.get(i) != ems.get(i));
        assert!(differs);
    }

    // ── get_slice ─────────────────────────────────────────────────

    #[test]
    fn test_get_slice_returns_correct_range() {
        let mut key = WzKey::new([0x4D, 0x23, 0xC7, 0x2B]);
        key.ensure_size(64);
        let expected: Vec<u8> = key.keys[10..30].to_vec();
        let slice = key.get_slice(10, 20);
        assert_eq!(slice.len(), 20);
        assert_eq!(slice, expected.as_slice());
    }

    #[test]
    fn test_get_slice_auto_expands() {
        let mut key = WzKey::new([0; 4]);
        let slice = key.get_slice(0, 100);
        assert_eq!(slice.len(), 100);
        assert!(key.keys.len() >= 100);
    }

    #[test]
    fn test_get_slice_at_offset() {
        let mut key = WzKey::new([0x4D, 0x23, 0xC7, 0x2B]);
        let snapshot: Vec<u8> = key.get_slice(50, 10).to_vec();
        assert_eq!(snapshot.len(), 10);
        // Verify it matches individual get() calls
        for (i, &b) in snapshot.iter().enumerate() {
            assert_eq!(b, key.get(50 + i));
        }
    }

    #[test]
    fn test_get_slice_empty() {
        let mut key = WzKey::new([0; 4]);
        let slice = key.get_slice(0, 0);
        assert!(slice.is_empty());
    }

    // ── Index trait ───────────────────────────────────────────────

    #[test]
    fn test_index_after_ensure_size() {
        let mut key = WzKey::new([0x4D, 0x23, 0xC7, 0x2B]);
        key.ensure_size(16);
        // Index trait should work for indices within the generated range
        let byte = key[0];
        let byte15 = key[15];
        // Verify consistency with get()
        assert_eq!(byte, key.get(0));
        assert_eq!(byte15, key.get(15));
    }

    #[test]
    #[should_panic]
    fn test_index_panics_without_ensure() {
        let key = WzKey::new([0; 4]);
        // keys buffer is empty, indexing should panic
        let _ = key[0];
    }

    // ── Zero IV produces zero keys ────────────────────────────────

    #[test]
    fn test_zero_iv_produces_consistent_keys() {
        let mut key1 = WzKey::new([0; 4]);
        let mut key2 = WzKey::new([0; 4]);
        let slice1 = key1.get_slice(0, 64);
        let slice2 = key2.get_slice(0, 64);
        assert_eq!(slice1, slice2);
    }
}