//! Secret sharing with authentication
//!
//! Internally this uses [sharks](https://docs.rs/sharks/0.5.0/sharks/) for Shamirs secret sharing.
//!
//! This is part of a work-in-progress Rust implementation of
//! the [Dark Crystal Key Backup Protocol](https://darkcrystal.pw/protocol-specification/).

use rand::{thread_rng, Rng};
use sharks::{Share, Sharks};
use std::collections::HashSet;
use std::convert::{TryFrom, TryInto};
use std::fmt;
use xsalsa20poly1305::aead::generic_array::GenericArray;
use zeroize::Zeroize;

pub mod encrypt;

/// Generate random indexes to take a subset of 255 shares
fn get_shares_to_take(num_shares: u8) -> HashSet<u8> {
    let mut rng = thread_rng();
    let mut shares_to_take = HashSet::new();
    while shares_to_take.len() < num_shares.into() {
        let next_index = rng.gen::<u8>();
        if next_index > 0 {
            shares_to_take.insert(next_index);
        }
    }
    shares_to_take
}

/// Create a set of shares for a given secret
pub fn share(secret: &[u8], num_shares: u8, threshold: u8) -> Result<Vec<Vec<u8>>, ShareError> {
    if num_shares < 2 {
        return Err(ShareError {
            message: "At least 2 shares required".to_string(),
        });
    }

    if threshold < 2 {
        return Err(ShareError {
            message: "Threshold must be at least 2".to_string(),
        });
    }

    if threshold > num_shares {
        return Err(ShareError {
            message: "Threshold must not be greater than the number of shares".to_string(),
        });
    }

    let sharks = Sharks(threshold);
    let dealer = sharks.dealer(&secret[..]);
    let shares_to_take = get_shares_to_take(num_shares);

    Ok(dealer
        .take(255)
        .map(|s| Vec::from(&s))
        .filter(|s| shares_to_take.contains(&s[0]))
        .collect())
}

/// Recover a secret from a given set of shares
pub fn combine(shares_bytes: Vec<Vec<u8>>) -> Result<Vec<u8>, RecoveryError> {
    let shares: Vec<Share> = shares_bytes
        .iter()
        .map(|s| Share::try_from(s.as_slice()).unwrap())
        .collect();

    let sharks = Sharks(shares.len().try_into().unwrap());
    match sharks.recover(&shares) {
        Ok(val) => Ok(val),
        Err(err) => Err(err.into()),
    }
}

/// Encrypt a secret and create shares of its key.
/// This gives authentication so we know whether recovery was successful
/// It also reduces duplication with long (> 32 bytes) secrets,
/// and improves security when using non-uniformly random secrets such
/// as passwords.
pub fn share_authenticated(
    secret: &[u8],
    num_shares: u8,
    threshold: u8,
) -> Result<(Vec<Vec<u8>>, Vec<u8>), ShareError> {
    let mut key = encrypt::generate_key();
    match share(&key, num_shares, threshold) {
        Ok(shares) => {
            let ciphertext = encrypt::encrypt(key, secret.to_vec()).unwrap();
            key.zeroize();
            Ok((shares, ciphertext))
        }
        Err(err) => Err(err),
    }
}

/// Combine a set of shares and ciphertext produced by share_authenticated
pub fn combine_authenticated(
    shares: Vec<Vec<u8>>,
    ciphertext: Vec<u8>,
) -> Result<Vec<u8>, RecoveryError> {
    let recovered_key = combine(shares)?;
    let key = GenericArray::from_slice(&recovered_key[..]);
    match encrypt::decrypt(*key, ciphertext) {
        Ok(val) => Ok(val),
        Err(err) => Err(err.into()),
    }
}

/// Error created when share fn fails
#[derive(Debug)]
pub struct ShareError {
    pub message: String,
}

impl fmt::Display for ShareError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Cannot create shares: {}", self.message)
    }
}

/// Error created when recovery fails
#[derive(Debug)]
pub struct RecoveryError {
    pub message: String,
}

