// Copyright (c) 2018 Henrique Sasaki Yuya (moriturus)
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#![feature(test, iterator_flatten)]

//! A Rust implementation of the [Lamport one-time signature scheme](https://en.wikipedia.org/wiki/Lamport_signature).
//!
//! # Notice
//!
//! This crate can build only on the `nightly` channel.
//!
//! # Usage
//!
//! ```
//! extern crate lamport_signature;
//! extern crate sha2;
//! extern crate rand;
//!
//! use lamport_signature::{PublicKey, PrivateKey, generate_keys};
//! use sha2::Sha256;
//! use rand::thread_rng;
//!
//! let mut rng = thread_rng();
//! let (mut private_key, public_key) = generate_keys::<Sha256, _>(&mut rng);
//!
//! let signature = private_key.sign(b"Hello, World!").expect("signing failed");
//! assert!(public_key.verify(&signature, b"Hello, World!"));
//! ```
//!
//! # Digest Algorithm
//!
//! [PrivateKey] and [PublicKey] should take the fixed output size digest algorithm types that provided by [RustCrypto/hashes](https://github.com/RustCrypto/hashes) as a type argument to construct.
//!
//! # RNG Algorithm
//!
//! [PrivateKey] and [PublicKey] also should take the specified RNG implemented in [rust-lang-nursery/rand](https://github.com/rust-lang-nursery/rand) as an argument to construct.

extern crate digest;
extern crate rand;
extern crate typenum;
#[macro_use]
extern crate log;

#[cfg(test)]
extern crate sha3;
#[cfg(test)]
extern crate test;

use digest::Digest;
use rand::RngCore;
use std::cmp::{Eq, Ord, Ordering, PartialEq, PartialOrd};
use std::hash::{Hash, Hasher};
use std::io::{self, Read};
use std::marker::PhantomData;
use typenum::Unsigned;

/// Errors in sign-verify scheme.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Error {
    /// An error indicating that the private key was reused.
    PrivateKeyReuseError,
}

/// A signature data generated by [PrivateKey].
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Signature {
    data: Vec<Vec<u8>>,
}

impl Signature {
    /// Constructs a [Signature] from the object implements [std::io::Read], like [std::fs::File].
    ///
    /// # Example
    ///
    /// ```
    /// extern crate sha2;
    /// extern crate lamport_signature;
    ///
    /// use std::fs::File;
    /// use std::io::Read;
    /// use sha2::Sha256;
    /// use lamport_signature::Signature;
    ///
    /// let mut file = File::open("signed_data").unwrap();
    /// assert!(Signature::from_bytes::<Sha256, _>(&mut file).is_ok());
    /// ```
    pub fn from_bytes<T, U>(bytes: &mut U) -> Result<Signature, io::Error>
    where
        T: Digest,
        U: Read,
    {
        let digest_size_in_bytes = T::OutputSize::to_usize();
        let digest_size_in_bits = digest_size_in_bytes * 8;
        let required_size_in_bytes = digest_size_in_bytes * digest_size_in_bits;
        debug!(
            "required size: {} bytes ({} bits)",
            required_size_in_bytes,
            required_size_in_bytes * 8
        );
        let mut buffer = vec![0u8; required_size_in_bytes];

        bytes.read_exact(&mut buffer)?;

        let data = unflatten_random_numbers::<T>(&mut buffer);
        Ok(Signature::new(data))
    }

    /// Converts the inner signature data into a linearized vector.
    pub fn to_bytes(&self) -> Vec<u8> {
        self.data.clone().into_iter().flatten().collect()
    }

    fn new(data: Vec<Vec<u8>>) -> Signature {
        Signature { data: data }
    }

    fn as_vec(&self) -> &Vec<Vec<u8>> {
        &self.data
    }
}

/// A one-time signing public key.
///
/// In general, a public key is generated by the paired [PrivateKey] or [generate_keys] function.
#[derive(Debug, Clone)]
pub struct PublicKey<T> {
    zero_values: Vec<Vec<u8>>,
    one_values: Vec<Vec<u8>>,
    algorithm: PhantomData<T>,
}

