use crate::bytes::Bytes;
use crate::rsa;
use crate::sha1;
use crate::srp;
use num_bigint::BigUint;
use num_bigint::RandBigInt;
use num_bigint::ToBigUint;
use openssl::bn::BigNum;
use openssl::sha::sha256;
use rand::Rng;

pub mod challenge33 {
    use num_bigint::BigUint;
    use std::fs;

    pub fn load_large_prime() -> BigUint {
        let s = fs::read_to_string("data/33.txt").expect("run from git root dir");
        let b: Vec<u8> = s
            .chars()
            .filter(|c| *c != '\n')
            .map(|c| c.to_digit(16).unwrap().try_into().unwrap())
            .collect();
        BigUint::from_radix_be(&b, 16).unwrap()
    }
}

pub fn challenge33() {
    fn expmod(base: u32, exp: u32, m: u32) -> u32 {
        if exp == 0 {
            return 1;
        }

        if exp % 2 == 0 {
            let s = expmod(base, exp / 2, m);
            (s * s) % m
        } else {
            (expmod(base, exp - 1, m) * base) % m
        }
    }

    fn simple_diffie_hellman() {
        #![allow(non_snake_case)]
        let mut rng = rand::thread_rng();

        // Set a variable "p" to 37 and "g" to 5.
        let p: u32 = 37;
        let g: u32 = 5;

        // Generate "a", a random number mod 37. Now generate "A", which is "g" raised to the "a"
        // power mode 37 --- A = (g**a) % p.
        let a: u32 = rng.gen::<u32>() % p;
        let A = expmod(g, a, p);

        // Do the same for "b" and "B".
        let b: u32 = rng.gen::<u32>() % p;
        let B = expmod(g, b, p);

        // "A" and "B" are public keys. Generate a session key with them; set "s" to "B" raised to
        // the "a" power mod 37 --- s = (B**a) % p.
        let s = expmod(B, a, p);

        // Do the same with A**b, check that you come up with the same "s".
        let s_ = expmod(A, b, p);

        assert_eq!(s, s_, "crypto is broken");
        // To turn "s" into a key, you can just hash it to create 128 bits of key material (or
        // SHA256 it to create a key for encrypting and a key for a MAC).
    }

    fn serious_diffie_hellman() {
        let p = challenge33::load_large_prime();
        let g = 2_u8.to_biguint().unwrap();

        let a = rsa::Keypair::make(&p, &g);
        let b = rsa::Keypair::make(&p, &g);

        let s = b.public.0.modpow(&a.private.0, &p);
        let s_ = a.public.0.modpow(&b.private.0, &p);
        assert_eq!(s, s_, "crypto is broken");
    }

    simple_diffie_hellman();
    serious_diffie_hellman();
    println!("[okay] Challenge 33: implemented Diffie-Hellman");
}

mod challenge34 {
    use crate::bytes::Bytes;
    use crate::cbc;
    use crate::rsa;
    use crate::sha1::Sha1;
    use num_bigint::BigUint;

    pub struct Bot {
        p: BigUint,
        keypair: rsa::Keypair,
        pub s: Option<BigUint>,
        pub key: Option<Bytes>,
    }

    impl Bot {
        pub fn new(p: BigUint, g: &BigUint) -> Self {
            let keypair = rsa::Keypair::make(&p, g);
            Bot {
                p,
                keypair,
                s: None,
                key: None,
            }
        }

        pub fn get_public_key(&self) -> rsa::PublicKey {
            self.keypair.public.clone()
        }

        pub fn exchange_keys(&mut self, public: &rsa::PublicKey) -> rsa::PublicKey {
            let s = public.0.modpow(&self.keypair.private.0, &self.p);
            self.s = Some(s);
            self.make_encryption_key();
            self.keypair.public.clone()
        }

        fn make_encryption_key(&mut self) {
            assert!(self.s.is_some(), "keys have not been exchanged");
            if self.key.is_some() {
                return;
            }
            let mut sha1 = Sha1::default();
            let key = &sha1.hash(&Bytes(self.s.clone().unwrap().to_bytes_be())).0[0..16];
            self.key = Some(Bytes(key.to_vec()));
        }

