bitbottle 0.9.1

use aes::{Aes128, BlockDecrypt, BlockEncrypt, NewBlockCipher};
use aes_gcm::{AeadInPlace, Aes128Gcm, NewAead};
use chacha20::XChaCha20;
use chacha20poly1305::XChaCha20Poly1305;
use cipher::{BlockCipherKey, NewCipher, StreamCipher};
use digest::generic_array::GenericArray;
use getrandom::getrandom;
use num_enum::TryFromPrimitive;
use crate::bottle_error::{BottleError, BottleResult};


#[allow(non_camel_case_types)]
#[derive(Clone, Copy, Debug, PartialEq, TryFromPrimitive)]
#[repr(u8)]
pub enum EncryptionAlgorithm {
    AES_128_GCM = 1,
    XCHACHA20_POLY1305 = 2,
}

const ALGORITHMS: [(EncryptionAlgorithm, &'static dyn Encryption); 2] = [
    (EncryptionAlgorithm::AES_128_GCM, &AES_ENCRYPTION),
    (EncryptionAlgorithm::XCHACHA20_POLY1305, &XCHACHA20_POLY1305_ENCRYPTION),
];

pub fn encryption_for(algorithm: EncryptionAlgorithm) -> BottleResult<&'static dyn Encryption> {
    ALGORITHMS.iter().find(|p| p.0 == algorithm).ok_or(BottleError::UnknownEncryption).map(|row| row.1)
}


pub fn cprng(buffer: &mut [u8]) -> BottleResult<()> {
    getrandom(buffer).map_err(|_| BottleError::NoRandom)
}

/// An encryption algorithm for `EncryptedBottle`.
pub trait Encryption {
    /// Full canonical name for this encryption method.
    fn name(&self) -> &'static str;

    /// Byte length of the secret key used for encryption.
    fn key_length(&self) -> usize;

    /// When using argon (or other KDF), how much material should be
    /// generated? For some algorithms (like xchacha20poly1305), the stream
    /// cipher also requires a nonce, so this can be larger than the usual
    /// key size.
    fn key_key_length(&self) -> usize;

    /// Byte length of the (CPRNG) nonce to be generated for each block.
    fn nonce_length(&self) -> usize;

    /// Byte length of the metadata preceding each encrypted block. This
    /// will encode things like a nonce and/or MAC.
    fn metadata_length(&self) -> usize;

    /// Create a cipher that can encrypt/decrypt blocks for this method.
    fn create(&self, key: &[u8]) -> Box<dyn EncryptionCipher>;

    /// For use with passwords: encrypt a randomly-generated key with the
    /// key generated from the password by argon.
    fn encrypt_key(&self, out_key: &mut [u8], key: &[u8], key_key: &[u8]);

    /// For use with passwords: decrypt a randomly-generated key with the
    /// key generated from the password by argon.
    fn decrypt_key(&self, out_key: &mut [u8], key: &[u8], key_key: &[u8]);
}

pub trait EncryptionCipher {
    /// Reference back to the `Encryption` that made me.
    fn encryption(&self) -> &'static dyn Encryption;

    /// Encrypt a block of data in place. If the encryption data has attached
    /// metadata, it will be put into `metadata`.
    fn encrypt_block(&mut self, data: &mut [u8], metadata: &mut [u8]) -> BottleResult<()>;

    /// Decrypt a block of data in place. Any metadata (associated data or
    /// MAC) should be in `metadata` in the same format returned from
    /// `encrypt_block`.
    fn decrypt_block(&mut self, data: &mut [u8], metadata: &[u8]) -> BottleResult<()>;
}


/// AES-128-GCM
pub struct AesEncryption {
}

impl AesEncryption {
    fn tag_length(&self) -> usize { 16 }
}

pub const AES_ENCRYPTION: AesEncryption = AesEncryption {};

impl Encryption for AesEncryption {
    fn name(&self) -> &'static str { "AES-128-GCM" }
    fn key_length(&self) -> usize { 16 }
    fn key_key_length(&self) -> usize { 16 }
    fn nonce_length(&self) -> usize { 12 }
    fn metadata_length(&self) -> usize { self.nonce_length() + self.tag_length() }

    fn create(&self, key: &[u8]) -> Box<dyn EncryptionCipher> {
        Box::new(AesEncryptionCipher::new(key))
    }

    fn encrypt_key(&self, out_key: &mut [u8], key: &[u8], key_key: &[u8]) {
        let mut encrypted_key: BlockCipherKey<Aes128> = GenericArray::clone_from_slice(key);
        let cipher = Aes128::new(GenericArray::from_slice(key_key));
        cipher.encrypt_block(&mut encrypted_key);
        out_key.copy_from_slice(&encrypted_key);
    }

    fn decrypt_key(&self, out_key: &mut [u8], key: &[u8], key_key: &[u8]) {
        let mut decrypted_key: BlockCipherKey<Aes128> = GenericArray::clone_from_slice(key);
        let cipher = Aes128::new(GenericArray::from_slice(key_key));
        cipher.decrypt_block(&mut decrypted_key);
        out_key.copy_from_slice(&decrypted_key);
    }
}

