//! Implements Bitransform for MD5 and its (unknown...) inverse
//! See https://en.wikipedia.org/wiki/MD5

extern crate rand;

use rand::prelude::*;

#[cfg(test)]
use std::fmt::Write;

use super::{
    Bitransform,
    BitransformDesc
};

// ----------------------------------------------------------------
// Based on pseudocode from https://en.wikipedia.org/wiki/MD5

const ROUND_SHIFTS: [u32; 0x40] = [
     7, 12, 17, 22,  7, 12, 17, 22,  7, 12, 17, 22,  7, 12, 17, 22,
     5,  9, 14, 20,  5,  9, 14, 20,  5,  9, 14, 20,  5,  9, 14, 20,
     4, 11, 16, 23,  4, 11, 16, 23,  4, 11, 16, 23,  4, 11, 16, 23,
     6, 10, 15, 21,  6, 10, 15, 21,  6, 10, 15, 21,  6, 10, 15, 21
];

const ROUND_CONSTS: [u32; 0x40] = [
    0xd76aa478, 0xe8c7b756, 0x242070db, 0xc1bdceee,
    0xf57c0faf, 0x4787c62a, 0xa8304613, 0xfd469501,
    0x698098d8, 0x8b44f7af, 0xffff5bb1, 0x895cd7be,
    0x6b901122, 0xfd987193, 0xa679438e, 0x49b40821,
    0xf61e2562, 0xc040b340, 0x265e5a51, 0xe9b6c7aa,
    0xd62f105d, 0x02441453, 0xd8a1e681, 0xe7d3fbc8,
    0x21e1cde6, 0xc33707d6, 0xf4d50d87, 0x455a14ed,
    0xa9e3e905, 0xfcefa3f8, 0x676f02d9, 0x8d2a4c8a,
    0xfffa3942, 0x8771f681, 0x6d9d6122, 0xfde5380c,
    0xa4beea44, 0x4bdecfa9, 0xf6bb4b60, 0xbebfbc70,
    0x289b7ec6, 0xeaa127fa, 0xd4ef3085, 0x04881d05,
    0xd9d4d039, 0xe6db99e5, 0x1fa27cf8, 0xc4ac5665,
    0xf4292244, 0x432aff97, 0xab9423a7, 0xfc93a039,
    0x655b59c3, 0x8f0ccc92, 0xffeff47d, 0x85845dd1,
    0x6fa87e4f, 0xfe2ce6e0, 0xa3014314, 0x4e0811a1,
    0xf7537e82, 0xbd3af235, 0x2ad7d2bb, 0xeb86d391
];

fn md5(rounds: usize, msg: &[u8]) -> Vec<u8> {
    // Initialize
    let (mut h0, mut h1, mut h2, mut h3) =
        (0x67452301, 0xefcdab89, 0x98badcfe, 0x10325476);

    let length0 = (msg.len() << 3) as u64; // original length in bits

    // Padding
    let mut msg_pad = msg.to_vec();
    msg_pad.push(0x80);
    let length1 = msg_pad.len();
    if (length1 & 0x3F) != 0x38 { // pad with zeros
        msg_pad.extend(vec![0u8; (((length1 + 8) >> 6) << 6) + 0x38 - length1]);
    }
    let mut length0_bytes = vec![];
    length0_bytes.extend((0..8).into_iter().map(|i| {
        ((length0 >> (i << 3)) & 0xFF) as u8
    }));
    msg_pad.extend(length0_bytes);

    // Process each 64-byte chunk
    for c in 0..(msg_pad.len() >> 6) {
        let chunk = & msg_pad[(c << 6)..((c + 1) << 6)];
        let mut w = [0u32; 0x10];
        for i in 0..0x10 {
            let w_bytes = & chunk[(i << 2)..((i + 1) << 2)];
            w[i] = (0..4).fold(0u32, |acc, i| acc | (w_bytes[i] as u32) << (i << 3) ); // little-endian
        }

        let (mut a, mut b, mut c, mut d) = (h0, h1, h2, h3);
        // Main loop
        for i in 0..rounds {
            let (mut f, g) = match i {
                0..=0xF => (d ^ (b & (c ^ d)), i),
                0x10..=0x1F => (c ^ (d & (b ^ c)), (5 * i + 1) & 0xF),
                0x20..=0x2F => (b ^ c ^ d, (3 * i + 5) & 0xF),
                _ => (c ^ (b | !d), (7 * i) & 0xF)
            };

            f += a + ROUND_CONSTS[i] + w[g];
            a = d;
            d = c;
            c = b;
            b += f.rotate_left(ROUND_SHIFTS[i]);
        }

        // Add chunk's hash to total
        h0 += a; h1 += b; h2 += c; h3 += d;
    }

    // Produce final hash
    let (h0, h1, h2, h3) = (h0.to_le_bytes(), h1.to_le_bytes(), h2.to_le_bytes(), h3.to_le_bytes());
    vec![h0, h1, h2, h3].concat()
}

