use crate::{bytes::Bytes, cbc, ctr, ecb, md4, parser, sha1, utils};
use std::path::Path;

pub fn challenge25() {
    // Now, write the code that allows you to "seek" into the ciphertext,
    // decrypt, and re-encrypt with different plaintext. Expose this as a
    // function, like, "edit(ciphertext, key, offset, newtext)".
    fn edit(ciphertext: &Bytes, key: &Bytes, offset: usize, newtext: &Vec<u8>) -> Bytes {
        let mut plaintext = ctr::decrypt(key, 0, ciphertext);
        assert!(
            offset + newtext.len() <= plaintext.len(),
            "challenge25 - edit - out of bounds"
        );
        plaintext.0[offset..(newtext.len() + offset)].copy_from_slice(&newtext[..]);
        ctr::encrypt(key, 0, &plaintext)
    }

    let cipher = utils::read_base64("data/25.txt");
    let key = Bytes::from_utf8("YELLOW SUBMARINE");
    let plaintext = ecb::decrypt(&key, &cipher);

    let key = Bytes::random(16);
    let nonce: u64 = 0; // otherwise edit would require the nonce too?

    // Imagine the "edit" function was exposed to attackers by means of an API
    // call that didn't reveal the key or the original plaintext; the attacker
    // has the ciphertext and controls the offset and "new text". Recover the
    // original plaintext.
    let ciphertext = ctr::encrypt(&key, nonce, &plaintext);
    let newtext = vec![b'a'; ciphertext.len()];
    let cipher_newtext = edit(&ciphertext, &key, 0, &newtext);
    let keystream = utils::xor(&newtext, &cipher_newtext.0);
    let recovered_plaintext = Bytes(utils::xor(&keystream, &ciphertext.0));
    assert_eq!(plaintext, recovered_plaintext);

    println!("[okay] Challenge 25: recovered AES CTR plaintext via edit");

    // A folkloric supposed benefit of CTR mode is the ability to easily "seek forward" into the
    // ciphertext; to access byte N of the ciphertext, all you need to be able to do is generate
    // byte N of the keystream. Imagine if you'd relied on that advice to, say, encrypt a disk.
    // Answer: you would have to run through the whole cipher stream till you reach the point where
    // you want to decrypt.
}

pub fn challenge26() {
    fn encrypt(input: &str, key: &Bytes) -> Bytes {
        let mut r = String::new();
        for c in input.chars() {
            assert!(!(c == ';' || c == '='), "encrypt: invalid char {}", c);
        }
        r.push_str("comment1=cooking%20MCs;userdata=");
        r.push_str(input);
        r.push_str(";comment2=%20like%20a%20pound%20of%20bacon");
        let data = Bytes(r.as_bytes().to_vec());
        ctr::encrypt(key, 0, &data)
    }

    let key = Bytes::random(16);

    // 0               16              32              48              64
    // 0..34..78..bc..f0..34..78..bc..f0..34..78..bc..f0..34..78..bc..f0..34..78..bc..f
    // comment1=cooking%20MCs;userdata=aaaaaaaaaaaaaaaa;comment2=%20like%20a%20pound%20of%20bacon
    // comment1=cooking%20MCs;userdata=fobar;admin=true;comment2=%20like%20a%20pound%20of%20bacon
    let input = "aaaaaaaaaaaaaaaa";
    let cipher = encrypt(input, &key);
    let keystream = utils::xor(&cipher.0[32..48], &Bytes::from_utf8(input).0);

    let input = "fobar;admin=true";
    let mut flipped_cipher = cipher.0[0..32].to_vec();
    flipped_cipher.append(&mut utils::xor(&keystream, input.as_bytes()));
    flipped_cipher.append(&mut cipher.0[48..cipher.len()].to_vec());

    let cleartext = ctr::decrypt(&key, 0, &Bytes(flipped_cipher));
    let dict = parser::parse_key_value(&cleartext.to_utf8());
    let admin_status = dict.get("admin").unwrap();
    assert_eq!(admin_status, "true");
    println!("[okay] Challenge 26: admin={}", admin_status);
}

