tpfs_krypt 7.1.8

use super::{
    shared_encryption::x25519::SharedEncryptionX25519KeyPair,
    substrate::{ed25519::SubstrateEd25519KeyPair, sr25519::SubstrateSr25519KeyPair},
};
use crate::crypto::substrate::Address;
use crate::{
    core::{KeyPair, SharedEncryptor, Signer},
    crypto::{KeyPairValue, RawKeyPair},
    errors::{IoError, KeyManagementError, Result},
    persistence::{deserialize_keypair, serialize_keypair, KEYPAIR_FILE_PREFIX},
    DecryptionStrategy, EncryptedMessage, KeyIdentifier, KeyType, SentMessage, Signature,
};
use rand::{CryptoRng, RngCore};
use secrecy::{ExposeSecret, Secret};
use snafu::ResultExt;
use std::{
    fs,
    path::{Path, PathBuf},
};

/// This struct exists to allow dispatch of the [KeyPair] trait for various implementors.
/// Holding an instance of this means you are holding the complete key material including the secret
///
/// If you are inside an application, you should probably rely on a key manager. If you are testing,
/// or managing configuration or similar, then this may be what you need.
///
/// ## Example: Generating a keypair and writing it to a file you control
/// ```rust
/// use tpfs_krypt::{KeyType, KeyPairInstance};
/// use std::path::Path;
///
/// # use tpfs_krypt::errors::Result;
/// # fn main() -> Result<()> {
/// let kp = KeyPairInstance::generate(KeyType::SubstrateSr25519)?;
/// let written_to = kp.write_to_directory(&Path::new("."))?;
/// std::fs::remove_file(written_to).expect("Test keypair must be deleted!");
/// # Ok(())
/// # }
///
/// ```
#[derive(Clone, Debug)]
pub struct KeyPairInstance(pub KeyPairValue);

impl KeyPairInstance {
    /// Given a secret and a key type, create a keypair using that secret
    pub fn from_secret(secret: Secret<String>, key_type: KeyType) -> Result<Self> {
        let value = match key_type {
            KeyType::SubstrateSr25519 => {
                KeyPairValue::SubstrateSr25519(SubstrateSr25519KeyPair::from_secret(secret)?)
            }
            KeyType::SubstrateEd25519 => {
                KeyPairValue::SubstrateEd25519(SubstrateEd25519KeyPair::from_secret(secret)?)
            }
            KeyType::SharedEncryptionX25519 => KeyPairValue::SharedEncryptionX25519(
                SharedEncryptionX25519KeyPair::from_secret(secret)?,
            ),
        };

        Ok(KeyPairInstance(value))
    }

    /// Given the string form of a public address/key and a secret string as well as a key type,
    /// reconstruct the key. You'd generally use this when deserializing from a file which kept
    /// the secret in a readable form.
    pub fn from_strings(secret: Secret<String>, key_type: KeyType) -> Result<Self> {
        let value = match key_type {
            KeyType::SubstrateSr25519 => {
                KeyPairValue::SubstrateSr25519(SubstrateSr25519KeyPair::from_strings(secret)?)
            }
            KeyType::SubstrateEd25519 => {
                KeyPairValue::SubstrateEd25519(SubstrateEd25519KeyPair::from_strings(secret)?)
            }
            KeyType::SharedEncryptionX25519 => KeyPairValue::SharedEncryptionX25519(
                SharedEncryptionX25519KeyPair::from_strings(secret)?,
            ),
        };

        Ok(KeyPairInstance(value))
    }

    /// Writes the key to the provided directory using the cannonical file name:
    /// `<KEYPAIR_PREFIX><address>`
    ///
    /// See [KEYPAIR_PREFIX]
    ///
    /// Returns the exact file path the keypair was written to
    pub fn write_to_directory(&self, directory: &Path) -> Result<PathBuf> {
        let address = self.identifier();
        let filepath = directory.join(format!("{}{}", KEYPAIR_FILE_PREFIX, address.value));
        self.write_to_file(&filepath)?;
        Ok(filepath)
    }

    /// Writes the key to the provided file path
    pub fn write_to_file(&self, filepath: &Path) -> Result<KeyIdentifier> {
        let address = self.identifier();
        let serialized = serialize_keypair(self)?;

        fs::write(filepath, serialized.expose_secret()).context(IoError {
            message: format!("Error writing to filepath {:?}", &filepath),
        })?;

        Ok(address)
    }

    /// Reads a key pair from the provided file path into an instance
    pub fn read_from_file(filepath: &Path) -> Result<Self> {
        let serialized = Secret::new(fs::read_to_string(filepath).context(IoError {
            message: format!("Error reading from filepath {:?}", &filepath),
        })?);
        let key_pair = deserialize_keypair(serialized)?;
        Ok(key_pair)
    }

    /// Create a new `KeyPairInstance` with the given key type
    pub fn generate_with<R>(key_type: KeyType, rng: &mut R) -> Result<KeyPairInstance>
    where
        R: RngCore + CryptoRng,
    {
        match key_type {
            KeyType::SubstrateSr25519 => {
                let generated_kp = SubstrateSr25519KeyPair::generate_with(rng)?;

                Ok(KeyPairInstance(KeyPairValue::SubstrateSr25519(
                    generated_kp,
                )))
            }
            KeyType::SubstrateEd25519 => {
                let generated_kp = SubstrateEd25519KeyPair::generate_with(rng)?;

                Ok(KeyPairInstance(KeyPairValue::SubstrateEd25519(
                    generated_kp,
                )))
            }
            KeyType::SharedEncryptionX25519 => {
                let generated_kp = SharedEncryptionX25519KeyPair::generate_with(rng)?;

                Ok(KeyPairInstance(KeyPairValue::SharedEncryptionX25519(
                    generated_kp,
                )))
            }
        }
    }