// ----------------------------------------------------------------

#[derive(Debug)]
pub struct Md5 {
    rounds: usize,
    msg_bytes: usize
}

impl Bitransform for Md5 {
    fn id() -> String {
        String::from("Md5")
    }

    fn info() -> String {
        format!("IN: msg, OUT: 128b MD5 of IN, PREFS: rounds (1-64), msg_bytes (>=1)")
    }

    fn new(prefs: &[usize]) -> Result<Box<dyn Bitransform>, String> {
        match prefs.len() {
            0 ..= 1 => Err(format!("Not enough prefs (must be 2)")),
            2 => Ok(Box::new(Self{rounds: prefs[0], msg_bytes: prefs[1]})),
            _ => Err(format!("Too many prefs (must be 2)"))
        }
    }

    fn desc(&self) -> BitransformDesc {
        BitransformDesc::new(Self::id(), &[self.rounds, self.msg_bytes])
    }

    fn inp_size(&self) -> usize {
        self.msg_bytes << 3
    }

    fn outp_size(&self) -> usize {
        0x80
    }

    fn calc_func(&self, args: &[u128]) -> Vec<u128> {
        let mut yb = vec![0u128; 0x80];
        for i in 0..0x80 {
            let mut msg = vec![0u8; args.len() >> 3];
            (0..args.len()).for_each(|j|
                msg[j >> 3] |= (((args[j] >> i) & 1) as u8) << (j & 7)
            );
            let y = md5(self.rounds, &msg);
            (0..0x80).for_each(|j|
                yb[j] |= (((y[j >> 3] >> (j & 7)) & 1) as u128) << i
            );
        }
        yb
    }

    fn rand_inp(&self, rng: &mut ThreadRng) -> Vec<u128> {
        let mut msgb = Vec::<u128>::with_capacity(self.inp_size());
        for _ in 0..self.inp_size() {
            msgb.push(rng.gen::<u128>())
        }
        msgb
    }

    fn score_inp_outp(&self, inp: &[u128], outp: &[u128]) -> f64 {
        let func = self.calc_func(inp);
        let mut score: i32 = 0;

        for i in 0..0x80 {
            score -= (func[i] ^ outp[i]).count_ones() as i32; // number of different bits
        }

        (score as f64) / (0x80 as f64)
    }

}

// ----------------------------------------------------------------

#[derive(Debug)]
pub struct InvMd5 {
    direct: Md5
}

impl Bitransform for InvMd5 {
    fn id() -> String {
        String::from("InvMd5")
    }

    fn info() -> String {
        format!("IN: 128b str, OUT: msg whose MD5 is IN, PREFS: rounds (1-64), msg_bytes (>=1)")
    }

    fn new(prefs: &[usize]) -> Result<Box<dyn Bitransform>, String> {
        match prefs.len() {
            0 ..= 1 => Err(format!("Not enough prefs (must be 2)")),
            2 => Ok(Box::new(Self{direct: Md5{rounds: prefs[0], msg_bytes: prefs[1]}})),
            _ => Err(format!("Too many prefs (must be 2)"))
        }
    }

    fn desc(&self) -> BitransformDesc {
        BitransformDesc::new(Self::id(), &[self.direct.rounds, self.direct.msg_bytes])
    }

    fn inp_size(&self) -> usize {
        0x80
    }

    fn outp_size(&self) -> usize {
        self.direct.msg_bytes << 3
    }

    fn calc_func(&self, args: &[u128]) -> Vec<u128> {
        self.direct.calc_func(args)
    }

    fn rand_inp(&self, rng: &mut ThreadRng) -> Vec<u128> {
        self.direct.calc_func(& self.direct.rand_inp(rng))
    }

    fn score_inp_outp(&self, inp: &[u128], outp: &[u128]) -> f64 {
        let func = self.direct.calc_func(outp);
        let mut score: i32 = 0;

        for i in 0..0x80 {
            score -= (func[i] ^ inp[i]).count_ones() as i32; // number of different bits
        }

        (score as f64) / (0x80 as f64)
    }

}

// ----------------------------------------------------------------

#[test]
fn test() {
    let pairs = [
        (String::from(""), "d41d8cd98f00b204e9800998ecf8427e"),
        (String::from("The quick brown fox jumps over the lazy dog"), "9e107d9d372bb6826bd81d3542a419d6")
    ];
    for (msg_str, true_hash_str) in pairs.iter() {
        let hash = md5(64, msg_str.as_bytes());
        let mut hash_str = String::from("");
        hash.into_iter().for_each(|b| {
            let _ = write!(hash_str, "{:02x}", b);
        });
        assert_eq!(hash_str.as_str(), *true_hash_str);
    }
}
