pkcrack 0.1.0

//! Cryptographic operations and precomputed tables for the attack

use crate::types::*;
use crate::crc::Crc32Computer;
use crate::error::Result;
use lazy_static::lazy_static;

/// Precomputed temporary value table for stage 1
/// temp_table[key3][i] where i ranges 0-63
static mut TEMP_TABLE: [[u32; 64]; 256] = [[0; 64]; 256];
static mut TEMP_TABLE_INITIALIZED: bool = false;

/// Inverse temporary value table (matches C code exactly)
static mut INV_TEMP_TABLE: [[[u8; 5]; 64]; 256] = [[[0; 5]; 64]; 256];
static mut INV_TEMP_TABLE_INITIALIZED: bool = false;

/// Entry counts for inverse temp table (matches C code)
static mut INV_ENTRIES: [[u16; 64]; 256] = [[0; 64]; 256];
static mut INV_ENTRIES_INITIALIZED: bool = false;

/// Multiplication table for constant operations
static mut MUL_TABLE: [u32; 256] = [0; 256];
static mut MUL_TABLE_INITIALIZED: bool = false;

/// Additional multiplication table mTab2
static mut MTAB2: [[u32; 2]; 256] = [[0; 2]; 256];
static mut MTAB2_INITIALIZED: bool = false;

/// CRC inverse table (matches C code)
static mut CRCINV_TABLE: [u32; 256] = [0; 256];
static mut CRCINV_TABLE_INITIALIZED: bool = false;

/// Initialize all precomputed tables
pub fn precompute_tables() -> Result<()> {
    precompute_temp_table()?;
    precompute_inv_temp_table()?;
    init_crcinv_table()?;
    init_mul_table()?;
    init_mtab2()?;
    Ok(())
}

/// Precompute temporary value table (matches C code exactly)
fn precompute_temp_table() -> Result<()> {
    unsafe {
        if TEMP_TABLE_INITIALIZED {
            return Ok(());
        }

        let mut num_entries: [usize; 256] = [0; 256];

        // Match C code logic exactly
        let mut temp = 0x10000;
        loop {
            let key3 = (((temp | 2) * (temp | 3)) >> 8) & 0xFF;

            // Check bounds before indexing
            if num_entries[key3] < 64 && key3 < 256 {
                TEMP_TABLE[key3][num_entries[key3]] = temp as u32;
                num_entries[key3] += 1;
            }

            if temp == 0 {
                break;
            }
            temp = temp.wrapping_sub(4);
        }

        TEMP_TABLE_INITIALIZED = true;
    }
    Ok(())
}

/// Initialize CRC inverse table (matches C code)
fn init_crcinv_table() -> Result<()> {
    unsafe {
        if CRCINV_TABLE_INITIALIZED {
            return Ok(());
        }

        // Build inverse CRC table
        for i in 0..256 {
            let mut crc = (i as u32) << 24;
            for _ in 0..8 {
                if crc & 0x80000000 != 0 {
                    crc = (crc << 1) ^ 0xedb88320; // CRC polynomial
                } else {
                    crc <<= 1;
                }
            }
            CRCINV_TABLE[i] = crc;
        }

        CRCINV_TABLE_INITIALIZED = true;
    }
    Ok(())
}

/// Precompute inverse temporary value table (matches C code exactly)
fn precompute_inv_temp_table() -> Result<()> {
    unsafe {
        if INV_TEMP_TABLE_INITIALIZED {
            return Ok(());
        }

        // Initialize invEntries to zeros
        for i in 0..256 {
            for j in 0..64 {
                INV_ENTRIES[i][j] = 0;
            }
        }

        // Match C code logic exactly
        for key3 in 0..256usize {
            for temp in 0..64 {
                let index = ((TEMP_TABLE[key3][temp] >> 10) & 63) as usize;
                INV_TEMP_TABLE[key3][index][INV_ENTRIES[key3][index] as usize] = temp as u8;
                INV_ENTRIES[key3][index] += 1;
            }
        }

        INV_TEMP_TABLE_INITIALIZED = true;
        INV_ENTRIES_INITIALIZED = true;
    }
    Ok(())
}

