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

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/SHA-1

fn sha1(rounds: usize, msg: &[u8]) -> Vec<u8> {
    // Initialize
    let (mut h0, mut h1, mut h2, mut h3, mut h4) =
        (0x67452301u32, 0xEFCDAB89u32, 0x98BADCFEu32, 0x10325476u32, 0xC3D2E1F0u32);

    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
    }));
    length0_bytes.reverse(); // big-endian
    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; 0x50];
        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[3 - i] as u32) << (i << 3) ); // big-endian
        }
        // Extend
        for i in 0x10..0x50 {
            w[i] = (w[i - 3] ^ w[i - 8] ^ w[i - 14] ^ w[i - 16]).rotate_left(1);
        }

        let (mut a, mut b, mut c, mut d, mut e) = (h0, h1, h2, h3, h4);
        // Main loop
        for i in 0..rounds {
            let (f, k) = match i {
                0..=19 => ((b & c) | ((!b) & d), 0x5A827999),
                20..=39 => ((b ^ c ^ d), 0x6ED9EBA1),
                40..=59 => ((b & c) | (b & d) | (c & d), 0x8F1BBCDC),
                _ => ((b ^ c ^ d), 0xCA62C1D6)
            };

            let temp = a.rotate_left(5) + f + e + k + w[i];
            e = d;
            d = c;
            c = b.rotate_left(30);
            b = a;
            a = temp;
        }

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

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

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

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

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

    fn info() -> String {
        format!("IN: msg, OUT: 160b SHA1 of IN, PREFS: rounds (1-80), 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 {
        0xA0
    }

    fn calc_func(&self, args: &[u128]) -> Vec<u128> {
        let mut yb = vec![0u128; 0xA0];
        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 = sha1(self.rounds, &msg);
            (0..0xA0).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..0xA0 {
            score -= (func[i] ^ outp[i]).count_ones() as i32; // number of different bits
        }

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

}

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

#[derive(Debug)]
pub struct InvSha1 {
    direct: Sha1
}

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

    fn info() -> String {
        format!("IN: 160b str, OUT: msg whose SHA1 is IN, PREFS: rounds (1-80), 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: Sha1{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 {
        0xA0
    }

    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..0xA0 {
            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(""), "da39a3ee5e6b4b0d3255bfef95601890afd80709"),
        (String::from("The quick brown fox jumps over the lazy dog"), "2fd4e1c67a2d28fced849ee1bb76e7391b93eb12")
    ];
    for (msg_str, true_hash_str) in pairs.iter() {
        let hash = sha1(80, 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);
    }
}
