//! An implementation of asymmetric cryptography using SECP256k1 ECDH with ChaCha20Poly1305
//! for the actual encryption.

use k256::{SecretKey, ecdh::diffie_hellman, elliptic_curve::sec1::ToEncodedPoint};
use sha2::Sha256;

use crate::{error::CryptoError, key::symmetric::SymmetricKey};

/// A keypair for use in asymmetric encryption operations.
pub struct AsymmetricKeypair {
    /// Private part of the key
    private: SecretKey,
    /// Public part of the key
    public: k256::PublicKey,
}

/// The public key part of an asymmetric keypair.
#[derive(Debug, PartialEq, Eq)]
pub struct PublicKey(k256::PublicKey);

impl AsymmetricKeypair {
    /// Generates a new random keypair
    pub fn new() -> Result<Self, CryptoError> {
        let mut raw_private = [0u8; 32];
        rand::fill(&mut raw_private);
        let sk = SecretKey::from_slice(&raw_private)
            .expect("Key is provided generated with fixed valid size");
        let pk = sk.public_key();

        Ok(Self {
            private: sk,
            public: pk,
        })
    }

    /// Create a new key from existing bytes.
    pub(crate) fn from_bytes(bytes: &[u8]) -> Result<Self, CryptoError> {
        if bytes.len() == 32 {
            let sk = SecretKey::from_slice(&bytes).expect("Key was checked to be a valid size");
            let pk = sk.public_key();
            Ok(Self {
                private: sk,
                public: pk,
            })
        } else {
            Err(CryptoError::InvalidKeySize)
        }
    }

    /// View the raw bytes of the private key of this keypair.
    pub(crate) fn as_raw_private_key(&self) -> Vec<u8> {
        self.private.as_scalar_primitive().to_bytes().to_vec()
    }

    /// Get the public part of this keypair.
    pub fn public_key(&self) -> PublicKey {
        PublicKey(self.public.clone())
    }

    /// Encrypt data for a receiver. First a shared secret is derived using the own private key and
    /// the receivers public key. Then, this shared secret is used for symmetric encryption of the
    /// plaintext. The receiver can decrypt this by generating the same shared secret, using his
    /// own private key and our public key.
    pub fn encrypt(
        &self,
        remote_key: &PublicKey,
        plaintext: &[u8],
    ) -> Result<Vec<u8>, CryptoError> {
        let mut symmetric_key = [0u8; 32];
        diffie_hellman(self.private.to_nonzero_scalar(), remote_key.0.as_affine())
            .extract::<Sha256>(None)
            .expand(&[], &mut symmetric_key)
            .map_err(|_| CryptoError::InvalidKeySize)?;

        let sym_key = SymmetricKey::from_bytes(&symmetric_key)?;

        sym_key.encrypt(plaintext)
    }

    /// Decrypt data from a sender. The remote key must be the public key of the keypair used by
    /// the sender to encrypt this message.
    pub fn decrypt(
        &self,
        remote_key: &PublicKey,
        ciphertext: &[u8],
    ) -> Result<Vec<u8>, CryptoError> {
        let mut symmetric_key = [0u8; 32];
        diffie_hellman(self.private.to_nonzero_scalar(), remote_key.0.as_affine())
            .extract::<Sha256>(None)
            .expand(&[], &mut symmetric_key)
            .map_err(|_| CryptoError::InvalidKeySize)?;

        let sym_key = SymmetricKey::from_bytes(&symmetric_key)?;

        sym_key.decrypt(ciphertext)
    }
}

impl PublicKey {
    /// Import a public key from raw bytes
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, CryptoError> {
        Ok(Self(k256::PublicKey::from_sec1_bytes(bytes)?))
    }

    /// Get the raw bytes of this `PublicKey`, which can be transferred to another party.
    ///
    /// The public key is SEC-1 encoded and compressed.
    pub fn as_bytes(&self) -> Box<[u8]> {
        self.0.to_encoded_point(true).to_bytes()
    }
}

#[cfg(test)]
mod tests {
    /// Export a public key and import it later
    #[test]
    fn import_public_key() {
        let kp = super::AsymmetricKeypair::new().expect("Can generate new keypair");
        let pk1 = kp.public_key();
        let pk_bytes = pk1.as_bytes();
        let pk2 = super::PublicKey::from_bytes(&pk_bytes).expect("Can import public key");

        assert_eq!(pk1, pk2);
    }
    /// Make sure 2 random keypairs derive the same shared secret (and thus encryption key), by
    /// encrypting a random message, decrypting it, and verifying it matches.
    #[test]
    fn encrypt_and_decrypt() {
        let kp1 = super::AsymmetricKeypair::new().expect("Can generate new keypair");
        let kp2 = super::AsymmetricKeypair::new().expect("Can generate new keypair");

        let pk1 = kp1.public_key();
        let pk2 = kp2.public_key();

        let message = b"this is a random message to encrypt and decrypt";

        let enc = kp1.encrypt(&pk2, message).expect("Can encrypt message");
        let dec = kp2.decrypt(&pk1, &enc).expect("Can decrypt message");

        assert_eq!(message.as_slice(), dec.as_slice());
    }

    /// Use a different public key for decrypting than the expected one, this should fail the
    /// decryption process as we use AEAD encryption with the symmetric key.
    #[test]
    fn decrypt_with_wrong_key() {
        let kp1 = super::AsymmetricKeypair::new().expect("Can generate new keypair");
        let kp2 = super::AsymmetricKeypair::new().expect("Can generate new keypair");
        let kp3 = super::AsymmetricKeypair::new().expect("Can generate new keypair");

        let pk2 = kp2.public_key();
        let pk3 = kp3.public_key();

        let message = b"this is a random message to encrypt and decrypt";

        let enc = kp1.encrypt(&pk2, message).expect("Can encrypt message");
        let dec = kp2.decrypt(&pk3, &enc);

        assert!(dec.is_err());
    }
}