#![no_std]
//! LIONESS implementation in Rust.
//!
//! # ☣️ Cryptographic hazmat ☣️
//!
//! This crate is not battle tested, nor is it audited. Its usage for critical systems is strongly
//! discouraged. It mainly exists as a learning exercise.
//!
//! # LIONESS
//!
//! [LIONESS](https://link.springer.com/content/pdf/10.1007/3-540-60865-6_48.pdf) is a block cipher
//! based on the Luby-Rackoff construction. It combines a hash function and a stream cipher to
//! produce a block permutation that allows for blocks of arbitrary(*) length to be encrypted.
//!
//! Internally, the cipher uses two instantiations of the keyed hash function, and two
//! instantiations of the stream cipher, leading to four separate round keys. The stream cipher and
//! the hash must be compatible: The output size of the hash must be the same as the key size of
//! the cipher.
//!
//! (*): The construction does not specify how short inputs (smaller than the hash output) should
//! be encrypted. As such, [Lioness] will error on such inputs, and you have to pad the input
//! manually.
//!
//! # Implementation
//!
//! The cipher is implemented as the [Lioness] struct, and is kept generic over the stream cipher
//! (using the [`cipher`](https://crates.io/crates/cipher) crate), as well as the hash function
//! (using the [`digest`](https://crates.io/crates/digest) crate). Rust's type system is used to
//! ensure that the cipher is compatible with the hash (the hash's output size must match the
//! cipher's key size).
//!
//! The construction is implemented in-place and works without allocations.
//!
//! This crate implements several convenience methods:
//!
//! To prevent users from having to manually create four round keys, we provide
//! [Lioness::new_dynamic]. This method internally uses `blake2` to create distinct round keys from
//! a given (potentially short) input key.
//!
//! To use ciphers that require initialization vectors (IVs), we provide [ZeroIv]. This struct
//! wraps a cipher that requires an IV, and supplies a zero-IV.
//!
//! To use unkeyed hash functions (like SHA), we provide [KeyedHash]. This struct wraps a hash
//! function and turns it into a keyed hash by prepending the key to the input data.
//!
//! # Examples
//!
//! A simple working combination for [Lioness] is to use
//! [ChaCha20](https://docs.rs/chacha20/latest/chacha20/type.ChaCha20.html) and
//! [Blake2s256](https://docs.rs/blake2/latest/blake2/type.Blake2sMac256.html) (32 byte key/hash
//! output). Note that we use the MAC variant of Blake2 to get a keyed version of the hash
//! directly:
//!
//! ```
//! use leona::{Lioness, ZeroIv};
//! use blake2::Blake2sMac256;
//! use chacha20::ChaCha20;
//!
//! type Cipher = Lioness<ZeroIv<ChaCha20>, Blake2sMac256>;
//!
//! let mut data = [0u8; 64];
//! data[..11].copy_from_slice(b"hello world");
//!
//! let cipher = Cipher::new_dynamic(b"secret");
//! cipher.encrypt(&mut data);
//! assert_ne!(&data[..11], b"hello world");
//! cipher.decrypt(&mut data);
//! assert_eq!(&data[..11], b"hello world");
//! ```
//!
//! You can also use AES in counter mode as a stream cipher, and SHA256 as a hash. Note how we have
//! to use the `Ctr32BE` wrapper to turn AES into a stream cipher, and [KeyedHash] to turn SHA into
//! a keyed hash. We choose Aes256, since its key size matches SHA256's output size (32 bytes):
//!
//! ```
//! use leona::{KeyedHash, Lioness, ZeroIv};
//! use aes::Aes256;
//! use ctr::Ctr32BE;
//! use sha2::Sha256;
//! use typenum::U32;
//!
//! type Cipher = Lioness<ZeroIv<Ctr32BE<Aes256>>, KeyedHash<U32, Sha256>>;
//!
//! let mut data = [0u8; 64];
//! data[..11].copy_from_slice(b"hello world");
//!
//! let cipher = Cipher::new_dynamic(b"secret");
//! cipher.encrypt(&mut data);
//! assert_ne!(&data[..11], b"hello world");
//! cipher.decrypt(&mut data);
//! assert_eq!(&data[..11], b"hello world");
//! ```
//!
//! Alternatively, you can use [Hmac](https://docs.rs/hmac/latest/hmac/index.html) to turn unkeyed
//! hash functions into keyed hashes:
//!
//! ```
//! use leona::{KeyedHash, Lioness, ZeroIv};
//! use aes::Aes256;
//! use ctr::Ctr32BE;
//! use hmac::Hmac;
//! use sha2::Sha256;
//!
//! type Cipher = Lioness<ZeroIv<Ctr32BE<Aes256>>, Hmac<Sha256>>;
//!
//! let mut data = [0u8; 64];
//! data[..11].copy_from_slice(b"hello world");
//!
//! let cipher = Cipher::new_dynamic(b"secret");
//! cipher.encrypt(&mut data);
//! assert_ne!(&data[..11], b"hello world");
//! cipher.decrypt(&mut data);
//! assert_eq!(&data[..11], b"hello world");
//! ```
//!
//! If the input data is shorter than the hash output, encryption will fail:
//!
//! ```
//! # use leona::{Lioness, ZeroIv};
//! # use blake2::Blake2sMac256;
//! # use chacha20::ChaCha20;
//! # type Cipher = Lioness<ZeroIv<ChaCha20>, Blake2sMac256>;
//! let mut data = [0u8; 16];
//! data[..11].copy_from_slice(b"hello world");
//!
//! let cipher = Cipher::new_dynamic(b"secret");
//! assert!(cipher.encrypt(&mut data).is_err());
//! ```
//!
//! If you need more control over the keys, use the [Lioness::new] constructor. In this case, you
//! can provide all round keys separately.
//!
//! # Alternatives
//!
//! LIONESS is slow, and it cannot encrypt short data. If you need a wide-block cipher, consider a
//! different construction such as `aez` (via [aez](https://crates.io/crates/aez) or
//! [zears](https://crates.io/crates/zears)):
//!
//! |              | `aes256`     | `chacha20`   |
//! |--------------|--------------|--------------|
//! | **`sha256`** | 645.58 MiB/s | 495.01 MiB/s |
//! | **`blake2`** | 285.56 MiB/s | 251.19 MiB/s |
//!
//! For comparison, `zears` achieves 5.8170 GiB/s (+simd, target-cpu=native).
use core::marker::PhantomData;

