use crate::bytes::Bytes;
use crate::bytes_base64::BytesBase64;
use crate::dsa;
use crate::rsa;
use crate::utils::bnclone;
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, &private_key).ok()?;
    let sig_ok = rsa::rsa_verify(&m, &public_key, &sig).ok()?;
    assert!(sig_ok, "RSA verify does not work");
    assert!(
        rsa::rsa_verify("other message", &public_key, &sig).ok()? == false,
        "RSA verify does not work"
    );

    let sig_ok = rsa::rsa_verify_insecure(&m, &public_key, &sig).ok()?;
    assert!(sig_ok, "RSA verify does not work");
    assert!(
        rsa::rsa_verify_insecure("other message", &public_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())
    }

    fn _cube(n: &BigNum) -> BigNum {
        n * &(n * n)
    }

    pub fn rsa_fake_sign(m: &str) -> Result<BigNum, ErrorStack> {
        let hash = sha256(m.as_bytes());
        let mut v = vec![0x0, 0x1, 0xff, 0x0];
        v.append(&mut hash.to_vec());
        while v.len() < 128 {
            v.push(0);
        }

        // Add one to the cube root to ensure that when the number is
        // cubed again it contains the desired signature.
        let sig_cubed = BigNum::from_slice(&v)?;
        let mut sig = cube_root(&sig_cubed)?;
        sig.add_word(1)?;

        Ok(sig)
    }

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

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

pub mod challenge43 {
    use crate::bytes::Bytes;
    use crate::dsa;
    use crate::rsa;
    use openssl::bn::BigNum;
    use openssl::error::ErrorStack;

    pub fn recover_x() -> Result<Option<BigNum>, ErrorStack> {
        // I used the parameters above.
        let params = dsa::DsaParameters::new()?;
        // I generated a keypair. My pubkey y is given:
        let y = BigNum::from_hex_str("84ad4719d044495496a3201c8ff484feb45b962e7302e56a392aee4abab3e4bdebf2955b4736012f21a08084056b19bcd7fee56048e004e44984e2f411788efdc837a0d2e5abb7b555039fd243ac01f0fb2ed1dec568280ce678e931868d23eb095fde9d3779191b8c0299d6e07bbb283e6633451e535c45513b2d33c99ea17")?;

        // I signed the following message.
        let msg = Bytes::from_utf8("For those that envy a MC it can be hazardous to your health\nSo be friendly, a matter of life and death, just like a etch-a-sketch\n");
        // My SHA1 for this string was d2d0714f014a9784047eaeccf956520045c45265.
        assert_eq!(
            dsa::h(&msg)?,
            BigNum::from_hex_str("d2d0714f014a9784047eaeccf956520045c45265")?
        );

        // They provide s and r as decimal integers and not hex strings. Took me a while
        // to notice that.
        let mut sig = dsa::DsaSig {
            r: BigNum::from_dec_str("548099063082341131477253921760299949438196259240")?,
            s: BigNum::from_dec_str("857042759984254168557880549501802188789837994940")?,
            k: BigNum::from_u32(0)?,
        };
        let msg_h = dsa::h(&msg)?;
        let r_inv = rsa::invmod(&sig.r, &params.q)?;

        // I signed this string with a broken implemention of DSA that generated "k"
        // values between 0 and 2^16. What's my private key?
        for k in 0..2_u32.pow(16) {
            // I get the signature.
            sig.k = BigNum::from_u32(k)?;
            let x = &(&r_inv * &(&(&sig.s * &sig.k) - &msg_h)) % &params.q;
            let keys = dsa::DsaKeys::new_from_x(&params, x)?;
            if y == keys.y {
                return Ok(Some(keys.x));
            }
        }
        Ok(None)
    }
}

pub fn challenge43() -> Option<()> {
    let msg = Bytes::from_utf8("hello, world!");
    let params = dsa::DsaParameters::new().ok()?;
    let keys = dsa::DsaKeys::new(&params).ok()?;
    let sig = keys.sign(&params, &msg).ok()?;
    let result = keys.verify(&params, &msg, &sig).ok()?;
    assert!(result, "verify failed unexpectedly");

    let recovered_x = dsa::recover_x(&params, &msg, &sig).ok()?;
    assert_eq!(recovered_x, keys.x, "DSA x recovery failed");
    let recovered_keys = dsa::DsaKeys::new_from_x(&params, recovered_x).ok()?;
    assert_eq!(recovered_keys.y, keys.y, "DSA y recovery failed");

    let msg = Bytes::from_utf8("hello world!");
    let result = keys.verify(&params, &msg, &sig).ok()?;
    assert!(!result, "verify succeeded unexpectedly");

    // Its SHA-1 fingerprint (after being converted to hex) is:
    // 0954edd5e0afe5542a4adf012611a91912a3ec16
    let x = challenge43::recover_x().ok()??;
    let x_as_hex_str = x.to_hex_str().ok()?.to_ascii_lowercase();
    assert_eq!(
        dsa::h(&Bytes::from_utf8(x_as_hex_str.as_ref())).ok()?,
        BigNum::from_hex_str("0954edd5e0afe5542a4adf012611a91912a3ec16").ok()?,
        "Recovery from none failed"
    );
    println!("[okay] Challenge 43: DSA key recovery from nonce");
    Some(())
}