impl<T> PartialEq for PublicKey<T> {
    fn eq(&self, other: &Self) -> bool {
        self.zero_values == other.zero_values
            && self.one_values == other.one_values
            && self.algorithm == other.algorithm
    }
}

impl<T> Eq for PublicKey<T> {}

impl<T> PartialOrd for PublicKey<T> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<T> Ord for PublicKey<T> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.zero_values
            .cmp(&other.zero_values)
            .then(self.one_values.cmp(&other.one_values))
            .then(self.algorithm.cmp(&other.algorithm))
    }
}

impl<T> Hash for PublicKey<T> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.zero_values.hash(state);
        self.one_values.hash(state);
        self.algorithm.hash(state);
    }
}

impl<T> PublicKey<T>
where
    T: Digest,
{
    /// Constructs a [PublicKey] from the object implements [std::io::Read], like [std::fs::File].
    ///
    /// # Example
    ///
    /// ```
    /// extern crate sha2;
    /// extern crate lamport_signature;
    ///
    /// use std::fs::File;
    /// use std::io::Read;
    /// use sha2::Sha256;
    /// use lamport_signature::PublicKey;
    ///
    /// let mut file = File::open("test_key.pub").unwrap();
    /// assert!(PublicKey::<Sha256>::from_bytes(&mut file).is_ok());
    /// ```
    pub fn from_bytes<U>(bytes: &mut U) -> Result<PublicKey<T>, io::Error>
    where
        U: Read,
    {
        let digest_size_in_bytes = T::OutputSize::to_usize();
        let digest_size_in_bits = digest_size_in_bytes * 8;
        let required_size_in_bytes = digest_size_in_bytes * digest_size_in_bits * 2;
        debug!(
            "required size: {} bytes ({} bits)",
            required_size_in_bytes,
            required_size_in_bytes * 8
        );
        let mut buffer = vec![0u8; required_size_in_bytes];

        bytes.read_exact(&mut buffer)?;

        let mut zero_values_merged = buffer;
        let one_values_merged = zero_values_merged.split_off(required_size_in_bytes / 2);

        let zero_values = unflatten_random_numbers::<T>(&zero_values_merged);
        let one_values = unflatten_random_numbers::<T>(&one_values_merged);

        Ok(PublicKey {
            zero_values: zero_values,
            one_values: one_values,
            algorithm: PhantomData,
        })
    }

    /// Converts the inner key data into a linearized vector.
    ///
    /// # Example
    ///
    /// ```
    /// extern crate sha2;
    /// extern crate rand;
    /// extern crate lamport_signature;
    ///
    /// use sha2::Sha256;
    /// use rand::{XorShiftRng, SeedableRng};
    /// use lamport_signature::{PublicKey, generate_keys};
    ///
    /// const SEED: [u8; 16] = [0; 16];
    /// let mut rng = XorShiftRng::from_seed(SEED);
    /// let (_, public_key) = generate_keys::<Sha256, _>(&mut rng);
    /// assert!(public_key.to_bytes().len() == 32 * 256 * 2);
    /// ```
    pub fn to_bytes(&self) -> Vec<u8> {
        self.zero_values
            .clone()
            .into_iter()
            .chain(self.one_values.clone().into_iter())
            .flatten()
            .collect()
    }

    /// Verifies the [Signature].
    ///
    /// # Example
    ///
    /// ```
    /// extern crate sha2;
    /// extern crate rand;
    /// extern crate lamport_signature;
    ///
    /// use sha2::Sha256;
    /// use rand::{XorShiftRng, SeedableRng};
    /// use lamport_signature::{PrivateKey, PublicKey, generate_keys};
    ///
    /// const SEED: [u8; 16] = [0; 16];
    /// let mut rng = XorShiftRng::from_seed(SEED);
    /// let (mut private_key, public_key) = generate_keys::<Sha256, _>(&mut rng);
    ///
    /// const MESSAGE: &[u8] = b"hello, world!";
    /// let signature = private_key.sign(MESSAGE).expect("failed to sign.");
    /// assert!(public_key.verify(&signature, MESSAGE));
    /// ```
    pub fn verify(&self, signature: &Signature, data: &[u8]) -> bool {
        if signature.as_vec().len() != T::OutputSize::to_usize() * 8 {
            return false;
        }

        let data_digest = T::digest(data);

        data_digest.into_iter().enumerate().all(|(i, byte)| {
            (0..8).all(|j| {
                let offset = i * 8 + j;
                if (byte & (1 << j)) > 0 {
                    let hashed_value = Vec::from(T::digest(&signature.as_vec()[offset]).as_ref());
                    hashed_value == self.one_values[offset]
                } else {
                    let hashed_value = Vec::from(T::digest(&signature.as_vec()[offset]).as_ref());
                    hashed_value == self.zero_values[offset]
                }
            })
        })
    }
}