        pub fn encrypt(&mut self, message: &Bytes) -> Bytes {
            assert!(self.key.is_some(), "keys have not been exchanged");

            const BLOCK_SIZE: usize = 16;
            let key = self.key.as_ref().unwrap();
            let mut iv = Bytes::random(16);
            let mut message = Bytes(message.0.clone());
            message.pad_pkcs7(BLOCK_SIZE);

            // AES-CBC(SHA1(s)[0:16], iv=random(16), msg) + iv
            let mut cipher = cbc::encrypt(key, &iv, &message).0;
            cipher.append(&mut iv.0);
            Bytes(cipher)
        }

        pub fn decrypt(&mut self, cipher: &Bytes) -> Bytes {
            const BLOCK_SIZE: usize = 16;
            assert!(self.key.is_some(), "keys have not been exchanged");

            let key = self.key.as_ref().unwrap();
            let cipher_len = cipher.len();

            let iv = Bytes(cipher.0[(cipher_len - BLOCK_SIZE)..cipher_len].to_vec());
            let cipher = Bytes(cipher.0[0..cipher_len - BLOCK_SIZE].to_vec());
            let mut message = cbc::decrypt(key, &iv, &cipher);
            message.remove_pkcs7(BLOCK_SIZE);
            message
        }
    }

    pub fn decrypt_with_s(cipher: &Bytes, s: &BigUint) -> Bytes {
        const BLOCK_SIZE: usize = 16;
        let mut sha1 = Sha1::default();
        let key = Bytes(sha1.hash(&Bytes(s.to_bytes_be())).0[0..16].to_vec());

        let cipher_len = cipher.len();
        let iv = Bytes(cipher.0[(cipher_len - BLOCK_SIZE)..cipher_len].to_vec());
        let cipher = Bytes(cipher.0[0..cipher_len - BLOCK_SIZE].to_vec());
        let mut message = cbc::decrypt(&key, &iv, &cipher);
        message.remove_pkcs7(BLOCK_SIZE);
        message
    }
}

pub fn challenge34() {
    fn echo() {
        let p = challenge33::load_large_prime();
        let g = 2_u8.to_biguint().unwrap();

        let mut a = challenge34::Bot::new(p.clone(), &g);

        // A->B: Send "p", "g", "A"
        let mut b = challenge34::Bot::new(p, &g);

        // B->A: Send "B"
        a.exchange_keys(&b.exchange_keys(&a.get_public_key()));
        assert_eq!(a.s, b.s, "crypto is broken");

        // A->B: Send AES-CBC(SHA1(s)[0:16], iv=random(16), msg) + iv
        let message_a = Bytes::from_utf8("Will this work? Let's make it longer.");
        let cipher_a = a.encrypt(&message_a);

        // B->A: Send AES-CBC(SHA1(s)[0:16], iv=random(16), A's msg) + iv
        let message_b = b.decrypt(&cipher_a);
        let cipher_b = b.encrypt(&message_b);
        let roundtrip = a.decrypt(&cipher_b);
        assert_eq!(message_a, roundtrip, "we broke it fam");
    }

    fn echo_mitm() {
        let p = challenge33::load_large_prime();
        let g = 2_u8.to_biguint().unwrap();
        let p_public = rsa::PublicKey(p.clone());

        // A->M Send "p", "g", "A"
        let mut a = challenge34::Bot::new(p.clone(), &g);
        let mut m = challenge34::Bot::new(p.clone(), &g);
        m.exchange_keys(&a.get_public_key());

        // M->B Send "p", "g", "p"
        // B->M Send "B"
        let mut b = challenge34::Bot::new(p, &g);
        b.exchange_keys(&p_public);

        // M->A Send "p"
        a.exchange_keys(&p_public);

        // A->M Send AES-CBC(SHA1(s)[0:16], iv=random(16), msg) + iv
        // M->B Relay that to B
        let message_a = Bytes::from_utf8("Will this work again? I don't doubt it.");
        let cipher_a = a.encrypt(&message_a);

        // B->M Send AES-CBC(SHA1(s)[0:16], iv=random(16), A's msg) + iv
        let message_b = b.decrypt(&cipher_a);
        let cipher_b = b.encrypt(&message_b);

        // M->A Relay that to A
        let roundtrip = a.decrypt(&cipher_b);
        assert_eq!(message_a, roundtrip, "we broke it fam");

        let restored_a = attack(&cipher_a);
        let restored_b = attack(&cipher_b);
        assert_eq!(message_a, restored_a, "attack didn't work");
        assert_eq!(message_b, restored_b, "attack didn't work");
    }

    fn attack(cipher: &Bytes) -> Bytes {
        // Do the DH math on this quickly to see what that does to the predictability of the key.
        // With p = A; s = (A ** b) % p becomes (p ** b) % p = 0 => s = 0
        let s = 0_u8.to_biguint().unwrap();
        challenge34::decrypt_with_s(cipher, &s)
    }

    echo();
    echo_mitm();
    println!("[okay] Challenge 34: implement MITM key-fixing attack on DH");
}

