//! Encrypt stdin → stdout under a fresh random key, for randomness testing.
//!
//! Usage: `cipher_encrypt <cipher>` — reads plaintext from stdin, writes
//! ciphertext to stdout. Block ciphers run in CTR mode with a random IV
//! drawn from the OS RNG so the output should look uniform regardless of
//! plaintext structure. Stream ciphers run in their native keystream mode.
//!
//! The R driver pipes the works of Shakespeare in and feeds the resulting
//! ciphertext to its statistical battery.

use std::io::{Read, Write};

use cryptography::chacha20::{ChaCha20, XChaCha20};
use cryptography::rabbit::Rabbit;
use cryptography::salsa20::Salsa20;
use cryptography::snow3g::Snow3g;
use cryptography::zuc::Zuc128;
use cryptography::BlockCipher;
use cryptography::{
    Aes128, Aes192, Aes256, Camellia128, Camellia192, Camellia256, Cast128, Des, Grasshopper,
    Magma, Present128, Present80, Seed, Serpent128, Serpent192, Serpent256, Simon128_128,
    Simon128_256, Simon32_64, Simon64_128, Sm4, Speck128_128, Speck128_256, Speck32_64,
    Speck64_128, TripleDes, Twofish128, Twofish256,
};

fn os_random(out: &mut [u8]) {
    let mut f = std::fs::File::open("/dev/urandom").expect("open /dev/urandom");
    f.read_exact(out).expect("read /dev/urandom");
}

/// Generic CTR-mode encryption: counter is a big-endian integer initialised
/// from a random IV; it is incremented once per block.
fn ctr_encrypt<C: BlockCipher>(cipher: &C, iv: &[u8], data: &mut [u8]) {
    assert_eq!(iv.len(), C::BLOCK_LEN);
    let n = C::BLOCK_LEN;
    let mut counter = vec![0u8; n];
    counter.copy_from_slice(iv);
    let mut keystream = vec![0u8; n];

    for chunk in data.chunks_mut(n) {
        keystream.copy_from_slice(&counter);
        cipher.encrypt(&mut keystream);
        for (b, k) in chunk.iter_mut().zip(keystream.iter()) {
            *b ^= *k;
        }
        // Increment counter as a big-endian integer.
        for byte in counter.iter_mut().rev() {
            let (next, carry) = byte.overflowing_add(1);
            *byte = next;
            if !carry {
                break;
            }
        }
    }
}

fn run_block<C: BlockCipher>(cipher: C, data: &mut [u8]) {
    let mut iv = vec![0u8; C::BLOCK_LEN];
    os_random(&mut iv);
    ctr_encrypt(&cipher, &iv, data);
}

fn random_array<const N: usize>() -> [u8; N] {
    let mut out = [0u8; N];
    os_random(&mut out);
    out
}

const CIPHERS: &[&str] = &[
    "aes128",
    "aes192",
    "aes256",
    "camellia128",
    "camellia192",
    "camellia256",
    "cast128",
    "des",
    "3des",
    "grasshopper",
    "magma",
    "present80",
    "present128",
    "seed",
    "serpent128",
    "serpent192",
    "serpent256",
    "sm4",
    "twofish128",
    "twofish256",
    "simon32_64",
    "simon64_128",
    "simon128_128",
    "simon128_256",
    "speck32_64",
    "speck64_128",
    "speck128_128",
    "speck128_256",
    "chacha20",
    "xchacha20",
    "salsa20",
    "rabbit",
    "zuc128",
    "snow3g",
];

fn print_help() {
    println!(
        "cipher_encrypt — randomness-test helper for `scripts/cipher_randomness.R`.\n\n\
         Usage: cipher_encrypt <cipher>\n\
                cipher_encrypt --help | -h\n\
                cipher_encrypt --list\n\n\
         Reads plaintext from stdin and writes the ciphertext under a fresh\n\
         OS-random key to stdout.  Block ciphers run in CTR mode with a random\n\
         IV; stream ciphers run in their native keystream mode.  The IV is not\n\
         emitted because the consumer is the R battery, not a real protocol.\n\n\
         This binary is for statistical testing only — it is not a general\n\
         encryption CLI.  Run with `--list` to see the supported cipher names."
    );
}

fn print_list() {
    for c in CIPHERS {
        println!("{c}");
    }
}