use blake2::Blake2b;
use cipher::{
    InOutBuf,
    stream::{StreamCipher, StreamCipherError},
};
use crypto_common::{KeyInit, KeyIvInit, KeySizeUser, OutputSizeUser};
use digest::{CustomizedInit, FixedOutput, Update};
use hybrid_array::{Array, ArraySize};
use thiserror::Error;
use typenum::{IsLessOrEqual, True, U64, Unsigned};
use zerocopy::FromZeros;

const SALT_1: &[u8] = b"k1";
const SALT_2: &[u8] = b"k2";
const SALT_3: &[u8] = b"k3";
const SALT_4: &[u8] = b"k4";

/// Errors that can occur during encryption and decryption.
#[derive(Error, Debug, Clone)]
pub enum Error {
    /// The given input data is shorter than the hash output.
    #[error("given input data is too short")]
    InputTooShort,
}

/// Main struct implementing the LIONESS construction.
///
/// `S` represents the stream cipher and `H` represents the hash algorithm.
///
/// See the crate level documentation for more details.
pub struct Lioness<S, H>
where
    S: KeySizeUser,
    H: KeySizeUser,
{
    k_1: Array<u8, S::KeySize>,
    k_2: Array<u8, H::KeySize>,
    k_3: Array<u8, S::KeySize>,
    k_4: Array<u8, H::KeySize>,
}

fn hash_key<L>(salt: &[u8], input: &[u8]) -> Array<u8, L>
where
    L: ArraySize + IsLessOrEqual<U64, Output = True>,
{
    let mut hasher = Blake2b::new_customized(salt);
    hasher.update(input);
    hasher.finalize_fixed()
}

fn xor_assign<L: ArraySize>(a: &mut L::ArrayType<u8>, b: &L::ArrayType<u8>) {
    a.as_mut()
        .iter_mut()
        .zip(b.as_ref().iter())
        .for_each(|(x, y)| *x ^= *y);
}

type Key<K> = <<K as KeySizeUser>::KeySize as ArraySize>::ArrayType<u8>;