pub fn challenge35() {
    fn echo(message: &Bytes, g: &BigUint, p: &BigUint) -> (Bytes, Bytes) {
        let mut a = challenge34::Bot::new(p.clone(), g);
        let mut b = challenge34::Bot::new(p.clone(), g);
        a.exchange_keys(&b.exchange_keys(&a.get_public_key()));
        assert_eq!(a.s, b.s, "crypto is broken");

        let cipher_a = a.encrypt(message);
        let message_b = b.decrypt(&cipher_a);
        let cipher_b = b.encrypt(&message_b);
        let roundtrip = a.decrypt(&cipher_b);
        assert_eq!(*message, roundtrip, "we broke it fam");
        (cipher_a, cipher_b)
    }

    fn attack_g1(cipher: &Bytes) -> Bytes {
        // A = (g ** a) % p, g = 1 -> A = 1
        // s = (A ** b) % p, A = 1 -> s = 1
        let s = 1_u8.to_biguint().unwrap();
        challenge34::decrypt_with_s(cipher, &s)
    }

    fn attack_g_is_p(cipher: &Bytes) -> Bytes {
        // A = (g ** a) % p, g = p -> A = 0
        // s = (A ** b) % p, A = 0 -> s = 0
        let s = 0_u8.to_biguint().unwrap();
        challenge34::decrypt_with_s(cipher, &s)
    }

    fn attack_g_is_p_minus_one(cipher: &Bytes, p: &BigUint) -> Bytes {
        // A = (g ** a) % p, g = p - 1 -> A = 1 | (p - 1)
        // s = (A ** b) % p, A = 1 | (p - 1) -> s = 1 | (p - 1)
        let s1 = 1_u8.to_biguint().unwrap();
        let s2 = p - 1_u8.to_biguint().unwrap();
        let m1 = challenge34::decrypt_with_s(cipher, &s1);
        let m2 = challenge34::decrypt_with_s(cipher, &s2);

        match m1.ascii_score().cmp(&m2.ascii_score()) {
            std::cmp::Ordering::Greater => m1,
            std::cmp::Ordering::Less => m2,
            std::cmp::Ordering::Equal => panic!("how can they be equal?"),
        }
    }

    // Do the MITM attack again, but play with "g". Write attacks for each.
    let m = Bytes::from_utf8("Quick to the point, to the point, no faking");
    let p = challenge33::load_large_prime();
    let g = 2_u8.to_biguint().unwrap();
    echo(&m, &g, &p);

    // g = 1
    let g = 1_u8.to_biguint().unwrap();
    let (c1, c2) = echo(&m, &g, &p);
    assert_eq!(attack_g1(&c1), m, "failed to attack g = 1");
    assert_eq!(attack_g1(&c2), m, "failed to attack g = 1");

    // g = p
    let g = p.clone();
    let (c1, c2) = echo(&m, &g, &p);
    assert_eq!(attack_g_is_p(&c1), m, "failed to attack g = p");
    assert_eq!(attack_g_is_p(&c2), m, "failed to attack g = p");

    // g = p - 1
    let g = p.clone() - 1_u8.to_biguint().unwrap();
    let (c1, c2) = echo(&m, &g, &p);
    assert_eq!(attack_g_is_p_minus_one(&c1, &p), m, "failed g = p - 1");
    assert_eq!(attack_g_is_p_minus_one(&c2, &p), m, "failed g = p - 1");

    println!("[okay] Challenge 35: implement MITM with malicious g attack on DH");
}

