use crate::bytes::Bytes;
use crate::bytes_base64::BytesBase64;
use crate::cbc;
use crate::ecb;
use crate::parser;
use rand::Rng;
use std::collections::HashMap;

pub fn challenge9() {
    let mut bytes = Bytes::from_utf8("YELLOW SUBMARINE");
    bytes.pad_pkcs7(20);
    println!("[okay] Challenge 9: {:?}", bytes.to_utf8());
}

pub fn challenge10() {
    fn read(path: &str) -> Bytes {
        let s = std::fs::read_to_string(path).unwrap();
        BytesBase64::from_base64(&s).to_bytes()
    }
    let iv = Bytes(vec![0; 16]);
    let key = Bytes::from_utf8("YELLOW SUBMARINE");

    let text = Bytes::from_utf8("aaaabbbbccccddddeeeeffffgggghhhh");
    let ciphertext = cbc::encrypt(&key, &iv, &text);
    let roundtrip = cbc::decrypt(&key, &iv, &ciphertext);
    if text == roundtrip {
        let ciphertext = read("data/10.txt");
        let cleartext = cbc::decrypt(&key, &iv, &ciphertext);
        let output = cleartext.to_utf8()[..16].to_string();
        println!("[okay] Challenge 10: {}", output);
    } else {
        println!("[fail] Challenge 10: rountrip failed");
    }
}

pub fn challenge11() {
    #[derive(Debug, PartialEq)]
    enum EncryptionType {
        CBC,
        ECB,
    }

    fn pad_data(mut v: Vec<u8>) -> Bytes {
        let mut pre_pad = Bytes::random_range(5, 10).0;
        let mut post_pad = Bytes::random_range(5, 10).0;
        pre_pad.append(&mut v);
        pre_pad.append(&mut post_pad);
        Bytes(pre_pad)
    }

    fn encryption_oracle(Bytes(data): &Bytes) -> (Bytes, EncryptionType) {
        // Write a function that encrypts data under an unknown key --- that is, a
        // function that generates a random key and encrypts under it.
        let key = Bytes::random(16);
        // Under the hood, have the function append 5-10 bytes (count chosen randomly)
        // before the plaintext and 5-10 bytes after the plaintext.
        let padded_data = pad_data(data.to_vec());
        // Now, have the function choose to encrypt under ECB 1/2 the time, and under CBC
        // the other half (just use random IVs each time for CBC). Use rand(2) to decide
        // which to use.
        let zero_or_one: u32 = rand::thread_rng().gen_range(0..2);
        let (data, encryption_type) = if zero_or_one == 1 {
            (ecb::encrypt(&key, &padded_data), EncryptionType::ECB)
        } else {
            let iv = Bytes::random(16);
            (cbc::encrypt(&key, &iv, &padded_data), EncryptionType::CBC)
        };
        (data, encryption_type)
    }

    fn cbcecb_detection_oracle(data: &Bytes) -> EncryptionType {
        // Detect the block cipher mode the function is using each time. You should end up
        // with a piece of code that, pointed at a block box that might be encrypting ECB
        // or CBC, tells you which one is happening.
        if data.has_duplicated_cycle(16) {
            EncryptionType::ECB
        } else {
            EncryptionType::CBC
        }
    }

    fn run_oracle(count: usize) {
        // I struggled for a while to understand how the detection oracle can
        // work.  The issue is that if we provide a random text to encrypt, we
        // won't be able to use duplicated blocks as a method for detecting the
        // cipher. I don't think there is a way around that, but if we can choose
        // the plaintext it becomes trivial.
        let text = Bytes::from_utf8("aaaabbbbccccddddaaaabbbbccccddddaaaabbbbccccdddd");
        let mut correct: usize = 0;
        for _ in 0..count {
            let (ciphertext, et) = encryption_oracle(&text);
            let get = cbcecb_detection_oracle(&ciphertext);
            if et == get {
                correct += 1;
            }
        }
        if correct == count {
            println!("[okay] Challenge 11: [{} / {}]", correct, count);
        } else {
            println!("[fail] Challenge 11: [{} / {}]", correct, count);
        }
    }
    run_oracle(10);
}

