pkcrack 0.1.0

//! Stage 3 of the pkcrack attack: Password search or ZIP decryption

use crate::types::*;
use crate::crypto::{get_key_manager, KeyManager};
use crate::zip::ZipProcessor;
use crate::error::{PkCrackError, Result};
use std::path::PathBuf;

/// Stage 3 processor for password search and decryption
pub struct Stage3Processor {
    key_manager: KeyManager,
    zip_processor: Option<ZipProcessor>,
}

impl Stage3Processor {
    /// Create new Stage 3 processor
    pub fn new() -> Result<Self> {
        Ok(Self {
            key_manager: get_key_manager()?.clone(),
            zip_processor: None,
        })
    }

    /// Create Stage 3 processor with ZIP processing capability
    pub fn with_zip_support() -> Result<Self> {
        Ok(Self {
            key_manager: get_key_manager()?.clone(),
            zip_processor: Some(ZipProcessor::new()?),
        })
    }

    /// Execute Stage 3: find password or decrypt ZIP
    pub fn execute(&self, solution: &Stage2Result, config: &crate::Config) -> Result<AttackResult> {
        // Reconstruct complete key state
        let mut key_state = self.reconstruct_key_state(solution)?;

        if let Some(output_zip) = &config.output_zip {
            // Decrypt ZIP file
            self.decrypt_zip(config, &mut key_state, output_zip.to_str().unwrap())?;
            Ok(AttackResult {
                keys: key_state,
                password: None, // Password not found, but we have the keys
                best_offset: solution.offset,
                num_candidates: 1,
            })
        } else {
            // Search for password
            let password = self.search_password(key_state)?;
            Ok(AttackResult {
                keys: key_state,
                password: Some(password),
                best_offset: solution.offset,
                num_candidates: 1,
            })
        }
    }

    /// Reconstruct complete key state from stage 2 solution
    fn reconstruct_key_state(&self, solution: &Stage2Result) -> Result<KeyState> {
        // Create initial state with the found keys
        let mut state = KeyState {
            key0: solution.key0,
            key1: solution.key1,
            key2: solution.key2,
        };

        Ok(state)
    }

    /// Search for password that generates the given key state
    fn search_password(&self, target_keys: KeyState) -> Result<String> {
        println!("Starting password search...");

        // Brute-force password search
        // This is computationally expensive for long passwords
        let mut password = String::new();

        // Try common password patterns first
        if let Some(found_password) = self.try_common_passwords(target_keys)? {
            return Ok(found_password);
        }

        // Try brute-force up to certain length
        if let Some(found_password) = self.brute_force_search(target_keys, 8)? {
            return Ok(found_password);
        }

        // If no password found, return error
        Err(PkCrackError::NoSolutionFound)
    }

    /// Try common password patterns
    fn try_common_passwords(&self, target_keys: KeyState) -> Result<Option<String>> {
        let common_passwords = vec![
            "password", "123456", "secret", "admin", "test",
            "qwerty", "abc123", "password123", "123456789",
            "welcome", "login", "letmein", "dragon",
        ];

        for password in common_passwords {
            let keys = self.key_manager.keys_from_password(password.as_bytes());

            // Check if keys match (allowing some tolerance for partial matches)
            if self.keys_match(keys, target_keys, 4) { // Allow 4 bits of difference
                println!("Found common password: {}", password);
                return Ok(Some(password.to_string()));
            }
        }

        Ok(None)
    }

    /// Brute-force password search up to specified length
    fn brute_force_search(&self, target_keys: KeyState, max_length: usize) -> Result<Option<String>> {
        // Character set to try
        let charset: Vec<char> = "abcdefghijklmnopqrstuvwxyz0123456789".chars().collect();

        // Start with length 1 passwords and increase
        for length in 1..=max_length {
            println!("Trying passwords of length {}...", length);

            if let Some(password) = self.brute_force_length(&charset, length, target_keys)? {
                return Ok(Some(password));
            }
        }

        Ok(None)
    }

