crypto-seal 0.4.0

use crate::{Result, error::Error};
#[cfg(feature = "std")]
use aes_gcm::aead::stream::{DecryptorBE32, EncryptorBE32};
use aes_gcm::{
    Aes256Gcm, Key, KeyInit, Nonce,
    aead::{Aead, Payload},
};
use alloc::string::String;
use alloc::vec::Vec;
use core::hash::Hash;
use curve25519_dalek::edwards::CompressedEdwardsY;
use ed25519_dalek::{Signature, SigningKey};
use hkdf::Hkdf;
use hmac::{Hmac, Mac};
use rand::RngCore;
use rand::rngs::OsRng;
use serde::{Deserialize, Deserializer, Serialize};
use sha2::Digest;
use sha2::Sha256;
use sha2::Sha512;
use signature::{Signer as _, Verifier as _};
#[cfg(feature = "std")]
use std::io;
use zeroize::{Zeroize, Zeroizing};

type HmacSha256 = Hmac<Sha256>;

/// Container of private keys
///
/// The following is supported
///
/// - [`ed25519_dalek`]
/// - [`k256`]
/// - [`p256`]
/// - [`p384`]
/// - [`aes_gcm::Aes256Gcm`]
#[derive(Clone)]
pub enum PrivateKey {
    Ed25519(ed25519_dalek::SigningKey),
    Secp256k1(k256::ecdsa::SigningKey),
    P256(p256::ecdsa::SigningKey),
    P384(p384::ecdsa::SigningKey),
    Aes256([u8; 32]),
}

impl core::fmt::Debug for PrivateKey {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let ty = match self {
            PrivateKey::Ed25519(_) => "ed25519",
            PrivateKey::Secp256k1(_) => "secp256k1",
            PrivateKey::P256(_) => "p256",
            PrivateKey::P384(_) => "p384",
            PrivateKey::Aes256(_) => "aes256",
        };

        write!(f, "{ty}")
    }
}

impl Default for PrivateKey {
    fn default() -> Self {
        Self::new_with(PrivateKeyType::default())
    }
}

impl Zeroize for PrivateKey {
    fn zeroize(&mut self) {
        match self {
            PrivateKey::Ed25519(key) => *key = SigningKey::from_bytes(&[0u8; 32]),
            PrivateKey::Secp256k1(key) => {
                if let Ok(dummy) = k256::ecdsa::SigningKey::from_slice(&[1u8; 32]) {
                    *key = dummy
                }
            }
            PrivateKey::P256(key) => {
                if let Ok(dummy) = p256::ecdsa::SigningKey::from_slice(&[1u8; 32]) {
                    *key = dummy
                }
            }
            PrivateKey::P384(key) => {
                if let Ok(dummy) = p384::ecdsa::SigningKey::from_slice(&[1u8; 48]) {
                    *key = dummy
                }
            }
            PrivateKey::Aes256(key) => key.zeroize(),
        }
    }
}

impl Drop for PrivateKey {
    fn drop(&mut self) {
        self.zeroize()
    }
}

/// Container of public keys
///
/// The following is supported
///
/// - [`ed25519_dalek`]
/// - [`k256`]
/// - [`p256`]
/// - [`p384`]
#[derive(Clone, Copy)]
pub enum PublicKey {
    Ed25519(ed25519_dalek::VerifyingKey),
    Secp256k1(k256::ecdsa::VerifyingKey),
    P256(p256::ecdsa::VerifyingKey),
    P384(p384::ecdsa::VerifyingKey),
}

impl PartialEq for PublicKey {
    fn eq(&self, other: &Self) -> bool {
        self.encode() == other.encode()
    }
}

impl Eq for PublicKey {}

impl Hash for PublicKey {
    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
        self.encode().hash(state);
    }
}

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

impl Ord for PublicKey {
    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
        self.encode().cmp(&other.encode())
    }
}

impl core::fmt::Debug for PublicKey {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "{}", self)
    }
}

impl core::fmt::Display for PublicKey {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "{}", bs58::encode(self.encode()).into_string())
    }
}

impl Serialize for PublicKey {
    fn serialize<S>(&self, serializer: S) -> core::result::Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        let pk_str = bs58::encode(self.encode()).into_string();
        serializer.serialize_str(&pk_str)
    }
}