pub fn challenge36() -> Option<()> {
    let mut rng = rand::thread_rng();

    let mut server = srp::Server::default();
    let email = "john1337@wayne.com";
    let password = "horse planet carpet country";
    let parameters = server.register(email, password)?;

    // Establish session
    let n = &parameters.n;
    let g = &parameters.g;
    let a = rng.gen_biguint_below(n);
    let a_public = g.modpow(&a, n);
    let session_keys = server.exchange(email, &a_public)?;

    // Client
    // Generate S = (B - k * g**x)**(a + u * x) % N
    // Generate K = SHA256(S)
    let u = srp::compute_u(&a_public, &session_keys.b_public)?;
    let x = srp::hash_password(session_keys.salt, password)?;
    let b_public = &session_keys.b_public;
    let k = &parameters.k;
    let s_client = (b_public - ((k * g.modpow(&x, n)) % n)).modpow(&(a + &u * x), n);
    let k_client = sha256(&s_client.to_bytes_be());

    // Server
    let k_server = server.get_k(email)?;
    assert_eq!(k_client, k_server);

    // I don't have HMAC-SHA256, so I will use HMAC-SHA1 instead.
    // C->S Send HMAC-SHA256(K, salt)
    let salt = Bytes(session_keys.salt.to_be_bytes().to_vec());
    let mac = sha1::hmac_sha1(&Bytes(k_client.to_vec()), &salt);
    assert!(server.log_in(email, &mac)?);

    println!("[okay] Challenge 36: implement secure remote password");
    Some(())
}

pub fn challenge37() -> Option<()> {
    let mut rng = rand::thread_rng();

    let mut server = srp::Server::default();
    let email = "john1337@wayne.com";
    let password = "horse planet carpet country";
    let parameters = server.register(email, password)?;

    let mut log_in_with_password = || -> Option<bool> {
        //  Get your SRP working in an actual client-server setting. "Log in" with a
        //  valid password using the protocol.
        let n = &parameters.n;
        let g = &parameters.g;
        let a = rng.gen_biguint_below(n);
        let a_public = g.modpow(&a, n);
        let session_keys = server.exchange(email, &a_public)?;

        // Client
        let u = srp::compute_u(&a_public, &session_keys.b_public)?;
        let x = srp::hash_password(session_keys.salt, password)?;
        let b_public = &session_keys.b_public;
        let k = &parameters.k;
        let s_client = (b_public - ((k * g.modpow(&x, n)) % n)).modpow(&(a + &u * x), n);
        let k_client = sha256(&s_client.to_bytes_be());

        let salt = Bytes(session_keys.salt.to_be_bytes().to_vec());
        let mac = sha1::hmac_sha1(&Bytes(k_client.to_vec()), &salt);
        server.log_in(email, &mac)
    };

    assert!(log_in_with_password()?);

    let mut log_in_without_password_a_zero = || -> Option<bool> {
        // Now log in without your password by having the client send 0 as its "A"
        // value.
        let a_public = 0_u8.to_biguint()?;
        let session_keys = server.exchange(email, &a_public)?;

        // What does this to the "S" value that both sides compute? It becomes zero.
        let s_client: BigUint = 0_u8.to_biguint()?;
        let k_client = sha256(&s_client.to_bytes_be());

        let salt = Bytes(session_keys.salt.to_be_bytes().to_vec());
        let mac = sha1::hmac_sha1(&Bytes(k_client.to_vec()), &salt);
        server.log_in(email, &mac)
    };

    assert!(log_in_without_password_a_zero()?);

    let mut log_in_without_password_multiple_n = || -> Option<bool> {
        // Now log in without your password by having the client send N, N*2, &c.
        let n = &parameters.n;
        let a_public = n * 2_u8.to_biguint()?;
        let session_keys = server.exchange(email, &a_public)?;

        // If A is a multiple of N, the term $(A * v**u) ** b % N$ also becomes zero.
        let s_client: BigUint = 0_u8.to_biguint()?;
        let k_client = sha256(&s_client.to_bytes_be());

        let salt = Bytes(session_keys.salt.to_be_bytes().to_vec());
        let mac = sha1::hmac_sha1(&Bytes(k_client.to_vec()), &salt);
        server.log_in(email, &mac)
    };

    assert!(log_in_without_password_multiple_n()?);

    println!("[okay] Challenge 37: break SRP with zero key");
    Some(())
}