    /// Brute-force passwords of specific length
    fn brute_force_length(&self, charset: &[char], length: usize, target_keys: KeyState) -> Result<Option<String>> {
        if length == 0 {
            return Ok(None);
        }

        let mut password = vec![' '; length];

        // Recursive brute-force
        if self.try_password_recursive(&mut password, 0, charset, target_keys)? {
            let password_str: String = password.into_iter().collect();
            return Ok(Some(password_str));
        }

        Ok(None)
    }

    /// Recursive helper for brute-force password generation
    fn try_password_recursive(&self, password: &mut [char], pos: usize, charset: &[char], target_keys: KeyState) -> Result<bool> {
        if pos >= password.len() {
            // Test this password
            let password_bytes: Vec<u8> = password.iter().map(|&c| c as u8).collect();
            let keys = self.key_manager.keys_from_password(&password_bytes);

            if self.keys_match(keys, target_keys, 0) { // Exact match
                println!("Found password: {}", password.iter().collect::<String>());
                return Ok(true);
            }
            return Ok(false);
        }

        for &c in charset {
            password[pos] = c;
            if self.try_password_recursive(password, pos + 1, charset, target_keys)? {
                return Ok(true);
            }
        }

        Ok(false)
    }

    /// Check if keys match within tolerance
    fn keys_match(&self, keys1: KeyState, keys2: KeyState, tolerance: u32) -> bool {
        let diff0 = (keys1.key0 ^ keys2.key0).count_ones();
        let diff1 = (keys1.key1 ^ keys2.key1).count_ones();
        let diff2 = (keys1.key2 ^ keys2.key2).count_ones();

        (diff0 + diff1 + diff2) <= tolerance
    }

    /// Decrypt ZIP file using found keys
    fn decrypt_zip(&self, config: &crate::Config, key_state: &mut KeyState, output_path: &str) -> Result<()> {
        let zip_processor = self.zip_processor.as_ref()
            .ok_or_else(|| PkCrackError::config("ZIP support not enabled"))?;

        let ciphertext_zip = config.ciphertext_zip.as_ref()
            .ok_or_else(|| PkCrackError::config("No ciphertext ZIP specified"))?;

        let ciphertext_filename = config.ciphertext_filename()
            .ok_or_else(|| PkCrackError::config("No ciphertext filename available"))?;

        // Extract and decrypt the ZIP
        zip_processor.decrypt_zip_to_file(
            ciphertext_zip,
            &ciphertext_filename,
            output_path,
            key_state,
        )?;

        println!("Successfully decrypted ZIP to: {}", output_path);
        Ok(())
    }

    /// Validate password against actual encryption
    pub fn validate_password(&self, password: &str, plaintext: &[u8], ciphertext: &[u8]) -> Result<bool> {
        let keys = self.key_manager.keys_from_password(password.as_bytes());
        Ok(self.key_manager.validate_keys(keys, plaintext, ciphertext))
    }

    /// Get password complexity estimate
    pub fn estimate_password_complexity(&self, password: &str) -> f64 {
        let charset_size = self.estimate_charset_size(password);
        let length = password.len();

        // Log base 2 of charset_size^length
        (length as f64) * (charset_size as f64).log2()
    }

    /// Estimate character set size used in password
    fn estimate_charset_size(&self, password: &str) -> usize {
        let mut has_lower = false;
        let mut has_upper = false;
        let mut has_digit = false;
        let mut has_special = false;

        for c in password.chars() {
            if c.is_lowercase() {
                has_lower = true;
            } else if c.is_uppercase() {
                has_upper = true;
            } else if c.is_ascii_digit() {
                has_digit = true;
            } else {
                has_special = true;
            }
        }

        let mut size = 0;
        if has_lower { size += 26; }
        if has_upper { size += 26; }
        if has_digit { size += 10; }
        if has_special { size += 32; } // Approximate

        size.max(1)
    }