pub struct AesEncryptionCipher {
    cipher: Aes128Gcm,
}

impl AesEncryptionCipher {
    fn new(key: &[u8]) -> AesEncryptionCipher {
        let cipher = Aes128Gcm::new(aes_gcm::Key::from_slice(key));
        AesEncryptionCipher { cipher }
    }
}

impl EncryptionCipher for AesEncryptionCipher {
    fn encryption(&self) -> &'static dyn Encryption {
        &AES_ENCRYPTION
    }

    fn encrypt_block(&mut self, data: &mut [u8], metadata: &mut [u8]) -> BottleResult<()> {
        assert!(metadata.len() == AES_ENCRYPTION.metadata_length());
        let nonce_len = AES_ENCRYPTION.nonce_length();
        let tag_len = AES_ENCRYPTION.tag_length();

        let assoc = [0u8; 0];
        let mut nonce_buffer = [0u8; 32];
        let nonce = &mut nonce_buffer[0 .. nonce_len];
        cprng(nonce)?;
        let nonce = GenericArray::clone_from_slice(nonce);

        let tag = self.cipher.encrypt_in_place_detached(&nonce, &assoc, data).map_err(|_| {
            BottleError::CipherError
        })?;
        metadata[0 .. nonce_len].copy_from_slice(&nonce);
        metadata[nonce_len .. nonce_len + tag_len].copy_from_slice(&tag);
        Ok(())
    }

    fn decrypt_block(&mut self, data: &mut [u8], metadata: &[u8]) -> BottleResult<()> {
        assert!(metadata.len() == AES_ENCRYPTION.metadata_length());
        let nonce_len = AES_ENCRYPTION.nonce_length();
        let tag_len = AES_ENCRYPTION.tag_length();

        let assoc = [0u8; 0];
        let nonce = GenericArray::from_slice(&metadata[0 .. nonce_len]);
        let tag = GenericArray::from_slice(&metadata[nonce_len .. nonce_len + tag_len]);

        self.cipher.decrypt_in_place_detached(nonce, &assoc, data, tag).map_err(|_| {
            BottleError::CipherError
        })?;
        Ok(())
    }
}


/// XCHACHA20-POLY1305
pub struct Xchacha2OPoly1305Encryption {
}

impl Xchacha2OPoly1305Encryption {
    fn tag_length(&self) -> usize { 16 }
}

pub const XCHACHA20_POLY1305_ENCRYPTION: Xchacha2OPoly1305Encryption = Xchacha2OPoly1305Encryption {};

impl Encryption for Xchacha2OPoly1305Encryption {
    fn name(&self) -> &'static str { "XCHACHA20-POLY1305" }
    fn key_length(&self) -> usize { 32 }
    fn key_key_length(&self) -> usize { 56 }
    fn nonce_length(&self) -> usize { 24 }
    fn metadata_length(&self) -> usize { self.nonce_length() + self.tag_length() }

    fn create(&self, key: &[u8]) -> Box<dyn EncryptionCipher> {
        Box::new(Xchacha2OPoly1305EncryptionCipher::new(key))
    }

    fn encrypt_key(&self, out_key: &mut [u8], key: &[u8], key_key: &[u8]) {
        let mut encrypted_key: chacha20::Key = GenericArray::clone_from_slice(key);
        let (key_key, nonce) = key_key.split_at(32);
        let mut cipher = XChaCha20::new(GenericArray::from_slice(key_key), GenericArray::from_slice(nonce));
        cipher.apply_keystream(&mut encrypted_key);
        out_key.copy_from_slice(&encrypted_key);
    }

    fn decrypt_key(&self, out_key: &mut [u8], key: &[u8], key_key: &[u8]) {
        let mut decrypted_key: chacha20::Key = GenericArray::clone_from_slice(key);
        let (key_key, nonce) = key_key.split_at(32);
        let mut cipher = XChaCha20::new(GenericArray::from_slice(key_key), GenericArray::from_slice(nonce));
        cipher.apply_keystream(&mut decrypted_key);
        out_key.copy_from_slice(&decrypted_key);
    }
}

pub struct Xchacha2OPoly1305EncryptionCipher {
    cipher: XChaCha20Poly1305,
}

impl Xchacha2OPoly1305EncryptionCipher {
    fn new(key: &[u8]) -> Xchacha2OPoly1305EncryptionCipher {
        let cipher = XChaCha20Poly1305::new(chacha20poly1305::Key::from_slice(key));
        Xchacha2OPoly1305EncryptionCipher { cipher }
    }
}