pub fn challenge27() {
    fn encrypt(key: &Bytes) -> Bytes {
        // AES-CBC(P_1, P_2, P_3) -> C_1, C_2, C_3
        let mut ct = Bytes::from_utf8("comment1=cooking%20MCs;userdata=secretsaucenouse");
        ct.pad_pkcs7(16);
        cbc::encrypt(key, key, &ct)
    }

    fn decrypt(key: &Bytes, cipher: &Bytes) -> Result<Bytes, Bytes> {
        // The CBC code from exercise 16 encrypts a URL string. Verify each byte
        // of the plaintext for ASCII compliance (ie, look for high-ASCII
        // values). Noncompliant messages should raise an exception or return an
        // error that includes the decrypted plaintext (this happens all the
        // time in real systems, for what it's worth).
        let mut cleartext = cbc::decrypt(key, key, cipher);
        cleartext.remove_pkcs7(16);
        match std::str::from_utf8(&cleartext.0) {
            Ok(_) => Ok(cleartext),
            Err(_) => Err(cleartext),
        }
    }

    fn modify(cipher: &Bytes) -> Bytes {
        // C_1, C_2, C_3 -> C_1, 0, C_1
        let mut c1 = cipher.get_block(0, 16).0;
        let mut modified = c1.clone();
        modified.append(&mut vec![0; 16]);
        modified.append(&mut c1);
        Bytes(modified)
    }

    fn recover_key(plaintext: &Bytes) -> Bytes {
        // P'_1 XOR P'_3
        let block_size = 16;
        Bytes(utils::xor(
            &plaintext.get_block(0, block_size).0,
            &plaintext.get_block(2, block_size).0,
        ))
    }

    let key = Bytes::random(16);

    // Use your code to encrypt a message that is at least 3 blocks long:
    let cipher = encrypt(&key);

    //  Modify the message (you are now the attacker):
    let cipher = modify(&cipher);
    let plaintext = decrypt(&key, &cipher);

    //  As the attacker, recovering the plaintext from the error, extract the key:
    let recovered_key = match plaintext {
        Ok(plaintext) => panic!("Ok: {:?}", plaintext.to_utf8()),
        Err(plaintext) => recover_key(&plaintext),
    };
    assert_eq!(key, recovered_key);
    println!("[okay] Challenge 27: recovered key successfully");
}

pub fn challenge28() {
    let mut sha1 = sha1::Sha1::default();
    assert_eq!(
        Bytes::from_hex("0098ba824b5c16427bd7a1122a5a442a25ec644d"),
        sha1.hash(&Bytes(vec![b'a'; 64]))
    );

    sha1.reset();
    assert_eq!(
        Bytes::from_hex("ad5b3fdbcb526778c2839d2f151ea753995e26a0"),
        sha1.hash(&Bytes(vec![b'a'; 128]))
    );

    sha1.reset();
    assert_eq!(
        Bytes::from_hex("7e240de74fb1ed08fa08d38063f6a6a91462a815"),
        sha1.hash(&Bytes(vec![b'a'; 3])),
    );

    sha1.reset();
    assert_eq!(
        Bytes::from_hex("da39a3ee5e6b4b0d3255bfef95601890afd80709"),
        sha1.hash(&Bytes(vec![])),
    );

    // Verify that you cannot tamper with the message without breaking the MAC
    // you've produced, and that you can't produce a new MAC without knowing the
    // secret key.
    let mut message = Bytes::from_utf8("love, love, love");
    let key = Bytes::from_utf8("kisses!");
    let mac = sha1::authenticate(&message, &key);
    message.flip_bit(2, 3);
    assert!(!sha1::verify(&message, &key, &mac));

    println!("[okay] Challenge 28: implemented SHA-1");
}

