//! Provides encryption of secret shares to specific recipients using [`crypto_box`](https://docs.rs/crate/crypto_box)
//!
//! Internally uses [`dark-crystal-secret-sharing-rust`](https://docs.rs/dark-crystal-secret-sharing-rust),
//! which uses [`sharks`](https://docs.rs/sharks/0.5.0/sharks/) for Shamirs secret sharing and [`xsalsa20poly1305`](https://docs.rs/xsalsa20poly1305/0.8.0/xsalsa20poly1305/)
//! for authenticated encryption.
//!
//! This is part of a work-in-progress Rust implementation of the [Dark Crystal Key Backup Protocol](https://darkcrystal.pw/protocol-specification/).

pub use crypto_box::aead::Error;
use crypto_box::aead::{generic_array::GenericArray, Aead};
pub use crypto_box::{Box, PublicKey, SecretKey};
pub use dark_crystal_secret_sharing_rust::{
    combine_authenticated, default_threshold, share_authenticated, thresold_sanity, RecoveryError,
    ShareError,
};
use rand::Rng;
use std::convert::TryInto;
use std::fmt;
use zeroize::Zeroize;

const PUBLIC_KEY_LENGTH: usize = 32;
const NONCE_LENGTH: usize = 24;

/// A set of encrypted shares, together with the public key used for encryption
/// and the encrypted secret
#[derive(Debug)]
pub struct EncryptedShareSet {
    pub ciphertext: Vec<u8>,
    pub encrypted_shares: Vec<Vec<u8>>,
    pub eph_public_key: PublicKey,
}

/// Create a set of shares and encrypt them to a given set of public keys
pub fn share_and_encrypt(
    public_keys: Vec<[u8; PUBLIC_KEY_LENGTH]>,
    secret: Vec<u8>,
    threshold: u8,
) -> Result<EncryptedShareSet, ShareAndEncryptError> {
    // TODO make a ShareError so these errors can be handled
    let num_shares = public_keys.len().try_into().unwrap();
    let (shares, ciphertext) = share_authenticated(&secret, num_shares, threshold)?;
    let mut encrypted_shares: Vec<Vec<u8>> = Vec::new();

    let mut rng = crypto_box::rand_core::OsRng;
    let eph_secret_key = SecretKey::generate(&mut rng);
    let eph_public_key = eph_secret_key.public_key();
    let mut eph_secret_key_bytes = eph_secret_key.as_bytes().clone();

    for share_index in 0..public_keys.len() {
        let share = &shares[share_index];
        let pk = PublicKey::from(public_keys[share_index]);
        let esk = SecretKey::from(eph_secret_key_bytes);
        encrypted_shares.push(encrypt(esk, pk, share.to_vec())?);
    }
    eph_secret_key_bytes.zeroize();

    Ok(EncryptedShareSet {
        encrypted_shares,
        ciphertext,
        eph_public_key,
    })
}

/// Create a set of shares and encrypt them to a given set of public keys
/// but make the shares shorted by using the nonce from the ciphertext
/// when encrypting the shares
pub fn share_and_encrypt_detached_nonce(
    public_keys: Vec<[u8; PUBLIC_KEY_LENGTH]>,
    secret: Vec<u8>,
    threshold: u8,
) -> Result<EncryptedShareSet, ShareAndEncryptError> {
    // TODO make a ShareError so these errors can be handled
    let num_shares = public_keys.len().try_into().unwrap();
    let (shares, ciphertext) = share_authenticated(&secret, num_shares, threshold)?;
    let mut encrypted_shares: Vec<Vec<u8>> = Vec::new();

    let mut rng = crypto_box::rand_core::OsRng;
    let eph_secret_key = SecretKey::generate(&mut rng);
    let eph_public_key = eph_secret_key.public_key();
    let mut eph_secret_key_bytes = eph_secret_key.as_bytes().clone();

    let nonce: [u8; NONCE_LENGTH] = ciphertext.clone()[..NONCE_LENGTH].try_into().unwrap();

    for share_index in 0..public_keys.len() {
        let share = &shares[share_index];
        let pk = PublicKey::from(public_keys[share_index]);
        let esk = SecretKey::from(eph_secret_key_bytes);
        encrypted_shares.push(encrypt_with_given_nonce(esk, pk, share.to_vec(), nonce)?);
    }
    eph_secret_key_bytes.zeroize();

    Ok(EncryptedShareSet {
        encrypted_shares,
        ciphertext,
        eph_public_key,
    })
}

