// License: see LICENSE file at root directory of `master` branch

#![cfg(target_endian="little")]

//! # Keccak

mod bit_str;
mod bits;
mod hash;
mod state;
mod tests;

use {
    alloc::{
        string::String,
        vec::Vec,
    },
    core::{
        convert::TryInto,
        mem,
    },

    self::{
        bit_str::BitStr,
        bits::Bits,
        state::State,
    },
};

#[cfg(feature="std")]
use {
    std::io::Write,

    crate::IoResult,
};

pub use self::hash::Hash;

/// # Keccak
///
/// ## Examples
///
/// ```
/// use zeros::Hash;
///
/// let mut keccak = Hash::Shake128.new_keccak();
/// keccak.update("test");
/// assert_eq!(
///     keccak.finish(),
///     &[
///         0xd3, 0xb0, 0xaa, 0x9c, 0xd8, 0xb7, 0x25, 0x56,
///         0x22, 0xce, 0xbc, 0x63, 0x1e, 0x86, 0x7d, 0x40,
///     ],
/// );
/// ```
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct Keccak {
    hash: Hash,
    bits: Bits,
    buf: Vec<u8>,
    buf_chunk: usize,
    output: BitStr,
}

impl Keccak {

    const BUF_CHUNK: usize = 8;

    /// # Makes new instance
    pub (crate) fn new(hash: Hash) -> Self {
        let bits = Bits::B1600;
        let buf = Vec::with_capacity(168);
        let buf_chunk = hash.rate() / 8;
        let output = bits.new_bit_str();

        Self {
            hash,
            bits,
            buf,
            buf_chunk,
            output,
        }
    }

    /// # Gets hash
    pub const fn hash(&self) -> &Hash {
        &self.hash
    }

    /// # Updates new data
    ///
    /// Returns length of input bytes.
    pub fn update<B>(&mut self, bytes: B) -> usize where B: AsRef<[u8]> {
        let result = {
            let bytes = bytes.as_ref();
            self.buf.extend(bytes);
            bytes.len()
        };

        while self.buf.len() >= self.buf_chunk {
            // Data
            let chunks = self.buf[..self.buf_chunk].chunks_exact(Self::BUF_CHUNK);
            if chunks.remainder().is_empty() == false {
                // Buf chunk is calculated from hash's rate -- which is tested that: rate % 64 == 0
                // So this should never happen. But just in case, panic.
                panic!("Invalid chunks");
            }
            for (idx, c) in chunks.enumerate() {
                self.output.mut_data()[idx] ^= u64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]]);
            }
            self.buf.drain(..self.buf_chunk);

            // Pad
            let pad_idx = self.hash.capacity() / 8;
            let data_len = self.output.data().len();
            for i in data_len - pad_idx .. data_len {
                self.output.mut_data()[i] ^= 0;
            }

            // Update
            self.keccak_p_1600_24();
        }

        result
    }

    /// # Calls `keccak_p()` with `1600` bits and `24` rounds
    fn keccak_p_1600_24(&mut self) {
        keccak_p(&Bits::B1600, 24, &mut self.output);
    }

    /// # Finishes and returns hash
    pub fn finish(mut self) -> Vec<u8> {
        let mut bit_str = BitStr::from(self.buf);
        for b in self.hash.suffix() {
            bit_str += *b;
        }
        // On start, buf has no data. If the user calls update(), buf will be handled in chunks. So here it should not be large.
        // After padding: data length = capacity + rate. If the code compiles, this unwrap() should work.
        bit_str.pad(&self.hash).unwrap();

        // Padding will make data length a positive multiple of rate. While rate is tested that: rate % 64 == 0
        // So unwrap() should work. But just in case, panic.
        self.buf = bit_str.try_into().unwrap();
        self.update(&[]);

        let output_bytes = self.hash.output_bytes();
        let mut result = Vec::with_capacity(output_bytes);
        loop {
            let mut count = (self.hash.rate() / 8).min(output_bytes.saturating_sub(result.len()));
            for u in self.output.data() {
                result.extend(&u.to_le_bytes());
                match count.checked_sub(mem::size_of::<u64>()) {
                    Some(0) | None => break,
                    Some(other) => count = other,
                };
            }

            if output_bytes <= result.len() {
                result.truncate(output_bytes);
                return result;
            }

            self.keccak_p_1600_24();
        }
    }

    /// # Finishes and returns hash as a hexadecimal string, in lower-case
    pub fn finish_as_hex(self) -> String {
        crate::bytes_to_hex(self.finish())
    }

}

impl From<Hash> for Keccak {

    fn from(hash: Hash) -> Self {
        Self::new(hash)
    }

}

impl From<&Hash> for Keccak {

    fn from(hash: &Hash) -> Self {
        Self::new(hash.clone())
    }

}

#[cfg(feature="std")]
impl Write for Keccak {

    fn write(&mut self, buf: &[u8]) -> IoResult<usize> {
        Ok(self.update(buf))
    }

    fn flush(&mut self) -> IoResult<()> {
        Ok(())
    }

}

/// # (See references)
fn keccak_p(bits: &Bits, rounds: u8, bit_str: &mut BitStr) {
    let mut state = State::new(bit_str, bits);

    let end = 12 + 2 * state.bits().l();
    for round_index in end - rounds .. end {
        state.theta();
        state.rho_and_pi();
        state.chi();
        state.iota(round_index);
    }
}

/// # (See references)
#[allow(unused)]
fn keccak_f(bits: &Bits, bit_str: &mut BitStr) {
    let rounds = 12 + 2 * bits.l();
    keccak_p(bits, rounds, bit_str);
}

