kaccy-core 0.2.0

//! Post-Quantum Cryptography Framework
//!
//! This module provides abstractions for post-quantum cryptographic algorithms
//! to prepare for quantum computing threats. Includes:
//! - Lattice-based cryptography (framework)
//! - Hash-based signatures
//! - Code-based cryptography (framework)
//! - Migration strategies from classical to quantum-resistant algorithms

use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::fmt;

/// Post-quantum cryptographic algorithm type
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum PQCAlgorithm {
    /// Lattice-based (e.g., CRYSTALS-Kyber, CRYSTALS-Dilithium)
    Lattice,
    /// Hash-based signatures (e.g., SPHINCS+)
    HashBased,
    /// Code-based (e.g., Classic McEliece)
    CodeBased,
    /// Multivariate (e.g., Rainbow)
    Multivariate,
}

impl fmt::Display for PQCAlgorithm {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            PQCAlgorithm::Lattice => write!(f, "Lattice-based"),
            PQCAlgorithm::HashBased => write!(f, "Hash-based"),
            PQCAlgorithm::CodeBased => write!(f, "Code-based"),
            PQCAlgorithm::Multivariate => write!(f, "Multivariate"),
        }
    }
}

/// Post-quantum key pair (abstraction)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PQKeyPair {
    /// Algorithm used to generate this key pair
    pub algorithm: PQCAlgorithm,
    /// Public key bytes
    pub public_key: Vec<u8>,
    /// Private key bytes
    pub private_key: Vec<u8>,
    /// Total key size in bytes
    pub key_size: usize,
}

impl PQKeyPair {
    /// Create a new PQ key pair (framework - would use actual PQC library)
    pub fn generate(algorithm: PQCAlgorithm) -> Result<Self> {
        let key_size = match algorithm {
            PQCAlgorithm::Lattice => 1568,      // CRYSTALS-Kyber-768 approx
            PQCAlgorithm::HashBased => 64,      // SPHINCS+-256
            PQCAlgorithm::CodeBased => 261120,  // Classic McEliece
            PQCAlgorithm::Multivariate => 1312, // Rainbow
        };

        // NOTE: In production, this would call actual PQC library
        // For now, we create placeholder keys
        Ok(Self {
            algorithm,
            public_key: vec![0u8; key_size / 2],
            private_key: vec![0u8; key_size / 2],
            key_size,
        })
    }

    /// Sign data (framework)
    pub fn sign(&self, _data: &[u8]) -> Result<Vec<u8>> {
        // NOTE: Would use actual PQC signature algorithm
        // Placeholder implementation
        let signature_size = match self.algorithm {
            PQCAlgorithm::Lattice => 2420,    // Dilithium3
            PQCAlgorithm::HashBased => 8080,  // SPHINCS+
            PQCAlgorithm::CodeBased => 32,    // Simplified
            PQCAlgorithm::Multivariate => 66, // Rainbow
        };

        Ok(vec![0u8; signature_size])
    }

    /// Verify signature (framework)
    pub fn verify(&self, _data: &[u8], signature: &[u8]) -> Result<bool> {
        // NOTE: Would use actual PQC verification
        // Placeholder: just check signature length
        Ok(!signature.is_empty())
    }
}

/// Hybrid cryptography combining classical and post-quantum
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HybridCrypto {
    /// Classical key (e.g., RSA, ECC)
    pub classical_public_key: Vec<u8>,
    /// Post-quantum key
    pub pq_key: PQKeyPair,
    /// Hybrid mode
    pub mode: HybridMode,
}

/// Hybrid mode for gradual migration
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum HybridMode {
    /// Prefer classical, fallback to PQ
    ClassicalPrimary,
    /// Use both (AND security)
    DualSignature,
    /// Prefer PQ, fallback to classical
    PQPrimary,
    /// PQ only (full migration)
    PQOnly,
}

impl HybridCrypto {
    /// Create a new hybrid cryptography configuration
    pub fn new(pq_algorithm: PQCAlgorithm, mode: HybridMode) -> Result<Self> {
        Ok(Self {
            classical_public_key: vec![0u8; 256], // Placeholder RSA-2048 key
            pq_key: PQKeyPair::generate(pq_algorithm)?,
            mode,
        })
    }