impl<'d> Deserialize<'d> for PublicKey {
    fn deserialize<D>(deserializer: D) -> core::result::Result<Self, D::Error>
    where
        D: Deserializer<'d>,
    {
        let pk_str = <String>::deserialize(deserializer)?;
        let bytes = bs58::decode(pk_str)
            .into_vec()
            .map_err(serde::de::Error::custom)?;
        PublicKey::decode(&bytes).map_err(serde::de::Error::custom)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum PublicKeyType {
    /// Ed25519 Public Key
    Ed25519,

    /// Secp256k1 Public Key
    Secp256k1,

    /// NIST P-256 Public Key
    P256,

    /// NIST P-384 Public Key
    P384,
}

impl TryFrom<u8> for PublicKeyType {
    type Error = Error;
    fn try_from(value: u8) -> core::result::Result<Self, Self::Error> {
        match value {
            0xa1 => Ok(PublicKeyType::Ed25519),
            0xb1 => Ok(PublicKeyType::Secp256k1),
            0xd1 => Ok(PublicKeyType::P256),
            0xe1 => Ok(PublicKeyType::P384),
            _ => Err(Error::InvalidPublicKey),
        }
    }
}

impl From<PublicKeyType> for u8 {
    fn from(value: PublicKeyType) -> Self {
        match value {
            PublicKeyType::Ed25519 => 0xa1,
            PublicKeyType::Secp256k1 => 0xb1,
            PublicKeyType::P256 => 0xd1,
            PublicKeyType::P384 => 0xe1,
        }
    }
}

impl From<ed25519_dalek::VerifyingKey> for PublicKey {
    fn from(pk: ed25519_dalek::VerifyingKey) -> Self {
        PublicKey::Ed25519(pk)
    }
}

impl TryFrom<PublicKey> for k256::ecdsa::VerifyingKey {
    type Error = Error;

    fn try_from(value: PublicKey) -> core::result::Result<Self, Self::Error> {
        match value {
            PublicKey::Secp256k1(pk) => Ok(pk),
            _ => Err(Error::InvalidPublicKey),
        }
    }
}

impl TryFrom<PublicKey> for p256::ecdsa::VerifyingKey {
    type Error = Error;

    fn try_from(value: PublicKey) -> core::result::Result<Self, Self::Error> {
        match value {
            PublicKey::P256(pk) => Ok(pk),
            _ => Err(Error::InvalidPublicKey),
        }
    }
}

impl TryFrom<PublicKey> for p384::ecdsa::VerifyingKey {
    type Error = Error;

    fn try_from(value: PublicKey) -> core::result::Result<Self, Self::Error> {
        match value {
            PublicKey::P384(pk) => Ok(pk),
            _ => Err(Error::InvalidPublicKey),
        }
    }
}

impl TryFrom<&PrivateKey> for x25519_dalek::StaticSecret {
    type Error = Error;

    fn try_from(value: &PrivateKey) -> core::result::Result<Self, Self::Error> {
        match value {
            PrivateKey::Ed25519(kp) => {
                let mut hasher: Sha512 = Sha512::new();
                hasher.update(kp.as_bytes());
                let hash = hasher.finalize();
                let mut new_sk: [u8; 32] = [0; 32];
                new_sk.copy_from_slice(&hash[..32]);
                let sk = x25519_dalek::StaticSecret::from(new_sk);
                new_sk.zeroize();
                Ok(sk)
            }
            _ => Err(Error::Unsupported),
        }
    }
}

impl TryFrom<PrivateKey> for x25519_dalek::StaticSecret {
    type Error = Error;

    fn try_from(value: PrivateKey) -> core::result::Result<Self, Self::Error> {
        TryFrom::try_from(&value)
    }
}

impl TryFrom<PublicKey> for ed25519_dalek::VerifyingKey {
    type Error = Error;

    fn try_from(value: PublicKey) -> core::result::Result<Self, Self::Error> {
        match value {
            PublicKey::Ed25519(pk) => Ok(pk),
            _ => Err(Error::InvalidPublicKey),
        }
    }
}

impl TryFrom<PublicKey> for x25519_dalek::PublicKey {
    type Error = Error;

    fn try_from(value: PublicKey) -> core::result::Result<Self, Self::Error> {
        match value {
            PublicKey::Ed25519(pk) => {
                let ep = CompressedEdwardsY(pk.to_bytes())
                    .decompress()
                    .ok_or(Error::Unsupported)?; //Note: This should not error here
                let mon = ep.to_montgomery();
                Ok(x25519_dalek::PublicKey::from(mon.0))
            }
            _ => Err(Error::InvalidPublicKey),
        }
    }
}

impl PublicKey {
    pub fn from_bytes(key_type: PublicKeyType, bytes: &[u8]) -> Result<PublicKey> {
        match key_type {
            PublicKeyType::Ed25519 => {
                let bytes: [u8; 32] = bytes.try_into()?;
                Self::from_ed25519_bytes(&bytes)
            }
            PublicKeyType::Secp256k1 => Self::from_secp256k1_bytes(bytes),
            PublicKeyType::P256 => Ok(PublicKey::P256(p256::ecdsa::VerifyingKey::from_sec1_bytes(
                bytes,
            )?)),
            PublicKeyType::P384 => Ok(PublicKey::P384(p384::ecdsa::VerifyingKey::from_sec1_bytes(
                bytes,
            )?)),
        }
    }

    pub fn from_ed25519_bytes(bytes: &[u8; 32]) -> Result<PublicKey> {
        let pk = ed25519_dalek::VerifyingKey::from_bytes(bytes)?;
        Ok(PublicKey::Ed25519(pk))
    }

    pub fn from_secp256k1_bytes(bytes: &[u8]) -> Result<PublicKey> {
        let public_key = k256::ecdsa::VerifyingKey::from_sec1_bytes(bytes)?;
        Ok(PublicKey::Secp256k1(public_key))
    }

    pub fn decode(bytes: &[u8]) -> Result<PublicKey> {
        let (ktype, key) = bytes.split_first().ok_or(Error::InvalidPublicKey)?;
        Self::from_bytes((*ktype).try_into()?, key)
    }

    pub fn encode(&self) -> Vec<u8> {
        let mut data = Vec::new();
        data.push(self.key_type().into());
        data.extend(self.to_bytes());
        data
    }

    /// Convert the [`PublicKey`] to a byte array
    pub fn to_bytes(&self) -> Vec<u8> {
        match self {
            PublicKey::Ed25519(public_key) => public_key.to_bytes().to_vec(),
            PublicKey::Secp256k1(public_key) => {
                public_key.to_encoded_point(true).as_bytes().to_vec()
            }
            PublicKey::P256(public_key) => public_key.to_encoded_point(true).as_bytes().to_vec(),
            PublicKey::P384(public_key) => public_key.to_encoded_point(true).as_bytes().to_vec(),
        }
    }

    pub fn key_type(&self) -> PublicKeyType {
        match self {
            PublicKey::Ed25519(_) => PublicKeyType::Ed25519,
            PublicKey::Secp256k1(_) => PublicKeyType::Secp256k1,
            PublicKey::P256(_) => PublicKeyType::P256,
            PublicKey::P384(_) => PublicKeyType::P384,
        }
    }
}

impl PublicKey {
    /// Verify the signature of the data provided using [`PrivateKey`]
    pub fn verify(&self, data: &[u8], signature: &[u8]) -> Result<()> {
        match self {
            PublicKey::Ed25519(pubkey) => {
                let signature = Signature::from_bytes(signature.try_into()?);
                pubkey.verify(data, &signature)?;
                Ok(())
            }
            PublicKey::Secp256k1(pubkey) => {
                let sig = k256::ecdsa::Signature::from_slice(signature)?;
                pubkey.verify(data, &sig)?;
                Ok(())
            }
            PublicKey::P256(pubkey) => {
                let sig = p256::ecdsa::Signature::from_slice(signature)?;
                pubkey.verify(data, &sig)?;
                Ok(())
            }
            PublicKey::P384(pubkey) => {
                let sig = p384::ecdsa::Signature::from_slice(signature)?;
                pubkey.verify(data, &sig)?;
                Ok(())
            }
        }
    }
}

#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
impl PublicKey {
    /// Verify the signature of the data from [`std::io::Read`] using [`PrivateKey`]
    pub fn verify_reader(&self, reader: &mut impl io::Read, signature: &[u8]) -> Result<()> {
        let mut data = Vec::new();
        reader.read_to_end(&mut data)?;
        self.verify(&data, signature)
    }
}
/// [`PrivateKey`] Types
#[derive(Debug, Copy, Clone, Default)]
pub enum PrivateKeyType {
    /// ED25519 Private Key
    #[default]
    Ed25519,

    /// AES-256 Private Key
    Aes256,

    /// Secp256k1 Private Key
    Secp256k1,

    /// NIST P-256 Private Key
    P256,

    /// NIST P-384 Private Key
    P384,
}

#[cfg(feature = "std")]
const WRITE_BUFFER_SIZE: usize = 512;
#[cfg(feature = "std")]
const READ_BUFFER_SIZE: usize = 528;
const NONCE_LEN: usize = 12;
const SALT_LEN: usize = 16;
const ENCRYPT_INFO: &[u8] = b"crypto-seal:aes256-gcm:v1";
#[cfg(feature = "std")]
const ENCRYPT_STREAM_INFO: &[u8] = b"crypto-seal:aes256-gcm-stream:v1";
const MAC_INFO: &[u8] = b"crypto-seal:hmac-sha256:v1";

impl TryFrom<u8> for PrivateKeyType {
    type Error = Error;
    fn try_from(value: u8) -> core::result::Result<Self, Self::Error> {
        match value {
            0xa1 => Ok(PrivateKeyType::Ed25519),
            0xb1 => Ok(PrivateKeyType::Secp256k1),
            0xc1 => Ok(PrivateKeyType::Aes256),
            0xd1 => Ok(PrivateKeyType::P256),
            0xe1 => Ok(PrivateKeyType::P384),
            _ => Err(Error::InvalidPrivateKey),
        }
    }
}

impl From<PrivateKeyType> for u8 {
    fn from(value: PrivateKeyType) -> Self {
        match value {
            PrivateKeyType::Ed25519 => 0xa1,
            PrivateKeyType::Secp256k1 => 0xb1,
            PrivateKeyType::Aes256 => 0xc1,
            PrivateKeyType::P256 => 0xd1,
            PrivateKeyType::P384 => 0xe1,
        }
    }
}

impl PrivateKey {
    /// Generates a new [`PrivateKey`] with randomly generated key
    pub fn new() -> Self {
        Self::default()
    }

    /// Generate a [`PrivateKey`] using [`PrivateKeyType`]
    pub fn new_with(key_type: PrivateKeyType) -> Self {
        match key_type {
            PrivateKeyType::Ed25519 => PrivateKey::Ed25519(SigningKey::generate(&mut OsRng)),
            PrivateKeyType::Aes256 => {
                let key_sized = generate::<32>();
                PrivateKey::Aes256(key_sized)
            }
            PrivateKeyType::Secp256k1 => {
                PrivateKey::Secp256k1(k256::ecdsa::SigningKey::random(&mut OsRng))
            }
            PrivateKeyType::P256 => PrivateKey::P256(p256::ecdsa::SigningKey::random(&mut OsRng)),
            PrivateKeyType::P384 => PrivateKey::P384(p384::ecdsa::SigningKey::random(&mut OsRng)),
        }
    }

    /// Import private key which is identified with [`PrivateKeyType`]
    pub fn import(key_type: PrivateKeyType, key: Vec<u8>) -> Result<Self> {
        let key = zeroize::Zeroizing::new(key);
        match key_type {
            PrivateKeyType::Ed25519 => {
                let key: [u8; 32] = key.as_slice().try_into()?;
                Ok(PrivateKey::Ed25519(ed25519_dalek::SigningKey::from_bytes(
                    &key,
                )))
            }
            PrivateKeyType::Aes256 => key
                .as_slice()
                .try_into()
                .map(PrivateKey::Aes256)
                .map_err(Error::from),
            PrivateKeyType::Secp256k1 => k256::ecdsa::SigningKey::from_slice(&key)
                .map(PrivateKey::Secp256k1)
                .map_err(Error::from),
            PrivateKeyType::P256 => p256::ecdsa::SigningKey::from_slice(&key)
                .map(PrivateKey::P256)
                .map_err(Error::from),
            PrivateKeyType::P384 => p384::ecdsa::SigningKey::from_slice(&key)
                .map(PrivateKey::P384)
                .map_err(Error::from),
        }
    }

    /// Imports a private key from bytes prefixed with a [`PrivateKeyType`] identifier
    pub fn decode<B: AsRef<[u8]>>(bytes: B) -> Result<PrivateKey> {
        let (ktype, key) = bytes
            .as_ref()
            .split_first()
            .ok_or(Error::InvalidPrivateKey)?;
        Self::import((*ktype).try_into()?, key.to_vec())
    }

    /// Exports the keys out as bytes.
    pub fn to_bytes(&self) -> Vec<u8> {
        match self {
            PrivateKey::Ed25519(kp) => kp.to_bytes().to_vec(),
            PrivateKey::Secp256k1(sk) => sk.to_bytes().as_slice().to_vec(),
            PrivateKey::P256(sk) => sk.to_bytes().as_slice().to_vec(),
            PrivateKey::P384(sk) => sk.to_bytes().as_slice().to_vec(),
            PrivateKey::Aes256(key) => key.to_vec(),
        }
    }

    /// Exports the key out with an identifier
    pub fn encode(&self) -> Vec<u8> {
        let mut data = Vec::new();
        data.push(self.key_type().into());
        data.extend(self.to_bytes());
        data
    }

    /// Provides the [`PrivateKeyType`] of the [`PrivateKey`]
    pub fn key_type(&self) -> PrivateKeyType {
        match self {
            PrivateKey::Aes256(_) => PrivateKeyType::Aes256,
            PrivateKey::Ed25519(_) => PrivateKeyType::Ed25519,
            PrivateKey::Secp256k1(_) => PrivateKeyType::Secp256k1,
            PrivateKey::P256(_) => PrivateKeyType::P256,
            PrivateKey::P384(_) => PrivateKeyType::P384,
        }
    }

    /// Provides the [`PublicKey`] of the [`PrivateKey`]
    /// Note: This will only work with asymmetric keys. Any symmetric keys will
    ///       return [`Error::Unsupported`]
    pub fn public_key(&self) -> Result<PublicKey> {
        match self {
            PrivateKey::Aes256(_) => Err(Error::Unsupported),
            PrivateKey::Ed25519(key) => Ok(key.verifying_key().into()),
            PrivateKey::Secp256k1(key) => Ok(PublicKey::Secp256k1(*key.verifying_key())),
            PrivateKey::P256(key) => Ok(PublicKey::P256(*key.verifying_key())),
            PrivateKey::P384(key) => Ok(PublicKey::P384(*key.verifying_key())),
        }
    }

    /// Sign the data provided using [`PrivateKey`]
    pub fn sign<B: AsRef<[u8]>>(&self, data: B) -> Result<Vec<u8>> {
        let data = data.as_ref();
        match self {
            PrivateKey::Aes256(key) => {
                let mac_key = derive_key(key, &[], MAC_INFO)?;
                let mut mac = <HmacSha256 as Mac>::new_from_slice(&*mac_key)
                    .map_err(|_| Error::EncryptionError)?;
                mac.update(data);
                Ok(mac.finalize().into_bytes().to_vec())
            }
            PrivateKey::Ed25519(key) => {
                let signature = key.sign(data);
                Ok(signature.to_bytes().to_vec())
            }
            PrivateKey::Secp256k1(key) => {
                let signature: k256::ecdsa::Signature = key.try_sign(data)?;
                Ok(signature.to_vec())
            }
            PrivateKey::P256(key) => {
                let signature: p256::ecdsa::Signature = key.try_sign(data)?;
                Ok(signature.to_vec())
            }
            PrivateKey::P384(key) => {
                let signature: p384::ecdsa::Signature = key.try_sign(data)?;
                Ok(signature.to_vec())
            }
        }
    }

    /// Sign the data from [`std::io::Read`] using [`PrivateKey`]
    #[cfg(feature = "std")]
    #[cfg_attr(docsrs, doc(cfg(feature = "std")))]
    pub fn sign_reader(&self, reader: &mut impl io::Read) -> Result<Vec<u8>> {
        match self {
            PrivateKey::Aes256(key) => {
                let mac_key = derive_key(key, &[], MAC_INFO)?;
                let mut mac = <HmacSha256 as Mac>::new_from_slice(&*mac_key)
                    .map_err(|_| Error::EncryptionError)?;
                let mut buffer = [0u8; WRITE_BUFFER_SIZE];
                loop {
                    match reader.read(&mut buffer) {
                        Ok(0) => break,
                        Ok(n) => mac.update(&buffer[..n]),
                        Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
                        Err(e) => return Err(Error::from(e)),
                    }
                }
                Ok(mac.finalize().into_bytes().to_vec())
            }
            _ => {
                let mut data = Vec::new();
                reader.read_to_end(&mut data)?;
                self.sign(&data)
            }
        }
    }

    /// Verify the signature of the data provided using [`PrivateKey`]
    pub fn verify(&self, data: &[u8], signature: &[u8]) -> Result<()> {
        match self {
            PrivateKey::Aes256(key) => {
                let mac_key = derive_key(key, &[], MAC_INFO)?;
                let mut mac = <HmacSha256 as Mac>::new_from_slice(&*mac_key)
                    .map_err(|_| Error::InvalidSignature)?;
                mac.update(data);
                mac.verify_slice(signature)
                    .map_err(|_| Error::InvalidSignature)
            }
            _ => {
                let public_key = self.public_key()?;
                public_key.verify(data, signature)
            }
        }
    }

    /// Verify the signature of the data from [`std::io::Read`] using [`PrivateKey`]
    #[cfg(feature = "std")]
    #[cfg_attr(docsrs, doc(cfg(feature = "std")))]
    pub fn verify_reader(&self, reader: &mut impl io::Read, signature: &[u8]) -> Result<()> {
        match self {
            PrivateKey::Aes256(key) => {
                let mac_key = derive_key(key, &[], MAC_INFO)?;
                let mut mac = <HmacSha256 as Mac>::new_from_slice(&*mac_key)
                    .map_err(|_| Error::InvalidSignature)?;
                let mut buffer = [0u8; WRITE_BUFFER_SIZE];
                loop {
                    match reader.read(&mut buffer) {
                        Ok(0) => break,
                        Ok(n) => mac.update(&buffer[..n]),
                        Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
                        Err(e) => return Err(Error::from(e)),
                    }
                }
                mac.verify_slice(signature)
                    .map_err(|_| Error::InvalidSignature)
            }
            _ => {
                let public_key = self.public_key()?;
                public_key.verify_reader(reader, signature)
            }
        }
    }
}

#[derive(Default, Copy, Clone, PartialEq)]
pub enum CarrierKeyType {
    /// Use AES256 key
    Direct { key: [u8; 32] },

    /// Use key exchange to generate a shared key
    Exchange { public_key: PublicKey },

    /// Use own private key
    /// > **Note** If public key encryption is used, this will use your own private/public key for key exchange
    /// > otherwise if its AES256, it will encrypt with that key itself
    #[default]
    None,
}

impl PrivateKey {
    /// Encrypt the data using [`PrivateKey`].
    /// If [`PrivateKeyType::Aes256`] is used, the `pubkey` will be ignored
    pub fn encrypt(&self, data: &[u8], pubkey: CarrierKeyType) -> Result<Vec<u8>> {
        self.encrypt_with_aad(data, pubkey, &[])
    }

    pub fn encrypt_with_aad(
        &self,
        data: &[u8],
        pubkey: CarrierKeyType,
        aad: &[u8],
    ) -> Result<Vec<u8>> {
        let ikm = self.fetch_encryption_key(pubkey)?;
        let salt = generate::<SALT_LEN>();
        let key = derive_key(ikm.as_slice(), &salt, ENCRYPT_INFO)?;
        let raw_nonce = generate::<NONCE_LEN>();
        let cipher = Aes256Gcm::new(Key::<Aes256Gcm>::from_slice(&*key));
        let nonce = Nonce::from_slice(&raw_nonce);
        let mut out = cipher
            .encrypt(nonce, Payload { msg: data, aad })
            .map_err(|_| Error::EncryptionError)?;
        out.extend_from_slice(&salt);
        out.extend_from_slice(&raw_nonce);
        Ok(out)
    }

    /// Decrypt the data using [`PrivateKey`].
    /// If [`PrivateKeyType::Aes256`] is used, the `pubkey` will be ignored
    pub fn decrypt(&self, data: &[u8], pubkey: CarrierKeyType) -> Result<Vec<u8>> {
        self.decrypt_with_aad(data, pubkey, &[])
    }

    pub fn decrypt_with_aad(
        &self,
        data: &[u8],
        pubkey: CarrierKeyType,
        aad: &[u8],
    ) -> Result<Vec<u8>> {
        if data.len() < SALT_LEN + NONCE_LEN {
            return Err(Error::DecryptionError);
        }
        let ikm = self.fetch_encryption_key(pubkey)?;
        let (rest, raw_nonce) = data.split_at(data.len() - NONCE_LEN);
        let (ciphertext, salt) = rest.split_at(rest.len() - SALT_LEN);
        let key = derive_key(ikm.as_slice(), salt, ENCRYPT_INFO)?;
        let cipher = Aes256Gcm::new(Key::<Aes256Gcm>::from_slice(&*key));
        let nonce = Nonce::from_slice(raw_nonce);
        cipher
            .decrypt(
                nonce,
                Payload {
                    msg: ciphertext,
                    aad,
                },
            )
            .map_err(|_| Error::DecryptionError)
    }
}

impl PrivateKey {
    /// Encrypt the data stream from [`std::io::Read`] to [`std::io::Write`] using [`PrivateKey`].
    /// If [`PrivateKeyType::Aes256`] is used, the `pubkey` will be ignored
    #[cfg(feature = "std")]
    #[cfg_attr(docsrs, doc(cfg(feature = "std")))]
    pub fn encrypt_stream(
        &self,
        reader: &mut impl io::Read,
        writer: &mut impl io::Write,
        pubkey: CarrierKeyType,
    ) -> Result<()> {
        let ikm = self.fetch_encryption_key(pubkey)?;
        let salt = generate::<SALT_LEN>();
        let key = derive_key(ikm.as_slice(), &salt, ENCRYPT_STREAM_INFO)?;
        let nonce = generate::<7>();

        let cipher = Aes256Gcm::new(Key::<Aes256Gcm>::from_slice(&*key));
        let mut buffer = [0u8; WRITE_BUFFER_SIZE];
        let mut stream = EncryptorBE32::from_aead(cipher, nonce.as_slice().into());
        writer.write_all(&salt)?;
        writer.write_all(&nonce)?;
        loop {
            let read_count = fill(reader, &mut buffer)?;
            if read_count == WRITE_BUFFER_SIZE {
                let ciphertext = stream
                    .encrypt_next(buffer.as_slice())
                    .map_err(|_| Error::EncryptionStreamError)?;
                writer.write_all(&ciphertext)?;
            } else {
                let ciphertext = stream
                    .encrypt_last(&buffer[..read_count])
                    .map_err(|_| Error::EncryptionStreamError)?;
                writer.write_all(&ciphertext)?;
                break;
            }
        }
        Ok(())
    }

    /// Decrypt the data stream from [`std::io::Read`] to [`std::io::Write`] using [`PrivateKey`].
    /// If [`PrivateKeyType::Aes256`] is used, the `pubkey` will be ignored
    #[cfg(feature = "std")]
    #[cfg_attr(docsrs, doc(cfg(feature = "std")))]
    pub fn decrypt_stream(
        &self,
        reader: &mut impl io::Read,
        writer: &mut impl io::Write,
        pubkey: CarrierKeyType,
    ) -> Result<()> {
        let ikm = self.fetch_encryption_key(pubkey)?;
        let mut salt = [0u8; SALT_LEN];
        reader.read_exact(&mut salt)?;
        let key = derive_key(ikm.as_slice(), &salt, ENCRYPT_STREAM_INFO)?;
        let mut nonce = vec![0u8; 7];
        reader.read_exact(&mut nonce)?;

        let cipher = Aes256Gcm::new(Key::<Aes256Gcm>::from_slice(&*key));

        let mut stream = DecryptorBE32::from_aead(cipher, nonce.as_slice().into());
        let mut buffer = [0u8; READ_BUFFER_SIZE];
        loop {
            let read_count = fill(reader, &mut buffer)?;
            if read_count == READ_BUFFER_SIZE {
                let plaintext = stream
                    .decrypt_next(buffer.as_slice())
                    .map_err(|_| Error::DecryptionStreamError)?;
                writer.write_all(&plaintext)?;
            } else {
                let plaintext = stream
                    .decrypt_last(&buffer[..read_count])
                    .map_err(|_| Error::DecryptionStreamError)?;
                writer.write_all(&plaintext)?;
                break;
            }
        }
        writer.flush()?;
        Ok(())
    }

    /// Used internally to obtain the encryption key
    fn fetch_encryption_key(&self, pubkey: CarrierKeyType) -> Result<Zeroizing<Vec<u8>>> {
        match pubkey {
            CarrierKeyType::Direct { key } => Ok(Zeroizing::new(key.to_vec())),
            CarrierKeyType::Exchange { public_key } => match self {
                PrivateKey::Aes256(key) => Ok(Zeroizing::new(key.to_vec())),
                PrivateKey::Secp256k1(sk) => {
                    let peer: k256::ecdsa::VerifyingKey = public_key.try_into()?;
                    let shared =
                        k256::ecdh::diffie_hellman(sk.as_nonzero_scalar(), peer.as_affine());
                    Ok(Zeroizing::new(
                        shared.raw_secret_bytes().as_slice().to_vec(),
                    ))
                }
                PrivateKey::Ed25519(_) => {
                    let static_key: x25519_dalek::StaticSecret = self.try_into()?;
                    let public_key: x25519_dalek::PublicKey = public_key.try_into()?;

                    let enc_key = static_key.diffie_hellman(&public_key);
                    Ok(Zeroizing::new(enc_key.as_bytes().to_vec()))
                }
                PrivateKey::P256(sk) => {
                    let peer: p256::ecdsa::VerifyingKey = public_key.try_into()?;
                    let shared =
                        p256::ecdh::diffie_hellman(sk.as_nonzero_scalar(), peer.as_affine());
                    Ok(Zeroizing::new(
                        shared.raw_secret_bytes().as_slice().to_vec(),
                    ))
                }
                PrivateKey::P384(sk) => {
                    let peer: p384::ecdsa::VerifyingKey = public_key.try_into()?;
                    let shared =
                        p384::ecdh::diffie_hellman(sk.as_nonzero_scalar(), peer.as_affine());
                    Ok(Zeroizing::new(
                        shared.raw_secret_bytes().as_slice().to_vec(),
                    ))
                }
            },
            CarrierKeyType::None => match self {
                PrivateKey::Aes256(key) => Ok(Zeroizing::new(key.to_vec())),
                PrivateKey::Secp256k1(sk) => {
                    let shared = k256::ecdh::diffie_hellman(
                        sk.as_nonzero_scalar(),
                        sk.verifying_key().as_affine(),
                    );
                    Ok(Zeroizing::new(
                        shared.raw_secret_bytes().as_slice().to_vec(),
                    ))
                }
                PrivateKey::Ed25519(_) => {
                    let static_key: x25519_dalek::StaticSecret = self.try_into()?;
                    let public_key: x25519_dalek::PublicKey =
                        x25519_dalek::PublicKey::from(&static_key);
                    let enc_key = static_key.diffie_hellman(&public_key);
                    Ok(Zeroizing::new(enc_key.as_bytes().to_vec()))
                }
                PrivateKey::P256(sk) => {
                    let shared = p256::ecdh::diffie_hellman(
                        sk.as_nonzero_scalar(),
                        sk.verifying_key().as_affine(),
                    );
                    Ok(Zeroizing::new(
                        shared.raw_secret_bytes().as_slice().to_vec(),
                    ))
                }
                PrivateKey::P384(sk) => {
                    let shared = p384::ecdh::diffie_hellman(
                        sk.as_nonzero_scalar(),
                        sk.verifying_key().as_affine(),
                    );
                    Ok(Zeroizing::new(
                        shared.raw_secret_bytes().as_slice().to_vec(),
                    ))
                }
            },
        }
    }
}

#[cfg(feature = "std")]
fn fill(reader: &mut impl io::Read, buffer: &mut [u8]) -> io::Result<usize> {
    let mut filled = 0;
    while filled < buffer.len() {
        match reader.read(&mut buffer[filled..]) {
            Ok(0) => break,
            Ok(n) => filled += n,
            Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
            Err(e) => return Err(e),
        }
    }
    Ok(filled)
}

fn derive_key(ikm: &[u8], salt: &[u8], info: &[u8]) -> Result<Zeroizing<[u8; 32]>> {
    let mut okm = Zeroizing::new([0u8; 32]);
    Hkdf::<Sha256>::new(Some(salt), ikm)
        .expand(info, &mut *okm)
        .map_err(|_| Error::EncryptionError)?;
    Ok(okm)
}

/// Used to generate random amount of data and store it in a Vec with a specific capacity
pub(crate) fn generate<const N: usize>() -> [u8; N] {
    let mut buffer: [u8; N] = [0u8; N];
    OsRng.fill_bytes(&mut buffer);
    buffer
}