#[test]
fn test_keccak() {
    assert_eq!(Keccak::BUF_CHUNK, 8);

    for h in &[Hash::Sha3_224, Hash::Sha3_256, Hash::Sha3_384, Hash::Sha3_512, Hash::Shake128, Hash::Shake256] {
        let keccak = h.new_keccak();
        assert!(keccak.buf_chunk > 0 && keccak.buf_chunk % Keccak::BUF_CHUNK == 0);
    }
}

#[test]
fn test_sha3_functions() {
    // Samples: https://csrc.nist.gov/projects/cryptographic-standards-and-guidelines/example-values

    for (data, result) in &[
        (
            &[][..],
            [
                0x6b, 0x4e, 0x03, 0x42, 0x36, 0x67, 0xdb, 0xb7, 0x3b, 0x6e, 0x15, 0x45, 0x4f, 0x0e,
                0xb1, 0xab, 0xd4, 0x59, 0x7f, 0x9a, 0x1b, 0x07, 0x8e, 0x3f, 0x5b, 0x5a, 0x6b, 0xc7,
            ],
        ),
        (
            b"haha",
            [
                0x8c, 0x3c, 0xfd, 0x38, 0x2b, 0xd3, 0x65, 0x8a, 0x97, 0xb8, 0x64, 0x99, 0x25, 0x5b,
                0x15, 0x17, 0xf1, 0x88, 0x0b, 0xb9, 0xc7, 0x0e, 0x27, 0x39, 0x1f, 0x58, 0x32, 0xa7,
            ],
        )
    ] {
        assert_eq!(Hash::Sha3_224.hash(data), result);
    }

    assert_eq!(Hash::Sha3_256.hash(&[]), &[
        0xa7, 0xff, 0xc6, 0xf8, 0xbf, 0x1e, 0xd7, 0x66, 0x51, 0xc1, 0x47, 0x56, 0xa0, 0x61, 0xd6, 0x62,
        0xf5, 0x80, 0xff, 0x4d, 0xe4, 0x3b, 0x49, 0xfa, 0x82, 0xd8, 0x0a, 0x4b, 0x80, 0xf8, 0x43, 0x4a,
    ]);
    assert_eq!(Hash::Sha3_384.hash(&[]), &[
        0x0c, 0x63, 0xa7, 0x5b, 0x84, 0x5e, 0x4f, 0x7d, 0x01, 0x10, 0x7d, 0x85, 0x2e, 0x4c, 0x24, 0x85,
        0xc5, 0x1a, 0x50, 0xaa, 0xaa, 0x94, 0xfc, 0x61, 0x99, 0x5e, 0x71, 0xbb, 0xee, 0x98, 0x3a, 0x2a,
        0xc3, 0x71, 0x38, 0x31, 0x26, 0x4a, 0xdb, 0x47, 0xfb, 0x6b, 0xd1, 0xe0, 0x58, 0xd5, 0xf0, 0x04,
    ]);
    assert_eq!(Hash::Sha3_512.new_keccak().finish(), &[
        0xa6, 0x9f, 0x73, 0xcc, 0xa2, 0x3a, 0x9a, 0xc5, 0xc8, 0xb5, 0x67, 0xdc, 0x18, 0x5a, 0x75, 0x6e,
        0x97, 0xc9, 0x82, 0x16, 0x4f, 0xe2, 0x58, 0x59, 0xe0, 0xd1, 0xdc, 0xc1, 0x47, 0x5c, 0x80, 0xa6,
        0x15, 0xb2, 0x12, 0x3a, 0xf1, 0xf5, 0xf9, 0x4c, 0x11, 0xe3, 0xe9, 0x40, 0x2c, 0x3a, 0xc5, 0x58,
        0xf5, 0x00, 0x19, 0x9d, 0x95, 0xb6, 0xd3, 0xe3, 0x01, 0x75, 0x85, 0x86, 0x28, 0x1d, 0xcd, 0x26,
    ]);
}

#[test]
fn test_shake_functions() {
    // Samples: https://csrc.nist.gov/projects/cryptographic-standards-and-guidelines/example-values

    for (data, result) in &[
        (
            &[][..],
            [0x7f, 0x9c, 0x2b, 0xa4, 0xe8, 0x8f, 0x82, 0x7d, 0x61, 0x60, 0x45, 0x50, 0x76, 0x05, 0x85, 0x3e],
        ),
        (
            b"haha",
            [0x05, 0xea, 0x99, 0x34, 0x23, 0x4a, 0x88, 0xa3, 0x1e, 0x67, 0xa1, 0x63, 0x8a, 0x79, 0x36, 0xd7],
        )
    ] {
        assert_eq!(Hash::Shake128.hash(data), result);
    }

    for (data, result) in &[
        (
            &[][..],
            [
                0x46, 0xb9, 0xdd, 0x2b, 0x0b, 0xa8, 0x8d, 0x13, 0x23, 0x3b, 0x3f, 0xeb, 0x74, 0x3e, 0xeb, 0x24,
                0x3f, 0xcd, 0x52, 0xea, 0x62, 0xb8, 0x1b, 0x82, 0xb5, 0x0c, 0x27, 0x64, 0x6e, 0xd5, 0x76, 0x2f,
            ],
        ),
        (
            b"haha",
            [
                0xeb, 0x07, 0x8b, 0x03, 0xd4, 0x6f, 0x6d, 0x91, 0xe2, 0x34, 0x80, 0xdc, 0xc4, 0xb7, 0x53, 0xc9,
                0x99, 0x87, 0x8f, 0xf4, 0x69, 0x06, 0xe7, 0x9b, 0x3e, 0x4b, 0x6b, 0x33, 0x6a, 0xc1, 0x36, 0xe1,
            ],
        )
    ] {
        assert_eq!(Hash::Shake256.hash(data), result);
    }
}