impl<S, H> Lioness<S, H>
where
    S: KeySizeUser<KeySize = H::OutputSize> + KeyInit + StreamCipher,
    H: KeySizeUser + KeyInit + FixedOutput + Update,
    S::KeySize: IsLessOrEqual<U64, Output = True>,
    H::KeySize: IsLessOrEqual<U64, Output = True>,
    Array<u8, S::KeySize>: FromZeros,
{
    /// Create a new [Lioness] instance from the given key.
    ///
    /// Unlike [Lioness::new], you are not required to pass the correct key sizes manually.
    /// Instead, `blake2` is used internally to derive keys of the correct length.
    ///
    /// Note that this only works for keys up to 64 bytes. If you have longer keys, you need to use
    /// [Lioness::new].
    pub fn new_dynamic(key: &[u8]) -> Self {
        Self::new(
            hash_key::<S::KeySize>(SALT_1, key).0,
            hash_key::<H::KeySize>(SALT_2, key).0,
            hash_key::<S::KeySize>(SALT_3, key).0,
            hash_key::<H::KeySize>(SALT_4, key).0,
        )
    }
}

impl<S, H> Lioness<S, H>
where
    S: KeySizeUser<KeySize = H::OutputSize> + KeyInit + StreamCipher,
    H: KeySizeUser + KeyInit + FixedOutput + Update,
    Array<u8, S::KeySize>: FromZeros,
{
    /// Create a new [Lioness] instance from the given round keys.
    ///
    /// `k_1` and `k_3` are used for the stream cipher and must match its key length, while `k_2`
    /// and `k_4` are used for the hash function.
    ///
    /// Don't be scared of the long type signatures. They are just normal arrays of `u8`:
    ///
    /// ```
    /// # use leona::{Lioness, ZeroIv};
    /// # use blake2::Blake2sMac256;
    /// # use chacha20::ChaCha20;
    /// type Cipher = Lioness<ZeroIv<ChaCha20>, Blake2sMac256>;
    /// Cipher::new([1; 32], [2; 32], [3; 32], [4; 32]);
    /// ```
    pub fn new(k_1: Key<S>, k_2: Key<H>, k_3: Key<S>, k_4: Key<H>) -> Self {
        Self {
            k_1: k_1.into(),
            k_2: k_2.into(),
            k_3: k_3.into(),
            k_4: k_4.into(),
        }
    }

    /// Encrypt the given buffer in-place.
    ///
    /// If the buffer is too short (shorter than the hash output), an error is returned.
    pub fn encrypt(&self, data: &mut [u8]) -> Result<(), Error> {
        let l_len = S::KeySize::USIZE;
        if data.len() < l_len {
            return Err(Error::InputTooShort);
        }

        assert!(data.len() >= l_len);

        // R = R + S(L + K_1)
        let mut key = <Array<u8, S::KeySize> as FromZeros>::new_zeroed();
        key.copy_from_slice(&data[..l_len]);
        xor_assign::<S::KeySize>(&mut key.0, &self.k_1.0);
        Self::apply_keystream(key.into(), &mut data[l_len..]);

        // L = L + H(K_2, R)
        let hash = Self::hash(&self.k_2, &data[l_len..]);
        data[..l_len]
            .iter_mut()
            .zip(hash)
            .for_each(|(x, y)| *x ^= y);

        // R = R + S(L + K_3)
        let mut key = <Array<u8, S::KeySize> as FromZeros>::new_zeroed();
        key.copy_from_slice(&data[..l_len]);
        xor_assign::<S::KeySize>(&mut key.0, &self.k_3.0);
        Self::apply_keystream(key.into(), &mut data[l_len..]);

        // L = L + H(K_4, R)
        let hash = Self::hash(&self.k_4, &data[l_len..]);
        data[..l_len]
            .iter_mut()
            .zip(hash)
            .for_each(|(x, y)| *x ^= y);

        Ok(())
    }

    /// Decrypt the given buffer in-place.
    ///
    /// If the buffer is too short (shorter than the hash output), an error is returned.
    pub fn decrypt(&self, data: &mut [u8]) -> Result<(), Error> {
        let l_len = S::KeySize::USIZE;
        if data.len() < l_len {
            return Err(Error::InputTooShort);
        }

        assert!(data.len() >= l_len);

        // L = L + H(K_4, R)
        let hash = Self::hash(&self.k_4, &data[l_len..]);
        data[..l_len]
            .iter_mut()
            .zip(hash)
            .for_each(|(x, y)| *x ^= y);

        // R = R + S(L + K_3)
        let mut key = <Array<u8, S::KeySize> as FromZeros>::new_zeroed();
        key.copy_from_slice(&data[..l_len]);
        xor_assign::<S::KeySize>(&mut key.0, &self.k_3.0);
        Self::apply_keystream(key.into(), &mut data[l_len..]);

        // L = L + H(K_2, R)
        let hash = Self::hash(&self.k_2, &data[l_len..]);
        data[..l_len]
            .iter_mut()
            .zip(hash)
            .for_each(|(x, y)| *x ^= y);

        // R = R + S(L + K_1)
        let mut key = <Array<u8, S::KeySize> as FromZeros>::new_zeroed();
        key.copy_from_slice(&data[..l_len]);
        xor_assign::<S::KeySize>(&mut key.0, &self.k_1.0);
        Self::apply_keystream(key.into(), &mut data[l_len..]);

        Ok(())
    }

    fn apply_keystream(key: Key<S>, data: &mut [u8]) {
        let mut cipher = S::new(&Array(key));
        cipher.apply_keystream(data);
    }

    fn hash(key: &Array<u8, H::KeySize>, data: &[u8]) -> Array<u8, H::OutputSize> {
        let mut hasher = H::new(key);
        hasher.update(data);
        hasher.finalize_fixed()
    }
}