pub fn challenge12() {
    fn read(path: &str) -> Bytes {
        let s = std::fs::read_to_string(path).unwrap();
        BytesBase64::from_base64(&s).to_bytes()
    }

    fn encryption_oracle(key: &Bytes, Bytes(data): &Bytes) -> Bytes {
        // Copy your oracle function to a new function that encrypts buffers under ECB mode using a consistent but unknown key
        // Now take that same function and have it append to the plaintext, BEFORE ENCRYPTING, the following string (from 12.txt):
        let mut data = data.to_vec();
        let mut string = read("data/12.txt");
        data.append(&mut string.0);
        let cipher = ecb::encrypt(&key, &Bytes(data));
        cipher
    }

    fn get_block_size(key: &Bytes) -> usize {
        // Detect cipher block size
        let mut v = vec![];
        let initial_cipher_len = encryption_oracle(&key, &Bytes(v.to_vec())).0.len();
        let mut new_cipher_len = initial_cipher_len;
        while initial_cipher_len == new_cipher_len {
            v.push(b'A');
            let cipher = encryption_oracle(&key, &Bytes(v.to_vec()));
            new_cipher_len = cipher.0.len();
        }
        let key_length = new_cipher_len - initial_cipher_len;
        key_length
    }

    fn is_encryption_ecb(key: &Bytes) -> bool {
        let data = Bytes::from_utf8("aaaabbbbccccddddaaaabbbbccccdddd");
        let cipher = encryption_oracle(&key, &data);
        if cipher.has_duplicated_cycle(16) {
            true
        } else {
            false
        }
    }

    fn decode(key: &Bytes) -> Bytes {
        let block_size = get_block_size(&key);
        let block_count = encryption_oracle(&key, &Bytes(vec![])).0.len() / block_size;
        let mut cleartext = vec![];

        for block_index in 0..block_count {
            let mut cleartext_block = vec![];
            for padding_length in (0..block_size).rev() {
                let padding_text = vec![b'-'; padding_length];
                let expected = encryption_oracle(&key, &Bytes(padding_text.to_vec()));
                let expected_block = expected.get_block(block_index, block_size);

                let mut known_text = if block_index == 0 {
                    padding_text.to_vec()
                } else {
                    let cleartext_offset =
                        ((block_index - 1) * block_size) + (block_size - padding_length);
                    cleartext[cleartext_offset..(cleartext_offset + padding_length)].to_vec()
                };
                known_text.append(&mut cleartext_block.to_vec());

                for i in 0..255 {
                    let mut guess_text = known_text.to_vec();
                    guess_text.push(i);
                    let cipher_block =
                        encryption_oracle(&key, &Bytes(guess_text)).get_block(0, block_size);
                    if cipher_block.0 == expected_block.0 {
                        cleartext_block.push(i);
                        break;
                    }
                }
            }
            cleartext.append(&mut cleartext_block);
        }
        Bytes(cleartext)
    }

    let key = Bytes::random(16); // consistent but unknown key
    assert_eq!(get_block_size(&key), 16); // 1. discover block size
    assert_eq!(is_encryption_ecb(&key), true); // 2. confirm oracle uses ecb
    let roundtrip_text = decode(&key); // 3.-6.
    let cleartext = read("data/12.txt");

    // 138 (instead of 139); I think we get one additional byte because we guess
    // the first padding byte. The right approach would be to remove the last
    // byte, encrypt it, and then compare it to the result of the encryption
    // oracle, but this approach is fine too.
    assert_eq!(roundtrip_text.0[..138], cleartext.0);
    println!(
        "[okay] Challenge 12: {}",
        roundtrip_text.to_utf8()[..17].to_string()
    );
}

pub fn challenge13() {
    fn profile_for(input: &str, key: &Bytes) -> Bytes {
        let mut r = String::new();
        for c in input.chars() {
            if !(c.is_ascii_alphabetic() || c == '.' || c == '@') {
                panic!("profile_for: invalid char {}", c);
            }
        }
        r.push_str("email=");
        r.push_str(input);
        r.push_str("&uid=1337&role=user");
        ecb::encrypt(&key, &Bytes(r.as_bytes().to_vec()))
    }

    fn decrypt(key: &Bytes, data: &Bytes) -> HashMap<String, String> {
        let c = ecb::decrypt(&key, &data);
        parser::parse_key_value(&c.to_utf8())
    }

    fn attack(key: &Bytes) -> Bytes {
        // Using only the user input to profile_for() (as an oracle to generate
        // "valid" ciphertexts) and the ciphertexts themselves, make a
        // role=admin profile.
        // (FelixM) I assume ECB and block_size = 16; we could figure
        // it out easily by adding enough 'a' to the email
        let mut r = vec![];

        //             ________________________________
        //             0..34..78..bc..f0..34..78..bc..f0..34..78..bc..f
        //             email=aaaaa@a.com&uid=1337&role=user
        let p = profile_for("aaaaa@a.com", &key);
        r.append(&mut p.0[0..32].to_vec());

        //                             ----------------
        //             0..34..78..bc..f0..34..78..bc..f0..34..78..bc..f
        //             email=aaaaaaa@a.admin&uid=1337&role=user
        let p = profile_for("aaaaaaa@a.admin", &key);
        r.append(&mut p.0[16..32].to_vec());

        //                                             ----------------
        //             0..34..78..bc..f0..34..78..bc..f0..34..78..bc..f
        //             email=aaaaaaaa@a.admin&uid=1337&role=user
        let p = profile_for("aaaaaaaa@a.admin", &key);
        r.append(&mut p.0[32..48].to_vec());

        Bytes(r)
    }

    let key = Bytes::random(16); // consistent but unknown key
    let profile = attack(&key);
    let dict = decrypt(&key, &profile);
    let role = dict.get("role").unwrap();
    assert_eq!(role, "admin");
    println!("[okay] Challenge 13: role={}", role);
}