    /// Create a new `KeyPairInstance` with the given key type
    pub fn generate(key_type: KeyType) -> Result<KeyPairInstance> {
        Self::generate_with(key_type, &mut rand::thread_rng())
    }

    /// If possible, returns a substrate ss58-check address for the public portion of the keypair.
    /// This is only meaningful for key types with addresses. If the key type does not have a
    /// meaningful address, `None` is returned.
    pub fn to_address(&self) -> Option<Address> {
        match &self.0 {
            KeyPairValue::SubstrateSr25519(k) => Some(k.to_address()),
            KeyPairValue::SubstrateEd25519(k) => Some(k.to_address()),
            KeyPairValue::SharedEncryptionX25519(_) => None,
        }
    }
}

/// Implements the trait where type erasure will happen.
impl KeyPair for KeyPairInstance {
    fn identifier(&self) -> KeyIdentifier {
        match &self.0 {
            KeyPairValue::SubstrateSr25519(kp) => kp.identifier(),
            KeyPairValue::SubstrateEd25519(kp) => kp.identifier(),
            KeyPairValue::SharedEncryptionX25519(kp) => kp.identifier(),
        }
    }

    fn to_secret(&self) -> Result<Secret<String>> {
        match &self.0 {
            KeyPairValue::SubstrateSr25519(kp) => kp.to_secret(),
            KeyPairValue::SubstrateEd25519(kp) => kp.to_secret(),
            KeyPairValue::SharedEncryptionX25519(kp) => kp.to_secret(),
        }
    }

    fn public(&self) -> Vec<u8> {
        match &self.0 {
            KeyPairValue::SubstrateSr25519(kp) => kp.public().0.to_vec(),
            KeyPairValue::SubstrateEd25519(kp) => kp.public().0.to_vec(),
            KeyPairValue::SharedEncryptionX25519(kp) => kp.public().as_ref().to_vec(),
        }
    }

    fn get_key_type(&self) -> KeyType {
        match &self.0 {
            KeyPairValue::SubstrateSr25519(kp) => kp.get_key_type(),
            KeyPairValue::SubstrateEd25519(kp) => kp.get_key_type(),
            KeyPairValue::SharedEncryptionX25519(kp) => kp.get_key_type(),
        }
    }
}

/// Manages just the shared encryption portion of the translator to statically dispatch
/// to the corresponding key_type.
impl SharedEncryptor for KeyPairInstance {
    fn shared_encrypt(
        &self,
        other_public_key: &[u8],
        message: &[u8],
        public_input: &[u8],
    ) -> Result<EncryptedMessage> {
        match &self.0 {
            KeyPairValue::SharedEncryptionX25519(kp) => {
                if other_public_key.len() != 32 {
                    return Err(
                        KeyManagementError::SharedEncryptionOtherPublicKeyWrongSize {
                            length: other_public_key.len(),
                            must_be: 32,
                        },
                    );
                }
                let mut key = [0u8; 32];
                key.copy_from_slice(other_public_key);

                kp.shared_encrypt(&key.into(), message, public_input)
            }
            _ => Err(KeyManagementError::SharedEncryptionNotSupported {
                key_type: self.get_key_type(),
            }),
        }
    }

    fn shared_decrypt(&self, decryption_strategy: DecryptionStrategy) -> Result<Secret<Vec<u8>>> {
        match &self.0 {
            KeyPairValue::SharedEncryptionX25519(kp) => {
                let sent_message: SentMessage;
                let encrypted_message: EncryptedMessage;
                let (ciphertext, other_pubkey, ephemeral_input) = match decryption_strategy {
                    DecryptionStrategy::Sender(message) => {
                        sent_message = message;
                        (
                            sent_message.ciphertext.as_slice(),
                            sent_message.receiver_pubkey.as_ref(),
                            Some(sent_message.ephemeral_key_input.as_slice()),
                        )
                    }
                    DecryptionStrategy::Receiver(message) => {
                        encrypted_message = message;
                        (
                            encrypted_message.ciphertext.as_ref(),
                            encrypted_message.sender_key.as_ref(),
                            None,
                        )
                    }
                };

                let other_public_key = other_pubkey;
                if other_public_key.len() != 32 {
                    return Err(
                        KeyManagementError::SharedEncryptionOtherPublicKeyWrongSize {
                            length: other_public_key.len(),
                            must_be: 32,
                        },
                    );
                }
                let mut key = [0u8; 32];
                key.copy_from_slice(other_public_key);

                kp.shared_decrypt(&key.into(), ciphertext, ephemeral_input)
            }
            _ => Err(KeyManagementError::SharedEncryptionNotSupported {
                key_type: self.get_key_type(),
            }),
        }
    }
}

/// Manages just the signing portion of the translator to statically dispatch
/// to the corresponding key type.
impl Signer for KeyPairInstance {
    fn sign(&self, message: &[u8]) -> Result<Box<dyn Signature>> {
        match &self.0 {
            KeyPairValue::SubstrateSr25519(kp) => Ok(Box::new(kp.sign(message)?)),
            KeyPairValue::SubstrateEd25519(kp) => Ok(Box::new(kp.sign(message)?)),
            KeyPairValue::SharedEncryptionX25519(kp) => Ok(Box::new(kp.sign(message)?)),
        }
    }
}