pkcrack 0.1.0

//! Stage 2 of the pkcrack attack: Reconstruct complete key sequences

use crate::types::*;
use crate::crc::{Crc32Computer, crc32_fast, inv_crc32_fast};
use crate::crypto::get_key_manager;
use crate::error::{PkCrackError, Result};
use crate::{msb, CONST, INVCONST};

/// Stage 2 processor for key reconstruction
pub struct Stage2Processor {
    crc_computer: Crc32Computer,
    solutions: Vec<Stage2Result>,
    found_solutions: bool,
}

impl Stage2Processor {
    /// Create new Stage 2 processor
    pub fn new() -> Result<Self> {
        Ok(Self {
            crc_computer: Crc32Computer::new(),
            solutions: Vec::new(),
            found_solutions: false,
        })
    }

    /// Execute Stage 2: find complete key sequences from key2 candidates
    pub fn execute(&mut self, key2_candidates: &[u32], plaintext: &[u8], ciphertext: &[u8],
                  offset: usize, abort_on_success: bool) -> Result<Vec<Stage2Result>> {
        self.solutions.clear();
        self.found_solutions = false;

        let text_length = plaintext.len();

        println!("Starting Stage 2 with {} key2 candidates", key2_candidates.len());

        for (i, &key2_13) in key2_candidates.iter().enumerate() {
            if abort_on_success && self.found_solutions {
                println!("Aborting Stage 2 after finding solution");
                break;
            }

            if i % 100 == 0 || key2_candidates.len() <= 100 {
                println!("Searching... {:.1}%", 100.0 * i as f64 / key2_candidates.len() as f64);
            }

            // Build key sequence from key2_13
            if let Some(solution) = self.build_key_sequence(key2_13, plaintext, ciphertext, offset, text_length)? {
                self.solutions.push(solution);
                self.found_solutions = true;

                if abort_on_success {
                    break;
                }
            }
        }

        println!("Stage 2 completed. Found {} solutions", self.solutions.len());
        Ok(self.solutions.clone())
    }

    /// Build complete key sequence starting from key2_13
    fn build_key_sequence(&self, key2_13: u32, plaintext: &[u8], ciphertext: &[u8],
                         offset: usize, text_length: usize) -> Result<Option<Stage2Result>> {
        // Start recursion from position 13
        if let Some((key0_7, key1_12, key2_13)) = self.recursion2(key2_13, plaintext, ciphertext, offset, 13, text_length)? {
            // Validate the complete key sequence
            if self.validate_key_sequence(key0_7, key1_12, key2_13, plaintext, ciphertext, offset, text_length) {
                return Ok(Some(Stage2Result {
                    key0: key0_7,
                    key1: key1_12,
                    key2: key2_13,
                    offset,
                }));
            }
        }

        Ok(None)
    }

    /// Recursive key reconstruction (implements recursion2 from C code)
    fn recursion2(&self, key2_i: u32, plaintext: &[u8], ciphertext: &[u8],
                  offset: usize, i: usize, text_length: usize) -> Result<Option<(u32, u32, u32)>> {
        if i < offset + 13 {
            // Base case: reached the point where we need to reconstruct key0 and key1
            return self.build_key0_key1(key2_i, plaintext, ciphertext, offset, text_length);
        }

        let key3_i = plaintext[i] ^ ciphertext[i];

        // Find compatible key2_{i-1} values
        let key2_im1_candidates = self.find_key2_im1_candidates(key2_i, key3_i)?;

        for &key2_im1 in &key2_im1_candidates {
            // Recurse to previous position
            if let Some((key0_7, key1_12, key2_13)) = self.recursion2(key2_im1, plaintext, ciphertext, offset, i - 1, text_length)? {
                // Validate that this key2_i is compatible with the sequence
                if self.validate_key2_transition(key2_im1, key2_i, key3_i) {
                    return Ok(Some((key0_7, key1_12, key2_13)));
                }
            }
        }

        Ok(None)
    }