/// Encrypt a given message using crypto_box
pub fn encrypt(
    secret_key: SecretKey,
    public_key: PublicKey,
    plaintext: Vec<u8>,
) -> Result<Vec<u8>, Error> {
    let alice_box = Box::new(&public_key, &secret_key);
    let mut rng = crypto_box::rand_core::OsRng;
    let nonce_bytes = rng.gen::<[u8; NONCE_LENGTH]>();
    let nonce = GenericArray::from_slice(&nonce_bytes);

    let mut ciphertext_with_nonce = nonce_bytes.to_vec();

    let ciphertext = alice_box.encrypt(&nonce, &plaintext[..])?;
    ciphertext_with_nonce.extend(ciphertext);
    Ok(ciphertext_with_nonce)
}

/// Decrypt a given ciphertext using crypto_box
pub fn decrypt(
    secret_key: SecretKey,
    public_key: &PublicKey,
    ciphertext_with_nonce: &Vec<u8>,
) -> Result<Vec<u8>, Error> {
    let ciphertext = &ciphertext_with_nonce[NONCE_LENGTH..];
    let nonce = GenericArray::from_slice(&ciphertext_with_nonce[..NONCE_LENGTH]);

    let bob_box = Box::new(public_key, &secret_key);
    bob_box.decrypt(&nonce, &ciphertext[..])
}

/// Encrypt a given message using crypto_box
/// using a given nonce rather than generating one
pub fn encrypt_with_given_nonce(
    secret_key: SecretKey,
    public_key: PublicKey,
    plaintext: Vec<u8>,
    nonce: [u8; NONCE_LENGTH],
) -> Result<Vec<u8>, Error> {
    let alice_box = Box::new(&public_key, &secret_key);
    let nonce = GenericArray::from_slice(&nonce);
    alice_box.encrypt(&nonce, &plaintext[..])
}

/// Decrypt a given ciphertext using crypto_box
/// using a given nonce rather than attaching one to the ciphertext
pub fn decrypt_with_given_nonce(
    secret_key: SecretKey,
    public_key: &PublicKey,
    ciphertext: &Vec<u8>,
    nonce: [u8; NONCE_LENGTH],
) -> Result<Vec<u8>, Error> {
    let nonce = GenericArray::from_slice(&nonce);
    let bob_box = Box::new(public_key, &secret_key);
    bob_box.decrypt(&nonce, &ciphertext[..])
}

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

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

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