/// A one-time signing private key.
#[derive(Debug, Clone)]
pub struct PrivateKey<T> {
    zero_values: Vec<Vec<u8>>,
    one_values: Vec<Vec<u8>>,
    used: bool,
    algorithm: PhantomData<T>,
}

impl<T> Drop for PrivateKey<T> {
    fn drop(&mut self) {
        zeroing(&mut self.zero_values);
        zeroing(&mut self.one_values);
    }
}

impl<T> PartialEq for PrivateKey<T> {
    fn eq(&self, other: &Self) -> bool {
        self.zero_values == other.zero_values
            && self.one_values == other.one_values
            && self.used == other.used
            && self.algorithm == other.algorithm
    }
}

impl<T> Eq for PrivateKey<T> {}

impl<T> PartialOrd for PrivateKey<T> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<T> Ord for PrivateKey<T> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.zero_values
            .cmp(&other.zero_values)
            .then(self.one_values.cmp(&other.one_values))
            .then(self.used.cmp(&other.used))
            .then(self.algorithm.cmp(&other.algorithm))
    }
}

impl<T> Hash for PrivateKey<T> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.zero_values.hash(state);
        self.one_values.hash(state);
        self.used.hash(state);
        self.algorithm.hash(state);
    }
}

impl<T> PrivateKey<T> {
    /// Whether or not the [PrivateKey] was used.
    pub fn used(&self) -> bool {
        self.used
    }
}

impl<T> PrivateKey<T>
where
    T: Digest,
{
    /// Constructs a [PrivateKey] with Digest algorithm type and the specified RNG.
    pub fn new<R>(rng: &mut R) -> PrivateKey<T>
    where
        R: RngCore,
    {
        PrivateKey {
            zero_values: generate_random_numbers::<T, _>(rng),
            one_values: generate_random_numbers::<T, _>(rng),
            used: false,
            algorithm: PhantomData,
        }
    }

    /// Generates a [PublicKey] from the [PrivateKey].
    pub fn public_key(&self) -> PublicKey<T> {
        let hashed_zero_values = hashed_random_numbers::<T>(&self.zero_values);
        let hashed_one_values = hashed_random_numbers::<T>(&self.one_values);

        PublicKey {
            zero_values: hashed_zero_values,
            one_values: hashed_one_values,
            algorithm: PhantomData,
        }
    }

    /// Signs the data.
    ///
    /// # Example
    ///
    /// ```
    /// extern crate sha2;
    /// extern crate rand;
    /// extern crate lamport_signature;
    ///
    /// use sha2::Sha256;
    /// use rand::{XorShiftRng, SeedableRng};
    /// use lamport_signature::PrivateKey;
    ///
    /// const SEED: [u8; 16] = [0; 16];
    /// let mut rng = XorShiftRng::from_seed(SEED);
    /// let mut private_key = PrivateKey::<Sha256>::new(&mut rng);
    /// const MESSAGE: &[u8] = b"hello, world!";
    /// assert!(private_key.sign(MESSAGE).is_ok());
    /// ```
    pub fn sign(&mut self, data: &[u8]) -> Result<Signature, Error> {
        if self.used {
            return Err(Error::PrivateKeyReuseError);
        }

        let data_hash = T::digest(data);

        let signature_len = data_hash.len() * 8;
        let signature_data = data_hash.iter().enumerate().fold(
            Vec::with_capacity(signature_len),
            |mut accm, (i, byte)| {
                let dispached = (0..8).map(|j| {
                    let offset = i * 8 + j;
                    if (byte & (1 << j)) > 0 {
                        self.one_values[offset].clone() // the bit is 1
                    } else {
                        self.zero_values[offset].clone() // the bit is 0
                    }
                });

                accm.append(&mut dispached.collect());
                accm
            },
        );

        self.used = true;
        Ok(Signature::new(signature_data))
    }
}