pub fn challenge29() {
    fn hash_fixated(bytes: &Bytes, fixture: &Bytes, byte_len: u64) -> Bytes {
        // Now, take the SHA-1 secret-prefix MAC of the message you want to forge --- this is just
        // a SHA-1 hash --- and break it into 32 bit SHA-1 registers (SHA-1 calls them "a", "b",
        // "c", &c).
        let mut s = sha1::Sha1::default();
        let fixate: Vec<u32> = fixture
            .0
            .chunks(4)
            .map(|c| u32::from_be_bytes(c.try_into().unwrap()))
            .collect();

        // Modify your SHA-1 implementation so that callers can pass in new
        // values for "a", "b", "c" &c (they normally start at magic numbers).
        // With the registers "fixated", hash the additional data you want to
        // forge.
        s.fix(fixate.try_into().unwrap(), byte_len);
        s.hash(bytes)
    }

    // use random
    let key = Bytes::random_range(2, 64);
    let message = Bytes::from_utf8(
        "comment1=cooking%20MCs;userdata=foo;comment2=%20like%20a%20pound%20of%20bacon",
    );
    let mac = sha1::authenticate(&message, &key);
    assert!(sha1::verify(&message, &key, &mac));

    let mut forged_message = vec![];
    let mut mac_forged = Bytes(vec![]);
    for key_len in 1..128 {
        // get padding for key || orig-message
        let key_guessed = vec![b'z'; key_len]; // key-guessed
        let mut bytes = key_guessed.clone();
        bytes.append(&mut message.0.clone()); // original-message
        let s1 = sha1::Sha1::default();
        let glue_padding = s1.get_padding(&bytes); // glue-padding

        // forget MAC via fixture: make sure to fix sha1.h *and* sha1.byte_length
        let byte_length = (key_guessed.len() + message.len() + glue_padding.len()) as u64;
        let new_message = b"admin=true".to_vec(); // new-message
        mac_forged = hash_fixated(&Bytes(new_message.clone()), &mac, byte_length);

        // forge message: original-message || glue-padding || new-message
        forged_message = message.0.clone();
        forged_message.append(&mut glue_padding.clone());
        forged_message.append(&mut new_message.clone());
        let r = sha1::verify(&Bytes(forged_message.clone()), &key, &mac_forged);
        if r {
            break;
        }
    }

    assert!(sha1::verify(&Bytes(forged_message), &key, &mac_forged));
    println!("[okay] Challenge 29: extended SHA-1 keyed message successfully");
}

pub fn challenge30() {
    fn test_md4() {
        assert_eq!(
            md4::hash(&Bytes::from_utf8("")),
            Bytes::from_hex("31d6cfe0d16ae931b73c59d7e0c089c0"),
        );
        assert_eq!(
            md4::hash(&Bytes::from_utf8("a")),
            Bytes::from_hex("bde52cb31de33e46245e05fbdbd6fb24"),
        );
        assert_eq!(
            md4::hash(&Bytes::from_utf8("abc")),
            Bytes::from_hex("a448017aaf21d8525fc10ae87aa6729d"),
        );
        assert_eq!(
            md4::hash(&Bytes::from_utf8("abcdefghijklmnopqrstuvwxyz")),
            Bytes::from_hex("d79e1c308aa5bbcdeea8ed63df412da9"),
        );
        assert_eq!(
            md4::hash(&Bytes(vec![b'a'; 1337])),
            Bytes::from_hex("9a4bceae0ae389c4653ad92cfd7bfc3e"),
        );
    }

    // extend MD4 copy and pasted from SHA-1 #allow!(dont-repeat-yourself)
    fn hash_fixated(bytes: &Bytes, fixture: &Bytes, byte_len: u64) -> Bytes {
        let mut m = md4::Md4Core::default();
        let fixate: Vec<u32> = fixture
            .0
            .chunks(4)
            .map(|c| u32::from_le_bytes(c.try_into().unwrap()))
            .collect();
        m.fix(fixate.try_into().unwrap(), byte_len);
        m.hash(bytes)
    }

    test_md4();
    let key = Bytes::random_range(2, 64);
    let message = Bytes::from_utf8(
        "comment1=cooking%20MCs;userdata=foo;comment2=%20like%20a%20pound%20of%20bacon",
    );
    let mac = md4::authenticate(&message, &key);
    assert!(md4::verify(&message, &key, &mac));

    let mut forged_message = vec![];
    let mut mac_forged = Bytes(vec![]);
    for key_len in 1..128 {
        // get padding for key || orig-message
        let key_guessed = vec![b'z'; key_len]; // key-guessed
        let mut bytes = key_guessed.clone();
        bytes.append(&mut message.0.clone()); // original-message
        let md4 = md4::Md4Core::default();
        let glue_padding = md4.get_padding(&bytes); // glue-padding

        // forget MAC via fixture: make sure to fix md4.state *and* md4.byte_length
        let byte_length = (key_guessed.len() + message.len() + glue_padding.len()) as u64;
        let new_message = b"admin=true".to_vec(); // new-message
        mac_forged = hash_fixated(&Bytes(new_message.clone()), &mac, byte_length);

        // forge message: original-message || glue-padding || new-message
        forged_message = message.0.clone();
        forged_message.append(&mut glue_padding.clone());
        forged_message.append(&mut new_message.clone());
        let r = md4::verify(&Bytes(forged_message.clone()), &key, &mac_forged);
        if r {
            break;
        }
    }

    assert!(md4::verify(&Bytes(forged_message), &key, &mac_forged));
    println!("[okay] Challenge 30: implemented and extended MD4 successfully");
}

mod challenge31 {
    use crate::{bytes::Bytes, sha1};
    use std::fs;
    use std::path::Path;
    use std::{thread, time};