pub fn challenge14() {
    fn read(path: &str) -> Bytes {
        let s = std::fs::read_to_string(path).unwrap();
        BytesBase64::from_base64(&s).to_bytes()
    }

    fn encryption_oracle(
        Bytes(random_prefix): &Bytes,
        random_key: &Bytes,
        Bytes(attacker_controlled): &Bytes,
    ) -> Bytes {
        // AES-128-ECB(random-prefix || attacker-controlled || target-bytes, random-key)
        // Thoughts: If I generate a random prefix for every encryption, then I don't
        // know how to decode it because I cannot really run experiments. If I generate
        // a single random prefix it becomes rather trivial. I just have to find out how
        // long it is and then adjust the decoding routine.
        let mut plaintext = random_prefix.to_vec();
        plaintext.append(&mut attacker_controlled.to_vec());
        let mut target_bytes = read("data/12.txt").0;
        plaintext.append(&mut target_bytes);
        let cipher = ecb::encrypt(&random_key, &Bytes(plaintext));
        cipher
    }

    fn get_block_size(prefix: &Bytes, key: &Bytes) -> usize {
        // Detect cipher block size this approach also confirms that it is ECB
        let v = vec![b'a'; 256];
        let cipher = encryption_oracle(prefix, key, &Bytes(v));

        for i in 1..10 {
            let block_size = i * 16;
            match get_duplicated_block_indices(&cipher, block_size) {
                Some((_, _)) => {
                    return block_size;
                }
                _ => (),
            }
        }
        0
    }

    fn get_duplicated_block_indices(cipher: &Bytes, block_size: usize) -> Option<(usize, usize)> {
        let chunks: Vec<&[u8]> = cipher.0.chunks(block_size).collect();
        for i in 0..chunks.len() {
            for j in (i + 1)..chunks.len() {
                if chunks[i] == chunks[j] {
                    return Some((i, j));
                }
            }
        }
        None
    }

    fn get_prefix_size(prefix: &Bytes, key: &Bytes) -> usize {
        let block_size = get_block_size(&prefix, &key);
        let duplicated_text = Bytes::from_utf8("aaaabbbbccccddddaaaabbbbccccdddd").0;

        for i in 0..block_size {
            let mut padding = vec![b'a'; i];
            padding.append(&mut duplicated_text.to_vec());
            let cipher = encryption_oracle(prefix, key, &Bytes(padding));
            match get_duplicated_block_indices(&cipher, block_size) {
                Some((first_block, _)) => {
                    return block_size * first_block - i;
                }
                _ => (),
            }
        }
        0
    }

    fn decode(prefix: &Bytes, key: &Bytes) -> Bytes {
        let block_size = get_block_size(prefix, key);
        let prefix_size = get_prefix_size(prefix, key);
        let prefix_padding_size = block_size - (prefix_size % block_size);
        let prefix_padding = Bytes(vec![b'a'; prefix_padding_size]);
        let block_count = encryption_oracle(&prefix, &key, &prefix_padding).len() / block_size;
        let first_block_index = (prefix_size + prefix_padding_size) / block_size;
        let mut cleartext = vec![];

        // AES-128-ECB(random-prefix || attacker-controlled || target-bytes, random-key)
        // rrrrrrrrrrrrrrrrrrrpppppppppppppaaaaaaaaaaaaaaathe real text
        // 0..34..78..bc..f0..34..78..bc..f0..34..78..bc..f0..34..78..bc..f0..34..78..bc..f
        // 0               1  ^            2               4               5
        //                    |            ^ first_block_index
        //                    \ prefix_padding
        for block_index in first_block_index..block_count {
            let mut cleartext_block = vec![];
            for padding_length in (0..block_size).rev() {
                let full_padding_text = vec![b'-'; prefix_padding_size + padding_length];
                let expected = encryption_oracle(prefix, key, &Bytes(full_padding_text.to_vec()));
                let expected_block = expected.get_block(block_index, block_size);

                let mut known_text = if block_index == first_block_index {
                    full_padding_text.to_vec()
                } else {
                    let mut prefix_padding = vec![b'-'; prefix_padding_size];
                    let cleartext_offset = ((block_index - first_block_index - 1) * block_size)
                        + (block_size - padding_length);
                    prefix_padding.append(
                        &mut cleartext[cleartext_offset..(cleartext_offset + padding_length)]
                            .to_vec(),
                    );
                    prefix_padding
                };
                known_text.append(&mut cleartext_block.to_vec());

                for i in 0..255 {
                    let mut guess_text = known_text.to_vec();
                    guess_text.push(i);
                    let cipher_block = encryption_oracle(prefix, key, &Bytes(guess_text))
                        .get_block(first_block_index, block_size);
                    if cipher_block.0 == expected_block.0 {
                        cleartext_block.push(i);
                        break;
                    }
                }
            }
            cleartext.append(&mut cleartext_block);
        }
        // We get last byte from cbs padding so remove it.
        Bytes(cleartext[0..(cleartext.len() - 1)].to_vec())
    }

    let prefix = Bytes::random_range(0, 200);
    let key = Bytes::random(16);
    assert_eq!(get_block_size(&prefix, &key), 16);
    let prefix_len = get_prefix_size(&prefix, &key);
    assert_eq!(prefix.len(), prefix_len);
    let roundtrip_text = decode(&prefix, &key);
    let cleartext = read("data/12.txt");
    assert_eq!(roundtrip_text, cleartext);
    println!(
        "[okay] Challenge 14: {}",
        roundtrip_text.to_utf8()[..17].to_string()
    );
}