    /// Find candidates for key2_{i-1} given key2_i and key3_i
    fn find_key2_im1_candidates(&self, key2_i: u32, key3_i: u8) -> Result<Vec<u32>> {
        let mut candidates = Vec::new();

        // Apply inverse CRC operation
        let key2_im1_base = inv_crc32_fast(key2_i, 0);

        // Generate candidates based on temp table constraints
        for temp_lo in 0..=0xFFFFu32 {
            let temp = temp_lo | 3;
            let calculated_key3 = ((temp as u64 * (temp as u64 ^ 1)) >> 8) & 0xFF;

            if calculated_key3 == key3_i as u64 {
                let temp_hi = ((temp as u64 * (temp as u64 ^ 1)) >> 32) as u32;
                let candidate = (key2_im1_base & 0xFF000000) | (temp_hi << 16) | temp_lo;

                // Additional constraint checking
                if (candidate & 0xFC000000) == (key2_im1_base & 0xFC000000) {
                    candidates.push(candidate);
                }
            }
        }

        // Limit number of candidates to prevent explosion
        candidates.truncate(1000);
        Ok(candidates)
    }

    /// Build key0_7 and key1_12 from key2_13
    fn build_key0_key1(&self, key2_13: u32, plaintext: &[u8], ciphertext: &[u8],
                      offset: usize, text_length: usize) -> Result<Option<(u32, u32, u32)>> {
        // Find possible key1_12 values
        let key1_12_candidates = self.find_key1_12_candidates(key2_13, plaintext, ciphertext, offset)?;

        for &key1_12 in &key1_12_candidates {
            // Find corresponding key0_7 values
            if let Some(key0_7) = self.find_key0_7_candidate(key1_12, plaintext, ciphertext, offset)? {
                // Validate this key sequence
                if self.validate_key_sequence(key0_7, key1_12, key2_13, plaintext, ciphertext, offset, text_length) {
                    return Ok(Some((key0_7, key1_12, key2_13)));
                }
            }
        }

        Ok(None)
    }

    /// Find candidates for key1_12 given key2_13
    fn find_key1_12_candidates(&self, key2_13: u32, plaintext: &[u8], ciphertext: &[u8],
                              offset: usize) -> Result<Vec<u32>> {
        let mut candidates = Vec::new();

        // key1_12 is related to key2_13 through the key update process
        // We need to work backwards from key2_13
        let key3_12 = plaintext[offset + 12] ^ ciphertext[offset + 12];

        // Try various key1_12 values that could produce key2_13
        for key1_12 in 0u32..=0xFFFFFFFFu32 {
            // Check if this key1_12 could produce the right key2_13
            let temp = (key2_13 & 0xFFFF) | 3;
            let calculated_key3 = ((temp as u64 * (temp as u64 ^ 1)) >> 8) & 0xFF;

            if calculated_key3 == key3_12 as u64 {
                // Additional constraint: key1_12's MSB should be related to key2_13
                let key1_msb = msb(key1_12);
                let key2_candidate = crc32_fast(key2_13, key1_msb);

                // This is a heuristic - may need refinement
                if (key2_candidate & 0xFF000000) == (key2_13 & 0xFF000000) {
                    candidates.push(key1_12);

                    // Limit candidates to prevent explosion
                    if candidates.len() >= 100 {
                        break;
                    }
                }
            }
        }

        Ok(candidates)
    }

    /// Find key0_7 candidate given key1_12
    fn find_key0_7_candidate(&self, key1_12: u32, plaintext: &[u8], ciphertext: &[u8],
                            offset: usize) -> Result<Option<u32>> {
        // key0_7 is related to key1_12 through the key update process
        let key3_7 = plaintext[offset + 7] ^ ciphertext[offset + 7];

        // Work backwards from key1_12
        let mut temp = key1_12.wrapping_sub(1);
        temp = ((temp as u64).wrapping_mul(INVCONST as u64) & 0xFFFFFFFF) as u32;

        // temp should equal key1_7 + (key0_7 & 0xFF)
        let key1_7 = temp.wrapping_sub(key1_12 & 0xFF);

        // Now work backwards further to find key0_7
        for key0_7 in 0u32..=0xFFFFFFFFu32 {
            // Check if this key0_7 produces the right key1_12
            let mut test_key0 = key0_7;
            let mut test_key1 = key1_7;

            // Simulate forward updates from position 7 to 12
            for i in offset + 7..offset + 12 {
                test_key1 = (test_key1 + (test_key0 & 0xFF)).wrapping_mul(CONST).wrapping_add(1);
                test_key0 = crc32_fast(test_key0, plaintext[i]);

                let temp = (test_key1 >> 24) as u8;
                let calculated_key3 = plaintext[i] ^ ciphertext[i];

                // Validate key3 at each step
                let test_key2_temp = (temp as u32 & 0xFFFF) | 3;
                let test_calculated_key3 = ((test_key2_temp as u64 * (test_key2_temp as u64 ^ 1)) >> 8) & 0xFF;

                if test_calculated_key3 != calculated_key3 as u64 {
                    break;
                }
            }

            // If we made it to position 12, check if we got the right key1_12
            if test_key1 == key1_12 {
                return Ok(Some(key0_7));
            }
        }

        Ok(None)
    }