    /// Sign with hybrid approach
    pub fn hybrid_sign(&self, data: &[u8]) -> Result<HybridSignature> {
        match self.mode {
            HybridMode::ClassicalPrimary | HybridMode::PQPrimary => {
                // Single signature (would choose based on mode)
                Ok(HybridSignature {
                    classical_sig: Some(vec![0u8; 256]),
                    pq_sig: None,
                    mode: self.mode,
                })
            }
            HybridMode::DualSignature => {
                // Both signatures required
                Ok(HybridSignature {
                    classical_sig: Some(vec![0u8; 256]),
                    pq_sig: Some(self.pq_key.sign(data)?),
                    mode: self.mode,
                })
            }
            HybridMode::PQOnly => {
                // Only PQ signature
                Ok(HybridSignature {
                    classical_sig: None,
                    pq_sig: Some(self.pq_key.sign(data)?),
                    mode: self.mode,
                })
            }
        }
    }
}

/// Hybrid signature containing both classical and PQ components
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HybridSignature {
    /// Classical (e.g. ECDSA) signature bytes, if present
    pub classical_sig: Option<Vec<u8>>,
    /// Post-quantum signature bytes, if present
    pub pq_sig: Option<Vec<u8>>,
    /// Hybrid mode that generated this signature
    pub mode: HybridMode,
}

/// Migration plan for transitioning to PQC
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PQMigrationPlan {
    /// Current migration phase
    pub current_phase: MigrationPhase,
    /// Target PQC algorithm to migrate to
    pub target_algorithm: PQCAlgorithm,
    /// Total planned timeline in days
    pub timeline_days: u32,
    /// Whether the migration maintains backward compatibility
    pub backwards_compatible: bool,
}

/// Migration phases
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum MigrationPhase {
    /// Preparation and testing
    Phase1Assessment,
    /// Deploy hybrid mode
    Phase2HybridDeploy,
    /// Gradual rollout
    Phase3GradualRollout,
    /// Full PQ deployment
    Phase4FullPQ,
    /// Deprecate classical
    Phase5DeprecateClassical,
}

impl PQMigrationPlan {
    /// Create a new migration plan starting at Phase 1 assessment
    pub fn new(target_algorithm: PQCAlgorithm, timeline_days: u32) -> Self {
        Self {
            current_phase: MigrationPhase::Phase1Assessment,
            target_algorithm,
            timeline_days,
            backwards_compatible: true,
        }
    }

    /// Advance to the next migration phase and return it
    pub fn advance_phase(&mut self) -> MigrationPhase {
        self.current_phase = match self.current_phase {
            MigrationPhase::Phase1Assessment => MigrationPhase::Phase2HybridDeploy,
            MigrationPhase::Phase2HybridDeploy => MigrationPhase::Phase3GradualRollout,
            MigrationPhase::Phase3GradualRollout => MigrationPhase::Phase4FullPQ,
            MigrationPhase::Phase4FullPQ => MigrationPhase::Phase5DeprecateClassical,
            MigrationPhase::Phase5DeprecateClassical => MigrationPhase::Phase5DeprecateClassical,
        };
        self.current_phase
    }

    /// Get the recommended hybrid mode for the current migration phase
    pub fn get_recommended_hybrid_mode(&self) -> HybridMode {
        match self.current_phase {
            MigrationPhase::Phase1Assessment => HybridMode::ClassicalPrimary,
            MigrationPhase::Phase2HybridDeploy => HybridMode::DualSignature,
            MigrationPhase::Phase3GradualRollout => HybridMode::PQPrimary,
            MigrationPhase::Phase4FullPQ | MigrationPhase::Phase5DeprecateClassical => {
                HybridMode::PQOnly
            }
        }
    }
}

/// Quantum threat assessment
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QuantumThreatAssessment {
    /// Overall quantum threat level
    pub threat_level: ThreatLevel,
    /// Estimated years until quantum computers can break current cryptography
    pub estimated_years_until_threat: u8,
    /// Names of assets that need quantum-safe protection
    pub critical_assets: Vec<String>,
    /// Recommended PQC algorithm based on threat assessment
    pub recommended_algorithm: PQCAlgorithm,
}

/// Threat level from quantum computing
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ThreatLevel {
    /// Quantum threat is 10+ years away
    Low,
    /// Quantum threat is 5–10 years away
    Medium,
    /// Quantum threat is 2–5 years away
    High,
    /// Quantum threat is less than 2 years away
    Critical,
}