pub mod challenge44 {
    use crate::bytes::Bytes;
    use crate::dsa;
    use crate::rsa;
    use openssl::bn::BigNum;
    use openssl::bn::BigNumContext;
    use std::io::{BufRead, BufReader};

    pub struct DsaSignedMsg {
        pub msg: Bytes,
        pub r: BigNum,
        pub s: BigNum,
        pub m: BigNum,
    }

    pub fn read_dsa_signed_messages() -> Vec<DsaSignedMsg> {
        let file = std::fs::File::open("data/44.txt").unwrap();
        let mut result = Vec::new();
        let mut lines: Vec<String> = BufReader::new(file).lines().map(|l| l.unwrap()).collect();
        // each message cosists of four lines: msg, s, r, m (sha1 hash of msg)
        for line in lines.chunks_mut(4) {
            line[0].remove_matches("msg: ");
            line[1].remove_matches("s: ");
            line[2].remove_matches("r: ");
            line[3].remove_matches("m: ");
            let msg = Bytes::from_utf8(&line[0]);
            let s = BigNum::from_dec_str(&line[1]).unwrap();
            let r = BigNum::from_dec_str(&line[2]).unwrap();
            let m = BigNum::from_hex_str(&line[3]).unwrap();
            assert_eq!(
                dsa::h(&msg).unwrap(),
                m,
                "Message hash from data/44.txt does not match"
            );
            let d = DsaSignedMsg { msg, r, s, m };
            result.push(d);
        }
        return result;
    }

    pub fn calculate_k(d1: &DsaSignedMsg, d2: &DsaSignedMsg) -> Option<BigNum> {
        // Recovers k if d1 and d2 have been created with the same k.
        let mut ctx = BigNumContext::new().ok()?;
        let params = dsa::DsaParameters::new().ok()?;
        let mut num = BigNum::new().ok()?;
        let mut den = BigNum::new().ok()?;
        let mut den_inv = BigNum::new().ok()?;
        let mut result = BigNum::new().ok()?;
        //     (m1 - m2)
        // k = --------- mod q
        //     (s1 - s2)
        num.mod_sub(&d1.m, &d2.m, &params.q, &mut ctx).ok()?;
        den.mod_sub(&d1.s, &d2.s, &params.q, &mut ctx).ok()?;
        if den_inv.mod_inverse(&den, &params.q, &mut ctx).is_err() {
            return None;
        }
        result.mod_mul(&num, &den_inv, &params.q, &mut ctx).ok()?;
        return Some(result);
    }

    pub fn find_x() -> Option<BigNum> {
        let mut ctx = BigNumContext::new().ok()?;
        let params = dsa::DsaParameters::new().ok()?;
        let msgs = read_dsa_signed_messages();
        let num_msgs = msgs.len();
        for i in 0..(num_msgs - 1) {
            for j in i..num_msgs {
                let k = calculate_k(&msgs[i], &msgs[j]);
                if k.is_none() {
                    continue;
                }
                let k = k.unwrap();
                let m = &msgs[i];
                let r_inv = rsa::invmod(&m.r, &params.q);
                let k_inv = rsa::invmod(&k, &params.q);

                // Reconstruct x from k. If two messages have been encrypted with the same k, we
                // can then use the x to get the same signature s and r. In that case we have found
                // the private key x and return it.
                if r_inv.is_ok() && k_inv.is_ok() {
                    let r_inv = r_inv.ok()?;
                    let k_inv = k_inv.ok()?;
                    let x = &(&r_inv * &(&(&m.s * &k) - &m.m)) % &params.q;
                    let mut r = BigNum::from_u32(0).ok()?;
                    r.mod_exp(&params.g, &k, &params.p, &mut ctx).ok()?;
                    r = &r % &params.q;
                    let s = &(&k_inv * &(&m.m + &(&x * &r))) % &params.q;
                    if s == m.s && r == m.r {
                        return Some(x);
                    }
                }
            }
        }
        return None;
    }
}

pub fn challenge44() -> Option<()> {
    let x = challenge44::find_x()?;
    let x_as_hex_str = x.to_hex_str().ok()?.to_ascii_lowercase();
    assert_eq!(
        dsa::h(&Bytes::from_utf8(x_as_hex_str.as_ref())).ok()?,
        BigNum::from_hex_str("ca8f6f7c66fa362d40760d135b763eb8527d3d52").ok()?,
        "Recovery from repeated nonce failed"
    );
    println!("[okay] Challenge 44: DSA nonce recovery from repeated nonce");
    Some(())
}

