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

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-2

const ROUND_CONSTS: [u32; 0x40] = [
    0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
    0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
    0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
    0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
    0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
    0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
    0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
    0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2
];

fn sha2(rounds: usize, msg: &[u8]) -> Vec<u8> {
    // Initialize
    let (mut h0, mut h1, mut h2, mut h3, mut h4, mut h5, mut h6, mut h7) =
        (0x6a09e667u32, 0xbb67ae85u32, 0x3c6ef372u32, 0xa54ff53au32, 0x510e527fu32, 0x9b05688cu32, 0x1f83d9abu32, 0x5be0cd19u32);

    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; 0x40];
        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..0x40 {
            let s0 = w[i - 15].rotate_right(7) ^ w[i - 15].rotate_right(18) ^ (w[i - 15] >> 3);
            let s1 = w[i - 2].rotate_right(17) ^ w[i - 2].rotate_right(19) ^ (w[i - 2] >> 10);
            w[i] = w[i - 16] + s0 + w[i - 7] + s1;
        }

        let (mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h) = (h0, h1, h2, h3, h4, h5, h6, h7);
        // Main loop
        for i in 0..rounds {
            let s1 = e.rotate_right(6) ^ e.rotate_right(11) ^ e.rotate_right(25);
            let ch = (e & f) ^ (!e & g);
            let temp1 = h + s1 + ch + ROUND_CONSTS[i]  + w[i];
            let s0 = a.rotate_right(2) ^ a.rotate_right(13) ^ a.rotate_right(22);
            let maj = (a & b) ^ (a & c) ^ (b & c);
            let temp2 = s0 + maj;

            h = g;
            g = f;
            f = e;
            e = d + temp1;
            d = c;
            c = b;
            b = a;
            a = temp1 + temp2;
        }

        // Add chunk's hash to total
        h0 += a; h1 += b; h2 += c; h3 += d; h4 += e; h5 += f; h6 += g; h7 += h;
    }

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

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

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

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

    fn info() -> String {
        format!("IN: msg, OUT: 256b SHA2 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 {
        0x100
    }

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

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

}

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

#[derive(Debug)]
pub struct InvSha2 {
    direct: Sha2
}

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

    fn info() -> String {
        format!("IN: 256b str, OUT: msg whose SHA2 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: Sha2{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 {
        0x100
    }

    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..0x100 {
            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(""), "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"),
        (String::from("The quick brown fox jumps over the lazy dog"), "d7a8fbb307d7809469ca9abcb0082e4f8d5651e46d3cdb762d02d0bf37c9e592")
    ];
    for (msg_str, true_hash_str) in pairs.iter() {
        let hash = sha2(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);
    }
}