/// Generates a [PrivateKey] and [PublicKey] at once.
pub fn generate_keys<T, R>(rng: &mut R) -> (PrivateKey<T>, PublicKey<T>)
where
    T: Digest,
    R: RngCore,
{
    let private_key = PrivateKey::<T>::new(rng);
    let public_key = private_key.public_key();
    (private_key, public_key)
}

fn generate_random_numbers<T, R>(rng: &mut R) -> Vec<Vec<u8>>
where
    T: Digest,
    R: RngCore,
{
    let size_in_bytes = T::OutputSize::to_usize();
    let outer_size_in_bits = size_in_bytes * 8;
    let outer = Vec::with_capacity(outer_size_in_bits);

    (0..outer_size_in_bits).fold(outer, |mut accm, _| {
        let mut inner = vec![0u8; size_in_bytes];
        rng.fill_bytes(&mut inner);
        accm.push(inner);
        accm
    })
}

fn hashed_random_numbers<T>(original: &[Vec<u8>]) -> Vec<Vec<u8>>
where
    T: Digest,
{
    let outer_size_in_bits = T::OutputSize::to_usize() * 8;
    let outer = Vec::with_capacity(outer_size_in_bits);

    (0..outer_size_in_bits).fold(outer, |mut accm, i| {
        let inner = Vec::from(T::digest(&original[i]).as_ref());
        accm.push(inner);
        accm
    })
}

fn unflatten_random_numbers<T>(origin: &[u8]) -> Vec<Vec<u8>>
where
    T: Digest,
{
    let len = origin.len();
    let output_size_in_bytes = T::OutputSize::to_usize();

    (0..len)
        .filter_map(|i| {
            if i % output_size_in_bytes == 0 {
                let sub_vec = (0..output_size_in_bytes).fold(Vec::new(), |mut accm, j| {
                    accm.push(origin[i + j]);
                    accm
                });
                Some(sub_vec)
            } else {
                None
            }
        })
        .collect()
}

fn zeroing(vec: &mut Vec<Vec<u8>>) {
    vec.iter_mut()
        .for_each(|inner| inner.iter_mut().for_each(|byte| *byte = 0));
}

#[cfg(test)]
mod tests {
    extern crate env_logger;
    use super::{generate_keys, PrivateKey, PublicKey, Signature};
    use rand::OsRng;
    use sha3::{Sha3_256, Sha3_512};
    use std::fs::File;
    use std::io::Read;

    #[test]
    fn test_public_key_conversion() {
        let _ = env_logger::try_init();

        let mut rng = OsRng::new().unwrap();
        let (_, original_public_key) = generate_keys::<Sha3_256, _>(&mut rng);

        let bytes = original_public_key.to_bytes();
        let mut bytes = &bytes[..];
        let restored_public_key = match PublicKey::<Sha3_256>::from_bytes(&mut bytes) {
            Ok(k) => k,
            Err(reason) => panic!("The public key cannot be restored: {}", reason),
        };

        assert!(restored_public_key == original_public_key);
    }