/// Initialize multiplication table
fn init_mul_table() -> Result<()> {
    unsafe {
        if MUL_TABLE_INITIALIZED {
            return Ok(());
        }

        for i in 0..256 {
            MUL_TABLE[i] = ((i as u64).wrapping_mul(CONST as u64) & 0xFFFFFFFF) as u32;
        }

        MUL_TABLE_INITIALIZED = true;
    }
    Ok(())
}

/// Initialize mTab2 table
fn init_mtab2() -> Result<()> {
    unsafe {
        if MTAB2_INITIALIZED {
            return Ok(());
        }

        for i in 0..256 {
            MTAB2[i][0] = ((i as u64).wrapping_mul(CONST as u64).wrapping_add(1) & 0xFFFFFFFF) as u32;
            MTAB2[i][1] = (((i as u64).wrapping_mul(CONST as u64)).wrapping_add(CONST as u64) & 0xFFFFFFFF) as u32;
        }

        MTAB2_INITIALIZED = true;
    }
    Ok(())
}

/// Get temporary value from precomputed table
#[inline]
pub fn get_temp_value(key3: u8, i: usize) -> Option<u32> {
    unsafe {
        if TEMP_TABLE_INITIALIZED && i < 64 {
            Some(TEMP_TABLE[key3 as usize][i])
        } else {
            None
        }
    }
}

/// Get inverse temporary values (matches C code exactly)
#[inline]
pub fn get_inv_temp_values(key3: u8, index: usize) -> Option<&'static [u8]> {
    unsafe {
        if INV_TEMP_TABLE_INITIALIZED && index < 64 {
            Some(&INV_TEMP_TABLE[key3 as usize][index])
        } else {
            None
        }
    }
}

/// Get inverse entry count (matches C code exactly)
#[inline]
pub fn get_inv_entries(key3: u8, index: usize) -> Option<u16> {
    unsafe {
        if INV_ENTRIES_INITIALIZED && index < 64 {
            Some(INV_ENTRIES[key3 as usize][index])
        } else {
            None
        }
    }
}

/// Get CRC inverse table value (matches C code exactly)
#[inline]
pub fn get_crcinv_table(byte: u8) -> Option<u32> {
    unsafe {
        if CRCINV_TABLE_INITIALIZED {
            Some(CRCINV_TABLE[byte as usize])
        } else {
            None
        }
    }
}

/// Get multiplication table value
#[inline]
pub fn get_mul_table_value(i: u8) -> Option<u32> {
    unsafe {
        if MUL_TABLE_INITIALIZED {
            Some(MUL_TABLE[i as usize])
        } else {
            None
        }
    }
}

/// Get mTab2 table values
#[inline]
pub fn get_mtab2_values(i: u8) -> Option<[u32; 2]> {
    unsafe {
        if MTAB2_INITIALIZED {
            Some(MTAB2[i as usize])
        } else {
            None
        }
    }
}

/// Key state manager with precomputed tables
#[derive(Clone)]
pub struct KeyManager {
    crc_computer: Crc32Computer,
}

impl KeyManager {
    pub fn new() -> Result<Self> {
        precompute_tables()?;
        Ok(Self {
            crc_computer: Crc32Computer::new(),
        })
    }

    /// Initialize keys from password
    pub fn keys_from_password(&self, password: &[u8]) -> KeyState {
        let mut state = KeyState {
            key0: 0x12345678,
            key1: 0x23456789,
            key2: 0x34567890,
        };

        for &byte in password {
            self.update_keys(&mut state, byte);
        }

        state
    }

    /// Update keys with a plaintext byte
    pub fn update_keys(&self, state: &mut KeyState, plain_byte: u8) {
        state.key0 = self.crc_computer.crc32(state.key0, plain_byte);
        state.key1 = (state.key1 + (state.key0 & 0xFF)).wrapping_mul(CONST).wrapping_add(1);
        state.key2 = self.crc_computer.crc32(state.key2, (state.key1 >> 24) as u8);
    }

    /// Get current key3 value
    pub fn get_key3(&self, state: &KeyState) -> u8 {
        let temp = (state.key2 & 0xFFFF) | 3;
        ((temp.wrapping_mul(temp ^ 1)) >> 8) as u8
    }