impl EncryptionCipher for Xchacha2OPoly1305EncryptionCipher {
    fn encryption(&self) -> &'static dyn Encryption {
        &XCHACHA20_POLY1305_ENCRYPTION
    }

    fn encrypt_block(&mut self, data: &mut [u8], metadata: &mut [u8]) -> BottleResult<()> {
        assert!(metadata.len() == XCHACHA20_POLY1305_ENCRYPTION.metadata_length());
        let nonce_len = XCHACHA20_POLY1305_ENCRYPTION.nonce_length();
        let tag_len = XCHACHA20_POLY1305_ENCRYPTION.tag_length();

        let assoc = [0u8; 0];
        let mut nonce_buffer = [0u8; 32];
        let nonce = &mut nonce_buffer[0 .. nonce_len];
        cprng(nonce)?;
        let nonce = GenericArray::clone_from_slice(nonce);

        let tag = self.cipher.encrypt_in_place_detached(&nonce, &assoc, data).map_err(|_| {
            BottleError::CipherError
        })?;
        metadata[0 .. nonce_len].copy_from_slice(&nonce);
        metadata[nonce_len .. nonce_len + tag_len].copy_from_slice(&tag);
        Ok(())
    }

    fn decrypt_block(&mut self, data: &mut [u8], metadata: &[u8]) -> BottleResult<()> {
        assert!(metadata.len() == XCHACHA20_POLY1305_ENCRYPTION.metadata_length());
        let nonce_len = XCHACHA20_POLY1305_ENCRYPTION.nonce_length();
        let tag_len = XCHACHA20_POLY1305_ENCRYPTION.tag_length();

        let assoc = [0u8; 0];
        let nonce = GenericArray::from_slice(&metadata[0 .. nonce_len]);
        let tag = GenericArray::from_slice(&metadata[nonce_len .. nonce_len + tag_len]);

        self.cipher.decrypt_in_place_detached(nonce, &assoc, data, tag).map_err(|_| {
            BottleError::CipherError
        })?;
        Ok(())
    }
}


#[cfg(test)]
mod test {
    use hex::decode;
    use crate::encryption::Xchacha2OPoly1305EncryptionCipher;

    use super::{AES_ENCRYPTION, AesEncryptionCipher, Encryption, EncryptionCipher, XCHACHA20_POLY1305_ENCRYPTION};

    #[test]
    fn aes128gcm_encrypt_block() {
        let key = [0u8; 16];
        let message = b"time can stand still inside my funeral home";
        let mut data = message.clone();
        let mut metadata = [0u8; 28];

        let mut cipher = AesEncryptionCipher::new(&key);
        cipher.encrypt_block(&mut data, &mut metadata).unwrap();
        assert_ne!(&data, message);
        cipher.decrypt_block(&mut data, &metadata).unwrap();
        assert_eq!(&data, message);
    }

    #[test]
    fn aes128gcm_encrypt_key() {
        let key = decode("173928ccc9d4897288a8674bd1327692").unwrap();
        let key_key = decode("9e57cde0a98c4baf66575733cd07f6a8").unwrap();

        let mut encrypted_key = vec![0u8; key.len()];
        AES_ENCRYPTION.encrypt_key(&mut encrypted_key, &key, &key_key);
        assert_ne!(&key, encrypted_key.as_slice());
        assert_ne!(key, key_key);

        let mut decrypted_key = vec![0u8; key.len()];
        AES_ENCRYPTION.decrypt_key(&mut decrypted_key, &encrypted_key, &key_key);
        assert_eq!(&key, decrypted_key.as_slice());
    }

    #[test]
    fn xchacha20poly1305_encrypt_block() {
        let key = [0u8; 32];
        let message = b"time can stand still inside my funeral home";
        let mut data = message.clone();
        let mut metadata = [0u8; 40];

        let mut cipher = Xchacha2OPoly1305EncryptionCipher::new(&key);
        cipher.encrypt_block(&mut data, &mut metadata).unwrap();
        assert_ne!(&data, message);
        cipher.decrypt_block(&mut data, &metadata).unwrap();
        assert_eq!(&data, message);
    }

    #[test]
    fn xchacha20poly1305_encrypt_key() {
        let key = decode("5e2f3c897085570344bffb0daaf5955ed0d2b2b050c982656786a8f33e65d866").unwrap();
        let key_key = decode(String::new() +
            "a68d5ba8c4a3c801307baed32675a8619b8fd95f03662640abe732f0e66bbea5" +
            "f6ce7a24086dce7eccd88eff170e6fbcadbdc8df5424d70b"
        ).unwrap();

        let mut encrypted_key = vec![0u8; key.len()];
        XCHACHA20_POLY1305_ENCRYPTION.encrypt_key(&mut encrypted_key, &key, &key_key);
        assert_ne!(&key, encrypted_key.as_slice());
        assert_ne!(key, key_key);

        let mut decrypted_key = vec![0u8; key.len()];
        XCHACHA20_POLY1305_ENCRYPTION.decrypt_key(&mut decrypted_key, &encrypted_key, &key_key);
        assert_eq!(&key, decrypted_key.as_slice());
    }
}