    /// Generate progress report for password search
    pub fn generate_progress_report(&self, current_password: &str, total_tried: u64, est_remaining: Option<u64>) -> String {
        let complexity = self.estimate_password_complexity(current_password);

        match est_remaining {
            Some(remaining) => {
                let progress_pct = (total_tried as f64 / (total_tried + remaining) as f64) * 100.0;
                format!(
                    "Password search progress: {:.2}% | Current: {} | Complexity: {:.1} bits | Tried: {} | Est. remaining: {}",
                    progress_pct, current_password, complexity, total_tried, remaining
                )
            }
            None => {
                format!(
                    "Password search | Current: {} | Complexity: {:.1} bits | Tried: {}",
                    current_password, complexity, total_tried
                )
            }
        }
    }
}

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

/// Convenience function to execute Stage 3
pub fn execute_stage3(solution: &Stage2Result, config: &crate::Config) -> Result<AttackResult> {
    let processor = if config.output_zip.is_some() {
        Stage3Processor::with_zip_support()?
    } else {
        Stage3Processor::new()?
    };

    processor.execute(solution, config)
}

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

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

    #[test]
    fn test_password_validation() {
        let processor = Stage3Processor::new().unwrap();
        let password = "test123";
        let plaintext = b"Hello, World!";

        // Encrypt plaintext with password-derived keys
        let keys = processor.key_manager.keys_from_password(password.as_bytes());
        let mut ciphertext = plaintext.to_vec();
        processor.key_manager.encrypt(&mut KeyState::new(keys.key0, keys.key1, keys.key2), &mut ciphertext);

        // Validate password
        let result = processor.validate_password(password, plaintext, &ciphertext);
        assert!(result.is_ok());
        assert!(result.unwrap());

        // Wrong password should fail
        let wrong_password = "wrongpass";
        let result = processor.validate_password(wrong_password, plaintext, &ciphertext);
        assert!(result.is_ok());
        assert!(!result.unwrap());
    }

    #[test]
    fn test_keys_match() {
        let processor = Stage3Processor::new().unwrap();
        let keys1 = KeyState::new(0x12345678, 0x87654321, 0xABCDEF00);
        let keys2 = KeyState::new(0x12345678, 0x87654321, 0xABCDEF00);
        let keys3 = KeyState::new(0x12345679, 0x87654321, 0xABCDEF00);

        assert!(processor.keys_match(keys1, keys2, 0)); // Exact match
        assert!(processor.keys_match(keys1, keys3, 1)); // Within tolerance
        assert!(!processor.keys_match(keys1, keys3, 0)); // Not within tolerance
    }

    #[test]
    fn test_password_complexity() {
        let processor = Stage3Processor::new().unwrap();

        let simple = "abc";
        let complex = "Abc123!@#";

        let simple_complexity = processor.estimate_password_complexity(simple);
        let complex_complexity = processor.estimate_password_complexity(complex);

        assert!(simple_complexity < complex_complexity);
        assert!(simple_complexity > 0.0);
        assert!(complex_complexity > 0.0);
    }

    #[test]
    fn test_generate_progress_report() {
        let processor = Stage3Processor::new().unwrap();

        let report = processor.generate_progress_report("test123", 1000, Some(9000));
        assert!(report.contains("Password search"));
        assert!(report.contains("test123"));
        assert!(report.contains("10.00%")); // 1000/(1000+9000) * 100

        let report2 = processor.generate_progress_report("abc", 500, None);
        assert!(report2.contains("Password search"));
        assert!(report2.contains("abc"));
    }

    #[test]
    fn test_common_passwords() {
        let processor = Stage3Processor::new().unwrap();

        // Create test keys from a common password
        let password = "password";
        let target_keys = processor.key_manager.keys_from_password(password.as_bytes());

        let result = processor.try_common_passwords(target_keys);
        assert!(result.is_ok());

        // Should find the common password
        let found = result.unwrap();
        assert!(found.is_some());
        assert_eq!(found.unwrap(), password);
    }

    #[test]
    fn test_brute_force_short() {
        let processor = Stage3Processor::new().unwrap();

        // Create test keys from a short password
        let password = "abc";
        let target_keys = processor.key_manager.keys_from_password(password.as_bytes());

        let result = processor.brute_force_search(target_keys, 3);
        assert!(result.is_ok());

        // Should find the short password
        let found = result.unwrap();
        assert!(found.is_some());
        assert_eq!(found.unwrap(), password);
    }
}