use crate::rsa;
use num_bigint::BigUint;
use openssl::bn::BigNum;
use openssl::bn::BigNumContext;
use openssl::error::ErrorStack;
use openssl::sha::sha256;

pub fn challenge41() -> Option<()> {
    let (public_key, private_key) = rsa::rsa_gen_keys().ok()?;

    let i = BigNum::from_u32(1337).ok()?;
    let c = rsa::rsa_encrypt_unpadded(&i, &public_key).ok()?;
    let m = rsa::rsa_decrypt_unpadded(&c, &private_key).ok()?;
    assert_eq!(i, m, "rsa is broken");

    let mut ctx = BigNumContext::new().ok()?;
    // Let S be a random number > 1 mod N. Doesn't matter what.
    let mut s = BigNum::new().ok()?;
    public_key.n.rand_range(&mut s).ok()?;
    // C' = ((S**E mod N) C) mod N
    let mut c2 = BigNum::new().ok()?;
    c2.mod_exp(&s, &public_key.e, &public_key.n, &mut ctx)
        .ok()?;
    let c2 = &(&c2 * &c) % &public_key.n;
    let p2 = rsa::rsa_decrypt_unpadded(&c2, &private_key).ok()?;

    //      P'
    // P = ---  mod N
    //      S
    let p2 = &(&p2 * &rsa::invmod(&s, &public_key.n).ok()?) % &public_key.n;
    assert_eq!(i, p2, "message recovery oracle failed");

    println!("[okay] Challenge 41: implement unpadded message recovery oracle");
    Some(())
}

pub fn challenge42() -> Option<()> {
    let (public_key, private_key) = rsa::rsa_gen_keys().ok()?;

    let i = BigNum::from_u32(1337).ok()?;
    let c = rsa::rsa_encrypt(&i, &public_key).ok()?;
    let m = rsa::rsa_decrypt(&c, &private_key).ok()?;
    assert_eq!(i, m, "rsa is broken");

    let m = "a message to verify";
    let sig = rsa::rsa_sign(&m, &public_key).ok()?;
    let sig_ok = rsa::rsa_verify(&m, &private_key, &sig).ok()?;
    assert!(sig_ok, "RSA verify does not work");
    assert!(
        rsa::rsa_verify("other message", &private_key, &sig).ok()? == false,
        "RSA verify does not work"
    );

    fn cube_root(n: &BigNum) -> Result<BigNum, ErrorStack> {
        let b = BigUint::from_bytes_be(&n.to_vec());
        let b = b.nth_root(3);
        BigNum::from_slice(&b.to_bytes_be())
    }

    pub fn rsa_fake_sign(m: &str) -> Result<BigNum, ErrorStack> {
        let hash = sha256(m.as_bytes());
        let padding_str_len = 1024;

        let mut v = vec![0x0; 3 + padding_str_len];
        v[0] = 0x0;
        v[1] = 0x1;
        for i in 2..padding_str_len + 2 {
            v[i] = 0x0;
        }
        v[padding_str_len + 1] = 0xFF;
        v[padding_str_len + 2] = 0x0;
        v.append(&mut hash.to_vec());
        let sig_cubed = BigNum::from_slice(&v)?;
        let sig = cube_root(&sig_cubed)?;
        Ok(sig)
    }

    let i = BigNum::from_u32(1337).ok()?;
    let c = rsa::rsa_encrypt_unpadded(&i, &public_key).ok()?;
    let m = rsa::rsa_decrypt_unpadded(&c, &private_key).ok()?;
    println!("i={i} c={c} m={m}");
    assert_eq!(i, m, "rsa is broken");

    let m = "hi mom";
    let sig = rsa_fake_sign(&m).ok()?;
    let sig_ok = rsa::rsa_verify_insecure(&m, &private_key, &sig).ok()?;
    assert!(sig_ok, "RSA verify does not work");

    println!("[xxxx] Challenge 42: Bleichenbacher's e=3 RSA Attack");
    Some(())
}