    pub fn verify(file: &Path, signature: &[u8], delay: u64) -> bool {
        // Have the server generate an HMAC key, and then verify that the "signature" on incoming
        // requests is valid for "file", using the "==" operator to compare the valid MAC for a
        // file with the "signature" parameter (in other words, verify the HMAC the way any normal
        // programmer would verify it).
        let key = Bytes::from_utf8("sosecretbb");
        let contents = fs::read_to_string(file);
        assert!(contents.is_ok(), "Could not read: {}", file.display());
        let contents = Bytes(contents.unwrap().as_bytes().to_vec());
        insecure_compare(&sha1::hmac_sha1(&key, &contents).0, signature, delay)
    }

    fn insecure_compare(a: &[u8], b: &[u8], delay: u64) -> bool {
        // Write a function, call it "insecure_compare", that implements the == operation by doing
        // byte-at-a-time comparisons with early exit (ie, return false at the first non-matching
        // byte).
        let delay = time::Duration::from_millis(delay);
        if a.len() != b.len() {
            return false;
        }
        for (a, b) in a.iter().zip(b.iter()) {
            if a != b {
                return false;
            }
            // In the loop for "insecure_compare", add a 50ms sleep after each byte.
            thread::sleep(delay);
        }
        true
    }

    pub fn _attack(file: &Path, delay: u64) -> Bytes {
        const BLOCK_SIZE: usize = 20;
        let mut sig = vec![0x0; BLOCK_SIZE];
        for i in 0..BLOCK_SIZE {
            let mut max_tuple: (u128, u8) = (u128::MIN, 0);
            for c in 0_u8..=255_u8 {
                let now = time::Instant::now();
                sig[i] = c;
                verify(file, &sig, delay);
                let elapsed = now.elapsed().as_micros();
                if elapsed > max_tuple.0 {
                    max_tuple = (elapsed, c);
                }
            }
            sig[i] = max_tuple.1;
        }
        Bytes(sig)
    }
}

pub fn challenge31() {
    let key = Bytes::from_utf8("YELLOW SUBMARINE");
    let message = Bytes::from_utf8("Attact at dawn after tomorrow when it's cold inside.");
    assert_eq!(
        sha1::hmac_sha1(&key, &message),
        Bytes::from_hex("8232f3d05afb6bce7e09fe764885cc158e435e36")
    );

    let path = Path::new("data/12.txt");
    let expected_sig = Bytes::from_hex("62f4527ea6cb716d0ad1ca0fc69135a49bc2d138");
    assert!(challenge31::verify(path, &expected_sig.0, 0), "Invalid");

    // Don't attack because it interrupts the flow of the other challenges by taking long.
    // const DELAY: u64 = 20;
    // let signature = challenge31::_attack(path, DELAY);
    // assert_eq!(expected_sig, signature, "Recovery of HMAC-SHA1 failed");

    println!("[okay] Challenge 31: recovered HMAC-SHA1 via timing attack");
}

pub fn challenge32() {
    use std::time;

    pub fn _attack(file: &Path, delay: u64) -> Bytes {
        const CYCLES_PER_BYTE: usize = 30;
        const BLOCK_SIZE: usize = 20;
        let mut sig = vec![0x0; BLOCK_SIZE];
        for i in 0..BLOCK_SIZE {
            let mut max_tuple: (u128, u8) = (u128::MIN, 0);
            for c in 0_u8..=255_u8 {
                let mut cummulated_us = 0;
                for _ in 0..CYCLES_PER_BYTE {
                    let now = time::Instant::now();
                    sig[i] = c;
                    challenge31::verify(file, &sig, delay);
                    cummulated_us += now.elapsed().as_micros();
                }
                if cummulated_us > max_tuple.0 {
                    max_tuple = (cummulated_us, c);
                }
            }
            sig[i] = max_tuple.1;
        }
        Bytes(sig)
    }

    // Don't attack because it interrupts the flow of the other challenges by taking long.
    // const DELAY: u64 = 1;
    // let path = Path::new("data/12.txt");
    // let expected_sig = Bytes::from_hex("62f4527ea6cb716d0ad1ca0fc69135a49bc2d138");
    // assert!(
    //     challenge31::_attack(path, DELAY) != expected_sig,
    //     "Recovery was successful"
    // );
    // assert!(
    //     _attack(path, DELAY) == expected_sig,
    //     "Recovery was not successful"
    // );

    println!("[okay] Challenge 32: recovered HMAC-SHA1 with slightly less artificial timing leak");
}