ruption_quantum_encrypt 0.1.1

use getrandom::getrandom;
use sha3::{Digest, Sha3_256, Shake256};
use sha3::digest::{Update, ExtendableOutput};
use std::io::Read;
use std::thread;
use std::time::{Instant, Duration};

// Toy parameters for demonstration; scale up for real security
const LATTICE_DIM: usize = 256;  // Lattice dimension (use 2048+ for billion-qubit resistance)
const CODE_LENGTH: usize = 512;  // Code length (use 8192+ for extreme security)

/// Public key for the quantum-secure encryption scheme.
///
/// Combines lattice-based and code-based cryptographic components.
/// This is a simplified representation; production use requires proper LWE and McEliece implementations.
pub struct PublicKey {
    _lattice_matrix: Vec<Vec<u8>>, // Simplified lattice public key (matrix), unused in toy version
    _code_generator: Vec<u8>,      // Simplified code-based generator, unused in toy version
}

/// Secret key for the quantum-secure encryption scheme.
///
/// Contains private data for lattice and code-based decryption.
/// This is a toy version; scale parameters for real-world security.
pub struct SecretKey {
    _lattice_secret: Vec<u8>,      // Lattice private key, unused in toy version
    _code_secret: Vec<u8>,         // Code private key, unused in toy version
}

/// Ciphertext produced during key encapsulation.
///
/// Holds encrypted data from both lattice and code-based components.
/// In practice, this would result from proper cryptographic operations.
pub struct Ciphertext {
    lattice_cipher: Vec<u8>,      // Lattice-based ciphertext
    code_cipher: Vec<u8>,         // Code-based ciphertext
}

/// Shared secret derived during encapsulation.
///
/// Used for symmetric encryption or key derivation, wrapped for type safety.
pub struct SharedSecret(Vec<u8>);

impl SharedSecret {
    /// Returns a reference to the shared secret bytes.
    ///
    /// Useful for comparing or using the secret in encryption.
    pub fn as_bytes(&self) -> &[u8] {
        &self.0
    }
}

/// A generator for high-quality randomness approximating true entropy.
///
/// Combines OS-provided entropy, timing jitter, and a quantum-inspired simulation.
/// Suitable for cryptographic key generation in a software-only environment.
///
/// # Examples
/// ```
/// use ruption_quantum_encrypt::TrueRandom;
/// let mut trng = TrueRandom::new();
/// let random_bytes = trng.generate(32);
/// assert_eq!(random_bytes.len(), 32);
/// ```
pub struct TrueRandom {
    entropy_pool: Vec<u8>,        // Pool of collected entropy
}

impl TrueRandom {
    /// Initializes a new randomness generator with system entropy.
    ///
    /// Seeds the entropy pool with 64 bytes from the OS's secure random source.
    /// Panics if entropy retrieval fails (rare on modern systems).
    pub fn new() -> Self {
        let mut initial_entropy = vec![0u8; 64];
        getrandom(&mut initial_entropy).expect("Failed to get system entropy");
        TrueRandom {
            entropy_pool: initial_entropy,
        }
    }

    /// Collects timing jitter from thread scheduling to enhance entropy.
    fn collect_jitter(&mut self) {
        let mut jitter = Vec::new();
        for _ in 0..10 {
            let start = Instant::now();
            thread::sleep(Duration::from_nanos(1));
            let elapsed = start.elapsed().as_nanos() as u8;
            jitter.push(elapsed);
        }
        self.entropy_pool.extend(jitter);
    }

    /// Simulates a quantum-inspired entropy source using system timing.
    ///
    /// Approximates unpredictable behavior in software; not true quantum randomness.
    fn quantum_sim_entropy(&mut self) -> Vec<u8> {
        let mut sim_entropy = Vec::new();
        let now = Instant::now().elapsed().as_nanos();
        let mut state = now as u64;

        for _ in 0..16 {
            state ^= state.wrapping_add(self.entropy_pool[state as usize % self.entropy_pool.len()] as u64);
            sim_entropy.push((state & 0xFF) as u8);
        }
        sim_entropy
    }

    /// Generates random bytes of the specified length.
    ///
    /// Mixes OS entropy, jitter, and simulated quantum entropy with SHA-3 for uniformity.
    ///
    /// # Arguments
    /// * `len` - The number of bytes to generate.
    ///
    /// # Returns
    /// A `Vec<u8>` of random bytes.
    pub fn generate(&mut self, len: usize) -> Vec<u8> {
        self.collect_jitter();
        let sim_entropy = self.quantum_sim_entropy();
        self.entropy_pool.extend(sim_entropy);

        let mut hasher = Sha3_256::new();
        Update::update(&mut hasher, &self.entropy_pool);
        let mixed = hasher.finalize();

        if len > mixed.len() {
            let mut xof = Shake256::default();
            xof.update(&mixed);
            let mut reader = xof.finalize_xof();
            let mut output = vec![0u8; len];
            reader.read_exact(&mut output).unwrap();
            output
        } else {
            mixed[..len].to_vec()
        }
    }
}

impl Default for TrueRandom {
    /// Provides a default instance of `TrueRandom`.
    ///
    /// Delegates to `new()` to initialize with system entropy.
    fn default() -> Self {
        Self::new()
    }
}