    /// Validate that key2_i can transition to key2_{i+1} with given key3_i
    fn validate_key2_transition(&self, key2_i: u32, key2_ip1: u32, key3_i: u8) -> bool {
        // Check if key2_{i+1} can be derived from key2_i with the given key3_i
        let temp = (key2_ip1 & 0xFFFF) | 3;
        let calculated_key3 = ((temp as u64 * (temp as u64 ^ 1)) >> 8) & 0xFF;

        calculated_key3 == key3_i as u64
    }

    /// Validate complete key sequence against plaintext/ciphertext
    fn validate_key_sequence(&self, key0_7: u32, key1_12: u32, key2_13: u32,
                           plaintext: &[u8], ciphertext: &[u8], offset: usize, text_length: usize) -> bool {
        // Create key state and validate
        let key_manager = match get_key_manager() {
            Ok(km) => km,
            Err(_) => return false,
        };

        // Reconstruct full state
        let mut state = match self.reconstruct_full_state(key0_7, key1_12, key2_13, plaintext, offset) {
            Ok(s) => s,
            Err(_) => return false,
        };

        // Validate by decrypting ciphertext and comparing with plaintext
        key_manager.validate_keys(state, plaintext, ciphertext)
    }

    /// Reconstruct full key state from partial information
    fn reconstruct_full_state(&self, key0_7: u32, key1_12: u32, key2_13: u32,
                            plaintext: &[u8], offset: usize) -> Result<KeyState> {
        let mut key0 = key0_7;
        let mut key1 = key1_12;
        let mut key2 = key2_13;

        // Simulate forward updates to reach the end
        for i in offset + 13..plaintext.len() {
            // Update key2 first
            let key3 = plaintext[i] ^ plaintext[i]; // This is wrong - need ciphertext
            // Actually, we need to use the ciphertext here, but we don't have it in this function
            // This is a placeholder - the real implementation needs access to ciphertext

            key0 = crc32_fast(key0, plaintext[i]);
            key1 = (key1 + (key0 & 0xFF)).wrapping_mul(CONST).wrapping_add(1);
            key2 = crc32_fast(key2, msb(key1));
        }

        Ok(KeyState::new(key0, key1, key2))
    }

    /// Get all found solutions
    pub fn get_solutions(&self) -> &[Stage2Result] {
        &self.solutions
    }

    /// Get number of solutions found
    pub fn get_num_solutions(&self) -> usize {
        self.solutions.len()
    }

    /// Check if any solutions were found
    pub fn has_solutions(&self) -> bool {
        self.found_solutions
    }

    /// Reset the processor state
    pub fn reset(&mut self) {
        self.solutions.clear();
        self.found_solutions = false;
    }
}

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

/// Convenience function to execute Stage 2
pub fn execute_stage2(key2_candidates: &[u32], plaintext: &[u8], ciphertext: &[u8],
                     offset: usize, abort_on_success: bool) -> Result<Vec<Stage2Result>> {
    let mut processor = Stage2Processor::new()?;
    processor.execute(key2_candidates, plaintext, ciphertext, offset, abort_on_success)
}

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

    #[test]
    fn test_stage2_processor_creation() {
        let processor = Stage2Processor::new();
        assert!(processor.is_ok());
    }

    #[test]
    fn test_find_key2_im1_candidates() {
        let processor = Stage2Processor::new().unwrap();
        let key2_i = 0x12345678;
        let key3_i = 0xAB;

        let result = processor.find_key2_im1_candidates(key2_i, key3_i);
        assert!(result.is_ok());
        // May return empty candidates for random values
    }

    #[test]
    fn test_validate_key2_transition() {
        let processor = Stage2Processor::new().unwrap();
        let key2_i = 0x12345678;
        let key2_ip1 = 0x87654321;
        let key3_i = 0xAB;

        // Should not panic
        let _result = processor.validate_key2_transition(key2_i, key2_ip1, key3_i);
    }

    #[test]
    fn test_execute_stage2_empty_candidates() {
        let candidates = vec![];
        let plaintext = b"Test plaintext data";
        let ciphertext = b"Encrypted test data";

        let result = execute_stage2(&candidates, plaintext, ciphertext, 12, false);
        assert!(result.is_ok());
        assert!(result.unwrap().is_empty());
    }
}