/// A helper struct that turns ciphers with IV into ciphers without IV.
///
/// This turns ciphers that require an IV (like `ChaCha20`), represented by
/// [KeyIvInit](https://docs.rs/cipher/latest/cipher/trait.KeyIvInit.html), into ciphers that only
/// require a key, represented by
/// [KeyInit](https://docs.rs/cipher/latest/cipher/trait.KeyInit.html). To do so, an all-zero IV is
/// supplied.
#[derive(Debug, Clone)]
pub struct ZeroIv<C>(C);

impl<C: KeySizeUser> KeySizeUser for ZeroIv<C> {
    type KeySize = C::KeySize;
}

impl<C: KeyIvInit> KeyInit for ZeroIv<C>
where
    Array<u8, C::IvSize>: FromZeros,
{
    fn new(key: &Array<u8, Self::KeySize>) -> Self {
        let iv = <Array<u8, C::IvSize> as FromZeros>::new_zeroed();
        Self(C::new(key, &iv))
    }
}

impl<C: StreamCipher> StreamCipher for ZeroIv<C> {
    fn check_remaining(&self, data_len: usize) -> Result<(), StreamCipherError> {
        self.0.check_remaining(data_len)
    }

    fn unchecked_apply_keystream_inout(&mut self, buf: InOutBuf<'_, '_, u8>) {
        self.0.unchecked_apply_keystream_inout(buf)
    }

    fn unchecked_write_keystream(&mut self, buf: &mut [u8]) {
        self.0.unchecked_write_keystream(buf)
    }
}

/// A helper struct that turns unkeyed hashes into keyed hashes.
///
/// This turns hashes that have no key (like SHA256) into hashes that have a key. Note that the
/// construction is rather simple: the key is prepended to the input data. This method is suggested
/// in the LIONESS paper
///
/// > The keyed hash function H_K(M):
/// > (a) is based on an unkeyed hash function H'(M), in which we append and/or prepend the key to
/// > the message [...]
#[derive(Debug, Clone)]
pub struct KeyedHash<L, H>(H, PhantomData<L>);

impl<L: ArraySize, H> KeySizeUser for KeyedHash<L, H> {
    type KeySize = L;
}

impl<L: ArraySize, H: Default + Update> KeyInit for KeyedHash<L, H> {
    fn new(key: &Array<u8, Self::KeySize>) -> Self {
        let mut hasher = H::default();
        hasher.update(key);
        KeyedHash(hasher, PhantomData)
    }
}

impl<L: ArraySize, H: OutputSizeUser> OutputSizeUser for KeyedHash<L, H> {
    type OutputSize = H::OutputSize;
}

impl<L: ArraySize, H: Update> Update for KeyedHash<L, H> {
    fn update(&mut self, data: &[u8]) {
        self.0.update(data)
    }
}

impl<L: ArraySize, H: FixedOutput> FixedOutput for KeyedHash<L, H> {
    fn finalize_into(self, buffer: &mut Array<u8, Self::OutputSize>) {
        self.0.finalize_into(buffer)
    }
}