impl From<ShareError> for ShareAndEncryptError {
    fn from(error: ShareError) -> Self {
        ShareAndEncryptError {
            message: error.to_string(),
        }
    }
}

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

    #[test]
    fn encryption() {
        let mut rng = crypto_box::rand_core::OsRng;
        let alice_secret_key = SecretKey::generate(&mut rng);
        let alice_public_key = alice_secret_key.public_key();

        let bob_secret_key = SecretKey::generate(&mut rng);
        let bob_public_key = bob_secret_key.public_key();

        let plaintext = b"hello";

        let ciphertext = encrypt(alice_secret_key, bob_public_key, plaintext.to_vec()).unwrap();
        let decrypted_plaintext = decrypt(bob_secret_key, &alice_public_key, &ciphertext).unwrap();

        assert_eq!(&plaintext[..], &decrypted_plaintext[..]);
    }

    #[test]
    fn encryption_with_given_nonce() {
        let mut rng = crypto_box::rand_core::OsRng;
        let alice_secret_key = SecretKey::generate(&mut rng);
        let alice_public_key = alice_secret_key.public_key();

        let bob_secret_key = SecretKey::generate(&mut rng);
        let bob_public_key = bob_secret_key.public_key();

        let nonce = rng.gen::<[u8; NONCE_LENGTH]>();
        let plaintext = b"hello";

        let ciphertext =
            encrypt_with_given_nonce(alice_secret_key, bob_public_key, plaintext.to_vec(), nonce)
                .unwrap();
        let decrypted_plaintext =
            decrypt_with_given_nonce(bob_secret_key, &alice_public_key, &ciphertext, nonce)
                .unwrap();

        assert_eq!(&plaintext[..], &decrypted_plaintext[..]);
    }

    #[test]
    fn test_share_and_encrypt() {
        let mut rng = crypto_box::rand_core::OsRng;
        let alice_secret_key = SecretKey::generate(&mut rng);
        let alice_public_key = alice_secret_key.public_key();

        let bob_secret_key = SecretKey::generate(&mut rng);
        let bob_public_key = bob_secret_key.public_key();

        let mut public_keys: Vec<[u8; PUBLIC_KEY_LENGTH]> = Vec::new();
        public_keys.push(*bob_public_key.as_bytes());
        public_keys.push(*alice_public_key.as_bytes());

        let original_secret = b"hello";
        let encrypted_share_set =
            share_and_encrypt(public_keys, original_secret[..].to_vec(), 2).unwrap();

        assert_eq!(encrypted_share_set.encrypted_shares.len(), 2);
        assert_eq!(encrypted_share_set.encrypted_shares[0].len(), 73);

        // Now decrypt the shares
        let mut decrypted_shares: Vec<Vec<u8>> = Vec::new();
        decrypted_shares.push(
            decrypt(
                alice_secret_key,
                &encrypted_share_set.eph_public_key,
                &encrypted_share_set.encrypted_shares[1],
            )
            .unwrap(),
        );
        decrypted_shares.push(
            decrypt(
                bob_secret_key,
                &encrypted_share_set.eph_public_key,
                &encrypted_share_set.encrypted_shares[0],
            )
            .unwrap(),
        );

        // Recover secret
        let recovered_secret =
            combine_authenticated(decrypted_shares, encrypted_share_set.ciphertext).unwrap();
        assert_eq!(recovered_secret, b"hello");
    }

    #[test]
    fn test_share_and_encrypt_detached_nonce() {
        let mut rng = crypto_box::rand_core::OsRng;
        let alice_secret_key = SecretKey::generate(&mut rng);
        let alice_public_key = alice_secret_key.public_key();

        let bob_secret_key = SecretKey::generate(&mut rng);
        let bob_public_key = bob_secret_key.public_key();

        let mut public_keys: Vec<[u8; PUBLIC_KEY_LENGTH]> = Vec::new();
        public_keys.push(*bob_public_key.as_bytes());
        public_keys.push(*alice_public_key.as_bytes());

        let original_secret = b"hello";
        let encrypted_share_set =
            share_and_encrypt_detached_nonce(public_keys, original_secret[..].to_vec(), 2).unwrap();

        assert_eq!(encrypted_share_set.encrypted_shares.len(), 2);
        assert_eq!(encrypted_share_set.encrypted_shares[0].len(), 49);

        // Now decrypt the shares
        let nonce = encrypted_share_set.ciphertext.clone()[..NONCE_LENGTH]
            .try_into()
            .unwrap();

        let mut decrypted_shares: Vec<Vec<u8>> = Vec::new();
        decrypted_shares.push(
            decrypt_with_given_nonce(
                alice_secret_key,
                &encrypted_share_set.eph_public_key,
                &encrypted_share_set.encrypted_shares[1],
                nonce,
            )
            .unwrap(),
        );
        decrypted_shares.push(
            decrypt_with_given_nonce(
                bob_secret_key,
                &encrypted_share_set.eph_public_key,
                &encrypted_share_set.encrypted_shares[0],
                nonce,
            )
            .unwrap(),
        );

        // Recover secret
        let recovered_secret =
            combine_authenticated(decrypted_shares, encrypted_share_set.ciphertext).unwrap();
        assert_eq!(recovered_secret, b"hello");
    }
}