impl QuantumThreatAssessment {
    /// Build a threat assessment from estimated timeline and secret longevity
    pub fn assess(years_until_threat: u8, has_long_term_secrets: bool) -> Self {
        let threat_level = match years_until_threat {
            0..=2 => ThreatLevel::Critical,
            3..=5 => ThreatLevel::High,
            6..=10 => ThreatLevel::Medium,
            _ => ThreatLevel::Low,
        };

        let recommended_algorithm = if has_long_term_secrets {
            // Prefer lattice-based for key exchange
            PQCAlgorithm::Lattice
        } else {
            // Hash-based for signatures
            PQCAlgorithm::HashBased
        };

        Self {
            threat_level,
            estimated_years_until_threat: years_until_threat,
            critical_assets: Vec::new(),
            recommended_algorithm,
        }
    }

    /// Add an asset name to the list of critical assets
    pub fn add_critical_asset(&mut self, asset: String) {
        self.critical_assets.push(asset);
    }
}

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

    #[test]
    fn test_pq_keypair_generation() {
        let keypair = PQKeyPair::generate(PQCAlgorithm::Lattice).unwrap();
        assert_eq!(keypair.algorithm, PQCAlgorithm::Lattice);
        assert!(!keypair.public_key.is_empty());
        assert!(!keypair.private_key.is_empty());
    }

    #[test]
    fn test_pq_signature() {
        let keypair = PQKeyPair::generate(PQCAlgorithm::HashBased).unwrap();
        let data = b"test message";

        let signature = keypair.sign(data).unwrap();
        assert!(!signature.is_empty());

        let valid = keypair.verify(data, &signature).unwrap();
        assert!(valid);
    }

    #[test]
    fn test_hybrid_crypto_dual_signature() {
        let hybrid = HybridCrypto::new(PQCAlgorithm::Lattice, HybridMode::DualSignature).unwrap();
        let data = b"important message";

        let sig = hybrid.hybrid_sign(data).unwrap();
        assert!(sig.classical_sig.is_some());
        assert!(sig.pq_sig.is_some());
    }

    #[test]
    fn test_hybrid_crypto_pq_only() {
        let hybrid = HybridCrypto::new(PQCAlgorithm::Lattice, HybridMode::PQOnly).unwrap();
        let data = b"pq only message";

        let sig = hybrid.hybrid_sign(data).unwrap();
        assert!(sig.classical_sig.is_none());
        assert!(sig.pq_sig.is_some());
    }

    #[test]
    fn test_migration_plan() {
        let mut plan = PQMigrationPlan::new(PQCAlgorithm::Lattice, 180);

        assert_eq!(plan.current_phase, MigrationPhase::Phase1Assessment);
        assert_eq!(
            plan.get_recommended_hybrid_mode(),
            HybridMode::ClassicalPrimary
        );

        plan.advance_phase();
        assert_eq!(plan.current_phase, MigrationPhase::Phase2HybridDeploy);
        assert_eq!(
            plan.get_recommended_hybrid_mode(),
            HybridMode::DualSignature
        );

        plan.advance_phase();
        assert_eq!(plan.current_phase, MigrationPhase::Phase3GradualRollout);

        plan.advance_phase();
        assert_eq!(plan.current_phase, MigrationPhase::Phase4FullPQ);
        assert_eq!(plan.get_recommended_hybrid_mode(), HybridMode::PQOnly);
    }

    #[test]
    fn test_quantum_threat_assessment() {
        let assessment = QuantumThreatAssessment::assess(3, true);
        assert_eq!(assessment.threat_level, ThreatLevel::High);
        assert_eq!(assessment.recommended_algorithm, PQCAlgorithm::Lattice);
    }

    #[test]
    fn test_threat_levels() {
        assert_eq!(
            QuantumThreatAssessment::assess(1, false).threat_level,
            ThreatLevel::Critical
        );
        assert_eq!(
            QuantumThreatAssessment::assess(7, false).threat_level,
            ThreatLevel::Medium
        );
        assert_eq!(
            QuantumThreatAssessment::assess(15, false).threat_level,
            ThreatLevel::Low
        );
    }

    #[test]
    fn test_critical_assets() {
        let mut assessment = QuantumThreatAssessment::assess(5, true);
        assessment.add_critical_asset("User private keys".to_string());
        assessment.add_critical_asset("Transaction signatures".to_string());

        assert_eq!(assessment.critical_assets.len(), 2);
    }

    #[test]
    fn test_algorithm_display() {
        assert_eq!(format!("{}", PQCAlgorithm::Lattice), "Lattice-based");
        assert_eq!(format!("{}", PQCAlgorithm::HashBased), "Hash-based");
    }
}