/// Generates a keypair for quantum-secure encryption.
///
/// Uses `TrueRandom` to produce unpredictable keys.
/// This is a simplified version; real-world use requires proper cryptographic math.
///
/// # Returns
/// A tuple `(PublicKey, SecretKey)` for use in encryption/decryption.
///
/// # Examples
/// ```
/// use ruption_quantum_encrypt::keypair;
/// let (pk, sk) = keypair();
/// ```
pub fn keypair() -> (PublicKey, SecretKey) {
    let mut trng = TrueRandom::new();

    let lattice_secret = trng.generate(LATTICE_DIM);
    let lattice_matrix = vec![trng.generate(LATTICE_DIM); LATTICE_DIM];

    let code_secret = trng.generate(CODE_LENGTH / 8);
    let code_generator = trng.generate(CODE_LENGTH);

    (
        PublicKey {
            _lattice_matrix: lattice_matrix,
            _code_generator: code_generator,
        },
        SecretKey {
            _lattice_secret: lattice_secret,
            _code_secret: code_secret,
        },
    )
}

/// Encapsulates a shared secret using the public key.
///
/// Produces a ciphertext and shared secret for secure key exchange.
/// This is a toy implementation; replace with real algorithms for production.
///
/// # Arguments
/// * `pk` - The recipient’s `PublicKey`.
///
/// # Returns
/// A tuple `(Ciphertext, SharedSecret)` with the encrypted data and secret.
///
/// # Examples
/// ```
/// use ruption_quantum_encrypt::{keypair, encapsulate};
/// let (pk, _sk) = keypair();
/// let (ct, ss) = encapsulate(&pk);
/// ```
pub fn encapsulate(_pk: &PublicKey) -> (Ciphertext, SharedSecret) {
    let mut trng = TrueRandom::new();

    let lattice_cipher = trng.generate(LATTICE_DIM);
    let code_cipher = trng.generate(CODE_LENGTH);

    let mut hasher = Sha3_256::new();
    Update::update(&mut hasher, &lattice_cipher);
    Update::update(&mut hasher, &code_cipher);
    let shared_secret = SharedSecret(hasher.finalize().to_vec());

    (
        Ciphertext {
            lattice_cipher,
            code_cipher,
        },
        shared_secret,
    )
}

/// Decapsulates the ciphertext to retrieve the shared secret.
///
/// Uses the secret key to recover the shared secret.
/// Simplified for demonstration; real decryption would use the secret key.
///
/// # Arguments
/// * `ct` - The `Ciphertext` to decapsulate.
/// * `sk` - The `SecretKey` for decryption.
///
/// # Returns
/// The `SharedSecret` derived from the ciphertext.
///
/// # Examples
/// ```
/// use ruption_quantum_encrypt::{keypair, encapsulate, decapsulate};
/// let (pk, sk) = keypair();
/// let (ct, ss1) = encapsulate(&pk);
/// let ss2 = decapsulate(&ct, &sk);
/// assert_eq!(ss1.as_bytes(), ss2.as_bytes());
/// ```
pub fn decapsulate(ct: &Ciphertext, _sk: &SecretKey) -> SharedSecret {
    let mut hasher = Sha3_256::new();
    Update::update(&mut hasher, &ct.lattice_cipher);
    Update::update(&mut hasher, &ct.code_cipher);
    SharedSecret(hasher.finalize().to_vec())
}

/// Derives multiple keys from a shared secret.
///
/// Uses SHAKE256 to generate multiple 256-bit keys for layered encryption.
///
/// # Arguments
/// * `shared_secret` - The `SharedSecret` to derive keys from.
/// * `num_keys` - Number of keys to generate.
///
/// # Returns
/// A `Vec<Vec<u8>>` of derived keys.
///
/// # Examples
/// ```
/// use ruption_quantum_encrypt::{keypair, encapsulate, derive_keys};
/// let (pk, _sk) = keypair();
/// let (_, ss) = encapsulate(&pk);
/// let keys = derive_keys(&ss, 3);
/// assert_eq!(keys.len(), 3);
/// assert_eq!(keys[0].len(), 32);
/// ```
pub fn derive_keys(shared_secret: &SharedSecret, num_keys: usize) -> Vec<Vec<u8>> {
    let mut keys = Vec::new();
    let mut xof = Shake256::default();
    xof.update(&shared_secret.0);
    let mut reader = xof.finalize_xof();

    for _ in 0..num_keys {
        let mut key = vec![0u8; 32];
        reader.read_exact(&mut key).unwrap();
        keys.push(key);
    }
    keys
}

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

    #[test]
    fn test_true_random() {
        let mut trng = TrueRandom::new();
        let rand1 = trng.generate(32);
        let rand2 = trng.generate(32);
        assert_eq!(rand1.len(), 32);
        assert_eq!(rand2.len(), 32);
        assert_ne!(rand1, rand2);
    }

    #[test]
    fn test_encryption() {
        let (pk, sk) = keypair();
        let (ct, ss1) = encapsulate(&pk);
        let ss2 = decapsulate(&ct, &sk);
        assert_eq!(ss1.as_bytes(), ss2.as_bytes());

        let keys = derive_keys(&ss1, 3);
        assert_eq!(keys.len(), 3);
        assert_eq!(keys[0].len(), 32);
    }
}