    /// Reverse key update operation (for attack)
    pub fn reverse_key_update(&self, state: &mut KeyState, plain_byte: u8) {
        // Reverse the key updates
        state.key2 = self.crc_computer.inv_crc32(state.key2, (state.key1 >> 24) as u8);

        // Reverse key1 update
        let mut temp = state.key1;
        temp = temp.wrapping_sub(1);
        // Find inverse of multiplication by CONST
        temp = ((temp as u64).wrapping_mul(INVCONST as u64) & 0xFFFFFFFF) as u32;
        temp = temp.wrapping_sub(state.key0 & 0xFF);
        state.key1 = temp;

        // Reverse key0 update
        state.key0 = self.crc_computer.inv_crc32(state.key0, plain_byte);
    }

    /// Encrypt data buffer
    pub fn encrypt(&self, state: &mut KeyState, data: &mut [u8]) {
        for byte in data.iter_mut() {
            *byte ^= self.get_key3(state);
            self.update_keys(state, *byte);
        }
    }

    /// Decrypt data buffer
    pub fn decrypt(&self, state: &mut KeyState, data: &mut [u8]) {
        for byte in data.iter_mut() {
            let decrypted = *byte ^ self.get_key3(state);
            self.update_keys(state, decrypted);
            *byte = decrypted;
        }
    }

    /// Validate that keys can correctly encrypt/decrypt test data
    pub fn validate_keys(&self, mut state: KeyState, plaintext: &[u8], ciphertext: &[u8]) -> bool {
        if plaintext.len() != ciphertext.len() {
            return false;
        }

        // Try to decrypt ciphertext and compare with plaintext
        for i in 0..plaintext.len() {
            let decrypted = ciphertext[i] ^ self.get_key3(&state);
            if decrypted != plaintext[i] {
                return false;
            }
            self.update_keys(&mut state, decrypted);
        }

        true
    }

    /// Generate random initial state for testing
    pub fn random_state(&self) -> KeyState {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        let mut hasher = DefaultHasher::new();
        std::time::SystemTime::now().hash(&mut hasher);
        let seed = hasher.finish() as u32;

        KeyState {
            key0: seed,
            key1: seed.wrapping_mul(0x12345678),
            key2: seed.wrapping_mul(0x87654321),
        }
    }
}

impl Default for KeyManager {
    fn default() -> Self {
        Self::new().expect("Failed to initialize KeyManager")
    }
}

/// Global key manager instance
lazy_static! {
    static ref KEY_MANAGER: KeyManager = KeyManager::new().expect("Failed to initialize KeyManager");
}

/// Get global key manager
pub fn get_key_manager() -> Result<&'static KeyManager> {
    Ok(&KEY_MANAGER)
}

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

    #[test]
    fn test_key_manager_initialization() {
        let manager = KeyManager::new().unwrap();
        assert!(get_key_manager().is_ok());
    }

    #[test]
    fn test_password_to_keys() {
        let manager = KeyManager::new().unwrap();
        let password = b"secret";
        let keys = manager.keys_from_password(password);

        // Should not panic and produce some values
        assert!(keys.key0 != 0 || keys.key1 != 0 || keys.key2 != 0);
    }

    #[test]
    fn test_encrypt_decrypt_roundtrip() {
        let manager = KeyManager::new().unwrap();
        let mut state = manager.keys_from_password(b"password");

        let plaintext = b"Hello, World!";
        let mut ciphertext = plaintext.to_vec();

        // Encrypt
        manager.encrypt(&mut state, &mut ciphertext);

        // Reset keys
        state = manager.keys_from_password(b"password");

        // Decrypt
        manager.decrypt(&mut state, &mut ciphertext);

        assert_eq!(plaintext, &ciphertext[..]);
    }

    #[test]
    fn test_key_validation() {
        let manager = KeyManager::new().unwrap();
        let password = b"testpassword";
        let mut state = manager.keys_from_password(password);

        let plaintext = b"Test data for validation";
        let mut ciphertext = plaintext.to_vec();

        // Encrypt
        manager.encrypt(&mut state, &mut ciphertext);

        // Validate
        let validation_state = manager.keys_from_password(password);
        assert!(manager.validate_keys(validation_state, plaintext, &ciphertext));
    }
}