impl fmt::Display for RecoveryError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Error during recovery {}", self.message)
    }
}

impl From<xsalsa20poly1305::aead::Error> for RecoveryError {
    fn from(error: xsalsa20poly1305::aead::Error) -> Self {
        RecoveryError {
            message: error.to_string(),
        }
    }
}

impl From<&str> for RecoveryError {
    fn from(error: &str) -> Self {
        RecoveryError {
            message: String::from(error),
        }
    }
}

/// Give a recommended threshold value for a given number of shares
pub fn default_threshold(number_of_shares: u8) -> u8 {
    if number_of_shares == 2 {
        return 2;
    };
    (number_of_shares as f32 * 0.75) as u8
}

/// Gives a threshold 'sanity' factor, given a threshold and number of shares
/// 0 is ideal.  Positive values represent the level of danger of
/// loosing access to the secret.
/// Negative values represent the level of danger of an attacker gaining it.
pub fn thresold_sanity(threshold: u8, number_of_shares: u8) -> i32 {
    (threshold as i32) - (default_threshold(number_of_shares) as i32)
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn passes_on_all_shares() {
        let original_secret = b"hello";
        let shares = share(&original_secret[..], 5, 3).unwrap();
        assert_eq!(shares.len(), 5);
        let recovered_secret = combine(shares).unwrap();
        assert_eq!(recovered_secret, b"hello");
    }

    #[test]
    fn passes_on_threshold_amount() {
        let original_secret = b"hello";
        let mut shares = share(&original_secret[..], 5, 3).unwrap();
        assert_eq!(shares.len(), 5);
        shares.remove(0);
        shares.remove(1);
        assert_eq!(shares.len(), 3);
        let recovered_secret = combine(shares).unwrap();
        assert_eq!(recovered_secret, b"hello");
    }

    #[test]
    fn authenticated() {
        let original_secret = b"hello";
        let (shares, ciphertext) = share_authenticated(&original_secret[..], 5, 3).unwrap();
        assert_eq!(shares.len(), 5);
        let recovered_secret = combine_authenticated(shares, ciphertext).unwrap();
        assert_eq!(recovered_secret, b"hello");
    }

    #[test]
    fn fails_on_insufficient_shares() {
        let original_secret = b"hello";
        let (mut shares, ciphertext) = share_authenticated(&original_secret[..], 5, 3).unwrap();
        assert_eq!(shares.len(), 5);
        shares.remove(0);
        shares.remove(1);
        shares.remove(2);
        assert_eq!(shares.len(), 2);

        let res = combine_authenticated(shares, ciphertext);
        assert!(res.is_err());
    }

    #[test]
    fn fails_on_impossible_threshold() {
        let original_secret = b"hello";
        let share_result = share(&original_secret[..], 5, 6);
        assert!(share_result.is_err());
    }

    #[test]
    fn fails_on_threshold_of_one() {
        let original_secret = b"hello";
        let share_result = share(&original_secret[..], 5, 1);
        assert!(share_result.is_err());
    }

    #[test]
    fn fails_on_creating_one_share() {
        let original_secret = b"hello";
        let share_result = share(&original_secret[..], 1, 1);
        assert!(share_result.is_err());
    }

    #[test]
    fn fails_on_bad_share() {
        let original_secret = b"hello";
        let (mut shares, ciphertext) = share_authenticated(&original_secret[..], 5, 3).unwrap();
        shares[0] = b"bad share".to_vec();

        let res = combine_authenticated(shares, ciphertext);
        assert!(res.is_err());
    }

    #[test]
    fn check_default_threshold() {
        assert_eq!(thresold_sanity(7, 10), 0);
        assert_eq!(thresold_sanity(5, 7), 0);
        assert_eq!(thresold_sanity(3, 5), 0);
        assert_eq!(thresold_sanity(2, 3), 0);
        assert_eq!(thresold_sanity(2, 2), 0);

        assert_eq!(thresold_sanity(2, 10), -5);
        assert_eq!(thresold_sanity(9, 10), 2);
    }
}