pub fn challenge15() {
    assert_eq!(
        Bytes::from_utf8("ICE ICE BABY\u{4}\u{4}\u{4}").has_valid_pkcs7(16),
        false
    );
    assert_eq!(
        Bytes::from_utf8("ICE ICE BABY\u{4}\u{4}\u{4}\u{4}").has_valid_pkcs7(16),
        true
    );
    assert_eq!(
        Bytes::from_utf8("ICE ICE BABY\u{3}\u{3}\u{4}\u{4}").has_valid_pkcs7(16),
        false
    );
    assert_eq!(
        Bytes::from_utf8("ICE ICE BABY!!!\u{0}").has_valid_pkcs7(16),
        false
    );
    assert_eq!(
        Bytes::from_utf8("ICE ICE BABY!!!\u{1}").has_valid_pkcs7(16),
        true
    );
    let mut bytes = Bytes::from_utf8("ICE ICE BABY\u{3}\u{3}\u{4}\u{4}");
    bytes.pad_pkcs7(16);
    assert_eq!(bytes.has_valid_pkcs7(16), true);
    println!("[okay] Challenge 15: PKCS7 works");
}

pub fn challenge16() {
    fn encrypt(input: &str, key: &Bytes, iv: &Bytes) -> Bytes {
        let mut r = String::new();
        for c in input.chars() {
            if c == ';' || c == '=' {
                panic!("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 mut cleartext = Bytes(r.as_bytes().to_vec());
        cleartext.pad_pkcs7(16);
        cbc::encrypt(&key, &iv, &cleartext)
    }

    let iv = Bytes::random(16);
    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=xxx=x;admin=true;comment2=%20like%20a%20pound%20of%20bacon
    // flip_bit(2, '9') = '='; flip_bit(1, '9') = ';'
    let mut cipher = encrypt("xxx9x9admin9true", &key, &iv);
    cipher.flip_bit(19, 2); // results in flipping same bit in byte of the following block
    cipher.flip_bit(21, 1);
    cipher.flip_bit(27, 2);
    let mut cleartext = cbc::decrypt(&key, &iv, &cipher);
    cleartext.remove_pkcs7(16);
    let cleartext_stripped = Bytes(cleartext.0[32..].to_vec());
    let dict = parser::parse_key_value(&cleartext_stripped.to_utf8());
    let admin_status = dict.get("admin").unwrap();
    assert_eq!(admin_status, "true");
    println!("[okay] Challenge 16: admin={}", admin_status);
}