use common::{aes_cbc_enc, aes_ecb_enc, gen_random_key, is_ecb, pkcs7_padding};
use rand::prelude::*;

fn encryption_oracle(plaintext: &[u8]) -> Vec<u8> {
    let key = gen_random_key();
    let mut rng = rand::rng();
    let enc_mode = rng.random_range(0..2); // 0 for ECB, 1 for CBC
    let plaintext = plaintext.to_vec();
    let random_prefix_len = rng.random_range(5..11);
    let random_suffix_len = rng.random_range(5..11);
    let mut random_padding = [0u8; 20];
    rng.fill(&mut random_padding);
    let random_prefix = random_padding[..random_prefix_len].to_vec();
    let random_suffix = random_padding[random_padding.len() - random_suffix_len..].to_vec();

    let mut data = Vec::new();
    data.extend(random_prefix);
    data.extend(plaintext);
    data.extend(random_suffix);

    pkcs7_padding(&mut data, 16);

    if enc_mode == 0 {
        // ECB mode in p7
        println!("Using ECB mode.");
        aes_ecb_enc(&data, &key).unwrap()
    } else {
        // CBC mode in p10
        let iv = [0u8; 16];
        println!("Using CBC mode.");
        aes_cbc_enc(&data, &key, &iv).unwrap()
    }
}

fn main() {
    let plaintext = b"YELLOW SUBMARINE".repeat(32);
    let encrypted_data = encryption_oracle(&plaintext);
    // in p8
    if is_ecb(&encrypted_data) {
        println!("ECB mode detected.");
    } else {
        println!("CBC mode detected.");
    }
}