    #[test]
    fn test_seedable_rng_generate_sha3_256_private_key() {
        use rand::{ChaChaRng, SeedableRng, XorShiftRng};

        const CHACHA_SEED: [u8; 32] = [
            1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24,
            25, 26, 27, 28, 29, 30, 31, 32,
        ];
        const XOR_SHIFT_SEED: [u8; 16] = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];

        let _ = env_logger::try_init();

        let mut chacha_rng = ChaChaRng::from_seed(CHACHA_SEED);
        let chacha_private_key = PrivateKey::<Sha3_256>::new(&mut chacha_rng);

        assert!(
            chacha_private_key.zero_values[0]
                == &[
                    177, 105, 126, 159, 198, 70, 30, 25, 131, 209, 49, 207, 105, 105, 28, 161, 167,
                    163, 252, 19, 79, 20, 152, 128, 232, 187, 43, 93, 35, 101, 225, 3
                ]
        );

        let mut xor_shift_rng = XorShiftRng::from_seed(XOR_SHIFT_SEED);
        let xor_shift_private_key = PrivateKey::<Sha3_256>::new(&mut xor_shift_rng);

        assert!(
            xor_shift_private_key.zero_values[0]
                == &[
                    7, 21, 0, 12, 172, 13, 7, 60, 103, 19, 8, 104, 13, 7, 31, 0, 36, 130, 187, 12,
                    114, 132, 213, 8, 36, 61, 110, 32, 75, 113, 177, 216
                ]
        );
    }

    #[test]
    fn test_generate_sha3_256_private_key() {
        let _ = env_logger::try_init();

        let mut rng = OsRng::new().unwrap();
        let private_key = PrivateKey::<Sha3_256>::new(&mut rng);

        assert!(!private_key.used());
        assert!(private_key.zero_values.len() == 256);
        assert!(private_key.zero_values.iter().all(|i| i.len() == 32));
        assert!(private_key.one_values.len() == 256);
        assert!(private_key.one_values.iter().all(|i| i.len() == 32));
    }

    #[test]
    fn test_generate_sha3_512_private_key() {
        let _ = env_logger::try_init();

        let mut rng = OsRng::new().unwrap();
        let private_key = PrivateKey::<Sha3_512>::new(&mut rng);

        assert!(!private_key.used());
        assert!(private_key.zero_values.len() == 512);
        assert!(private_key.zero_values.iter().all(|i| i.len() == 64));
        assert!(private_key.one_values.len() == 512);
        assert!(private_key.one_values.iter().all(|i| i.len() == 64));
    }

    #[test]
    fn test_sha3_256_xor_shift_rng_public_key_from_bytes_then_verify() {
        let _ = env_logger::try_init();

        let mut key_file = File::open("test_key.pub").expect("could not open test_key.pub file.");
        let public_key =
            PublicKey::<Sha3_256>::from_bytes(&mut key_file).expect("could not read from the file");

        let mut signed_data_file =
            File::open("signed_data").expect("could not open signed_data file.");
        let signature = Signature::from_bytes::<Sha3_256, _>(&mut signed_data_file)
            .expect("could not read from the file.");

        assert!(public_key.verify(&signature, b"hello, world!"));
    }

    #[test]
    fn test_sha3_256_xor_shift_rng_public_key_to_bytes() {
        use rand::{SeedableRng, XorShiftRng};

        const SEED: [u8; 16] = [0; 16];

        let _ = env_logger::try_init();

        let mut rng = XorShiftRng::from_seed(SEED);
        let (_, public_key) = generate_keys::<Sha3_256, _>(&mut rng);
        let lineared = public_key.to_bytes();

        let mut key_file = File::open("test_key.pub").expect("could not open test_key.pub file.");
        let mut test_public_key = [0u8; (256 / 8) * 256 * 2];
        key_file
            .read_exact(&mut test_public_key)
            .expect("could not read test_key.pub file.");

        assert!(lineared[0] == test_public_key[0]);
        assert!(lineared[((256 / 8) * 256 * 2) - 1] == test_public_key[((256 / 8) * 256 * 2) - 1]);
        assert!(lineared.len() == (256 / 8) * 256 * 2);
    }
}