fn main() {
    let name = std::env::args().nth(1).unwrap_or_else(|| {
        eprintln!("usage: cipher_encrypt <cipher>  (try `--help`)");
        std::process::exit(2);
    });

    match name.as_str() {
        "-h" | "--help" => {
            print_help();
            return;
        }
        "--list" => {
            print_list();
            return;
        }
        _ => {}
    }

    let mut data = Vec::new();
    std::io::stdin().read_to_end(&mut data).expect("read stdin");

    match name.to_ascii_lowercase().as_str() {
        // ── AES ───────────────────────────────────────────────────────────────
        "aes128" => run_block(Aes128::new(&random_array::<16>()), &mut data),
        "aes192" => run_block(Aes192::new(&random_array::<24>()), &mut data),
        "aes256" => run_block(Aes256::new(&random_array::<32>()), &mut data),
        // ── Camellia ──────────────────────────────────────────────────────────
        "camellia128" => run_block(Camellia128::new(&random_array::<16>()), &mut data),
        "camellia192" => run_block(Camellia192::new(&random_array::<24>()), &mut data),
        "camellia256" => run_block(Camellia256::new(&random_array::<32>()), &mut data),
        // ── CAST-128 ──────────────────────────────────────────────────────────
        "cast128" => run_block(Cast128::new(&random_array::<16>()), &mut data),
        // ── DES / 3DES ────────────────────────────────────────────────────────
        "des" => {
            // Loop until we draw a non-weak key (probability of failure ≈ 2⁻⁵²).
            let cipher = loop {
                let k = random_array::<8>();
                if let Ok(c) = Des::new(&k) {
                    break c;
                }
            };
            run_block(cipher, &mut data);
        }
        "3des" => {
            let cipher = loop {
                let k = random_array::<24>();
                if let Ok(c) = TripleDes::new_3key(&k) {
                    break c;
                }
            };
            run_block(cipher, &mut data);
        }
        // ── Grasshopper / Magma ───────────────────────────────────────────────
        "grasshopper" => run_block(Grasshopper::new(&random_array::<32>()), &mut data),
        "magma" => run_block(Magma::new(&random_array::<32>()), &mut data),
        // ── PRESENT ───────────────────────────────────────────────────────────
        "present80" => run_block(Present80::new(&random_array::<10>()), &mut data),
        "present128" => run_block(Present128::new(&random_array::<16>()), &mut data),
        // ── SEED / Serpent / SM4 / Twofish ────────────────────────────────────
        "seed" => run_block(Seed::new(&random_array::<16>()), &mut data),
        "serpent128" => run_block(Serpent128::new(&random_array::<16>()), &mut data),
        "serpent192" => run_block(Serpent192::new(&random_array::<24>()), &mut data),
        "serpent256" => run_block(Serpent256::new(&random_array::<32>()), &mut data),
        "sm4" => run_block(Sm4::new(&random_array::<16>()), &mut data),
        "twofish128" => run_block(Twofish128::new(&random_array::<16>()), &mut data),
        "twofish256" => run_block(Twofish256::new(&random_array::<32>()), &mut data),
        // ── Simon / Speck (representative variants across block sizes) ────────
        "simon32_64" => run_block(Simon32_64::new(&random_array::<8>()), &mut data),
        "simon64_128" => run_block(Simon64_128::new(&random_array::<16>()), &mut data),
        "simon128_128" => run_block(Simon128_128::new(&random_array::<16>()), &mut data),
        "simon128_256" => run_block(Simon128_256::new(&random_array::<32>()), &mut data),
        "speck32_64" => run_block(Speck32_64::new(&random_array::<8>()), &mut data),
        "speck64_128" => run_block(Speck64_128::new(&random_array::<16>()), &mut data),
        "speck128_128" => run_block(Speck128_128::new(&random_array::<16>()), &mut data),
        "speck128_256" => run_block(Speck128_256::new(&random_array::<32>()), &mut data),
        // ── Stream ciphers ────────────────────────────────────────────────────
        "chacha20" => {
            let key = random_array::<32>();
            let nonce = random_array::<12>();
            ChaCha20::new(&key, &nonce).apply_keystream(&mut data);
        }
        "xchacha20" => {
            let key = random_array::<32>();
            let nonce = random_array::<24>();
            XChaCha20::new(&key, &nonce).apply_keystream(&mut data);
        }
        "salsa20" => {
            let key = random_array::<32>();
            let nonce = random_array::<8>();
            Salsa20::new(&key, &nonce).apply_keystream(&mut data);
        }
        "rabbit" => {
            let key = random_array::<16>();
            let iv = random_array::<8>();
            Rabbit::new(&key, &iv).apply_keystream(&mut data);
        }
        "zuc128" => {
            let key = random_array::<16>();
            let iv = random_array::<16>();
            Zuc128::new(&key, &iv).fill(&mut data);
        }
        "snow3g" => {
            let key = random_array::<16>();
            let iv = random_array::<16>();
            Snow3g::new(&key, &iv).fill(&mut data);
        }
        other => {
            eprintln!("unknown cipher: {other}");
            std::process::exit(1);
        }
    };

    let stdout = std::io::stdout();
    let mut h = stdout.lock();
    h.write_all(&data).expect("write stdout");
}