pub fn challenge38() -> Option<()> {
    let mut rng = rand::thread_rng();

    // Client and Server
    let n = challenge33::load_large_prime();
    let g = 2_u8.to_biguint()?;
    let password = "bar";

    // Server: x = SHA256(salt|password); v = g**x % n;
    let salt: u32 = rng.gen();
    let x = srp::hash_password(salt, &password)?;
    let v = g.modpow(&x, &n);

    // C->C: I, A = g**a % n;
    let a = rng.gen_biguint_below(&n);
    let a_public = g.modpow(&a, &n);

    // S->C: salt, B = g**b % n, u = 128 bit random number
    // Pose as the server and use arbitrary values for b, B, u, and salt.
    let b = 11_u8.to_biguint()?;
    let b_public = g.modpow(&b, &n);
    let u: u128 = rng.gen();

    // Client: x = SHA256(salt|password); S = B**(a + ux) % n; K = SHA256(S);
    let x = srp::hash_password(salt, &password)?;
    let s_client = b_public.modpow(&(a + &u * x), &n);
    let k_client = sha256(&s_client.to_bytes_be());

    // Server: S = (A * v ** u)**b % n; K = SHA256(S);
    let s_server = (&a_public * v.modpow(&u.to_biguint()?, &n)).modpow(&b, &n);
    let k_server = sha256(&s_server.to_bytes_be());

    assert_eq!(k_client, k_server);

    // C->S
    //   Send HMAC-SHA256(K, salt)
    let salt_bytes = Bytes(salt.to_be_bytes().to_vec());
    let mac_client = sha1::hmac_sha1(&Bytes(k_client.to_vec()), &salt_bytes);

    // S->C
    let mac_server = sha1::hmac_sha1(&Bytes(k_server.to_vec()), &salt_bytes);
    assert_eq!(mac_server, mac_client, "HMAC verification failed");

    fn crack_password(
        mac_client: &Bytes,
        salt: u32,
        a_public: &BigUint,
        u: u128,
    ) -> Option<String> {
        // Now, run the protocol as a MITM attacker: pose as the server and use arbitrary values
        // for b, B, u, and salt.  Crack the password from A's HMAC-SHA256(K, salt).
        let dict = ["foo", "bar", "lol"];
        let g = 2_u8.to_biguint()?;
        let n = challenge33::load_large_prime();
        let b = 11_u8.to_biguint()?;
        let salt_bytes = Bytes(salt.to_be_bytes().to_vec());

        for password in dict {
            let x = srp::hash_password(salt, &password)?;
            let v = g.modpow(&x, &n);
            let s_attack = (a_public * v.modpow(&u.to_biguint()?, &n)).modpow(&b, &n);
            let k_attack = sha256(&s_attack.to_bytes_be());
            let mac_attack = sha1::hmac_sha1(&Bytes(k_attack.to_vec()), &salt_bytes);
            if mac_attack == *mac_client {
                return Some(password.to_owned());
            }
        }

        None
    }

    let cracked_password = crack_password(&mac_client, salt, &a_public, u);
    match cracked_password {
        Some(p) if p == password => {
            println!("[okay] Challenge 38: offline dictionary attack on SRP");
        }
        _ => println!("[FAIL] Challenge 38: offline dictionary attack on SRP"),
    };

    Some(())
}