pub fn challenge45() -> Option<()> {
    // Part 1: 0 mod p
    //
    // Use the parameters from the previous exercise, but substitute 0 for "g". Generate a
    // signature. You will notice something bad.
    //
    // Observation:
    //
    // - If g=0, because of r := (g^k mod p) mod q, it follows r=0.
    // - When verifying, v := ((g^u1 * y^u2) % p) % q).
    // - Therefore, v=0, and v == r is always true for any signature.
    //
    // Our implementation catches this case, both for signing and verifying.

    // Part 2: 1 mod p
    let mut params = dsa::DsaParameters::new().ok()?;
    fn generate_magic_signature(
        params: &mut dsa::DsaParameters,
    ) -> Result<dsa::DsaSig, ErrorStack> {
        let mut ctx = BigNumContext::new()?;

        //Now, try (p+1) as "g". With this "g", you can generate a magic signature s, r for any DSA
        //public key that will validate against any string.
        params.g = &params.p + &BigNum::from_u32(1)?;

        // For arbitrary z:
        // r = ((y**z) % p) % q
        let mut r = BigNum::from_u32(0)?;
        let z = BigNum::from_u32(1337)?;
        r.mod_exp(&params.g, &z, &params.p, &mut ctx)?;
        r = &r % &params.q;

        //       r
        // s =  --- % q
        //       z
        let s = &rsa::invmod(&z, &params.q)? * &r;

        // k doesn't matter
        let k = BigNum::from_u32(0)?;

        Ok(dsa::DsaSig { r, s, k })
    }

    let sig = generate_magic_signature(&mut params).ok()?;
    let keys = dsa::DsaKeys::new(&params).ok()?;

    let msg = Bytes::from_utf8("Hello, world");
    let result = keys.verify(&params, &msg, &sig).ok()?;
    assert!(result, "verify failed unexpectedly");

    let msg = Bytes::from_utf8("Goodbye, world");
    let result = keys.verify(&params, &msg, &sig).ok()?;
    assert!(result, "verify failed unexpectedly");

    println!("[okay] Challenge 45: DSA parameter tampering");
    Some(())
}

pub fn challenge46() -> Option<()> {
    let m_b64 = BytesBase64::from_base64("VGhhdCdzIHdoeSBJIGZvdW5kIHlvdSBkb24ndCBwbGF5IGFyb3VuZCB3aXRoIHRoZSBGdW5reSBDb2xkIE1lZGluYQ").ok()?;
    let m = BigNum::from_slice(&m_b64.to_bytes().0).ok()?;
    let (public_key, private_key) = rsa::rsa_gen_keys().ok()?;

    let c = rsa::rsa_encrypt_unpadded(&m, &public_key).ok()?;
    let n = BigNum::from_slice(&public_key.n.to_vec()).ok()?;
    assert!(
        rsa::rsa_decrypt_unpadded(&c, &private_key).ok()? == m,
        "rsa does not work"
    );

    let parity_oracle = |cipher: &BigNum| -> bool {
        // Return true or false based on whether the decrypted plaintext was even or odd.
        let cleartext: BigNum = rsa::rsa_decrypt_unpadded(cipher, &private_key).unwrap();
        !cleartext.is_bit_set(0)
    };

    let c_a = rsa::rsa_encrypt_unpadded(&BigNum::from_u32(97).ok()?, &public_key).ok()?;
    assert!(!parity_oracle(&c_a), "parity oracle odd doesn't work");

    let c_b = rsa::rsa_encrypt_unpadded(&BigNum::from_u32(98).ok()?, &public_key).ok()?;
    assert!(parity_oracle(&c_b), "parity oracle even doesn't work");


    // Double by multiplying with (2**e) % n
    let mut ctx = BigNumContext::new().ok()?;
    let two = BigNum::from_u32(2).ok()?;
    let mut multiplier = BigNum::new().ok()?;
    multiplier.mod_exp(&two, &public_key.e, &n, &mut ctx).ok()?;
    let double = |c: BigNum| -> BigNum {
        return &(&c * &multiplier) % &private_key.n;
    };


    let solve = |mut cipher: BigNum, n: BigNum | -> Result<BigNum, ErrorStack> {
        // - You can repeatedly apply this heuristic, once per bit of the message, checking your oracle
        //   function each time.
        // - Your decryption function starts with bounds for the plaintext of [0,n].
        // - Each iteration of the decryption cuts the bounds in half; either the upper bound is reduced
        //   by half, or the lower bound is.
        // - After log2(n) iterations, you have the decryption of the message.
        let mut div = BigNum::from_u32(1)?;
        let mut lower = BigNum::from_u32(0)?;
        let mut upper = BigNum::from_u32(1)?;
        let num_bits = n.num_bits();

        for _ in 0..num_bits {
            div = &div << 1_i32;
            lower = &lower << 1_i32;
            upper = &upper << 1_i32;
            let new = &(&lower + &upper) >> 1_i32;
            cipher = double(cipher);
            if parity_oracle(&cipher) {
                upper = new;
            } else {
                lower = new;
            }
        }
        Ok(&(&upper * &n) / &div)
    };

    let s = solve(bnclone(&c), bnclone(&n)).ok()?;
    assert!(m == s, "RSA parity attack did not work");

    println!("[okay] Challenge 46: RSA parity oracle");
    Some(())
}

pub fn challenge47() -> Option<()> {
    println!("[xxxx] Challenge 47: Bleichenbacher's PKCS 1.5 Padding Oracle (Simple Case)");
    None
}