pub fn challenge39() -> Option<()> {
    // I recommend you not bother with primegen,
    // but do take the time to get your own EGCD and
    // invmod algorithm working.
    let a = BigNum::from_u32(17).ok()?;
    let n = BigNum::from_u32(3120).ok()?;
    let r = BigNum::from_u32(2753).ok()?;
    assert_eq!(rsa::invmod(&a, &n).ok()?, r, "invmod does not work");

    let (public_key, private_key) = rsa::rsa_gen_keys().ok()?;

    // Test this out with a number, like "42".
    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 i = "Party, party, party, all night long. Yuah!";
    let c = rsa::rsa_encrypt_str(&i, &public_key).ok()?;
    let m = rsa::rsa_decrypt_str(&c, &private_key).ok()?;
    assert_eq!(i, &m, "string rsa is broken");

    println!("[okay] Challenge 39: implement RSA");
    Some(())
}

pub fn challenge40() -> Option<()> {
    let m = BigNum::from_u32(1337).ok()?;

    // 1. Capturing any 3 of the ciphertexts and their corresponding pubkeys
    let (p_0, _) = rsa::rsa_gen_keys().ok()?;
    let (p_1, _) = rsa::rsa_gen_keys().ok()?;
    let (p_2, _) = rsa::rsa_gen_keys().ok()?;
    let c_0 = rsa::rsa_encrypt(&m, &p_0).ok()?;
    let c_1 = rsa::rsa_encrypt(&m, &p_1).ok()?;
    let c_2 = rsa::rsa_encrypt(&m, &p_2).ok()?;

    // 2. Using the CRT to solve for the number represented by the three ciphertexts (which are
    //    residues mod their respective pubkeys)
    let rsa::RsaPublicKey { e: _, n: n_0 } = p_0;
    let rsa::RsaPublicKey { e: _, n: n_1 } = p_1;
    let rsa::RsaPublicKey { e: _, n: n_2 } = p_2;

    // N_012 is the product of all three moduli
    let n_012 = &(&n_0 * &n_1) * &n_2;
    let m_s_0 = &n_1 * &n_2;
    let m_s_1 = &n_0 * &n_2;
    let m_s_2 = &n_0 * &n_1;

    // result =
    //   (c_0 * m_s_0 * invmod(m_s_0, n_0)) +
    //   (c_1 * m_s_1 * invmod(m_s_1, n_1)) +
    //   (c_2 * m_s_2 * invmod(m_s_2, n_2)) mod N_012
    let result = &(&(&(&(&c_0 * &m_s_0) * &rsa::invmod(&m_s_0, &n_0).ok()?)
        + &(&(&c_1 * &m_s_1) * &rsa::invmod(&m_s_1, &n_1).ok()?))
        + &(&(&c_2 * &m_s_2) * &rsa::invmod(&m_s_2, &n_2).ok()?))
        % &n_012;

    // 3. Taking the cube root of the resulting number
    fn cube_root(n: &BigNum) -> Option<BigNum> {
        let b = BigUint::from_bytes_be(&n.to_vec());
        let b = b.nth_root(3);
        BigNum::from_slice(&b.to_bytes_be()).ok()
    }

    let m_cubed = &(&m * &m) * &m;
    assert_eq!(m_cubed, result, "CRT implementation did not work");

    let c = cube_root(&result)?;
    assert_eq!(c, m, "cube root implementation did not work");

    println!("[okay] Challenge 40: implement an E=3 RSA Broadcast attack");
    Some(())
}