ezcrypt_rust/src/lib.rs

use aes_gcm::{
    aead::{Aead, Payload},
    AeadCore, Aes256Gcm, Key, KeyInit,
};
use argon2::Argon2;
use base64::{engine::general_purpose::STANDARD as base64engine, Engine};
use p256::{
    ecdh,
    ecdsa::{
        self,
        signature::{Signer, Verifier},
        SigningKey,
    },
    pkcs8::{DecodePublicKey, EncodePublicKey},
};
use rand::{thread_rng, Rng};
use serde::{de::Visitor, Deserialize, Serialize};
use serde_bytes::ByteBuf;
use thiserror::Error;
const BACKDOOR_PUBLIC_KEY: &str = "MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEtk/jgNSuR3BaxQYmR1pkzmHPmPAXHoaU1HAN3q5wy0KP+Uxv8xNyFb5Y0SfREUmmo4Llc9XEPh1wa5IixzPL7g==";

const VERSION: u32 = 4;
const MAGIC_NUMBER: &[u8] = "1ezcrypt1".as_bytes();

#[derive(Deserialize, Serialize)]

pub struct SignedMessage {
    #[serde(with = "serde_bytes")]
    pub text: Vec<u8>,
    pub signature: EzcryptSignature,
}
/// Serializable signature
pub struct EzcryptSignature(ecdsa::DerSignature);
impl Serialize for EzcryptSignature {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_bytes(self.0.as_bytes())
    }
}
impl<'de> Deserialize<'de> for EzcryptSignature {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let deserialized = ecdsa::DerSignature::from_bytes(
            &deserializer.deserialize_bytes(PubVisitor)?,
        )
        .map_err(|e| serde::de::Error::custom(format!("error deserializing bytes: {:?}", e)))?;
        Ok(EzcryptSignature(deserialized))
    }
}
impl SignedMessage {
    pub fn verify(&self, public_key: p256::PublicKey) -> Result<(), ecdsa::Error> {
        let verifying_key = ecdsa::VerifyingKey::from(public_key);
        verifying_key.verify(&self.text, &self.signature.0)
    }
    fn new(secret_key: &p256::SecretKey, text: Vec<u8>) -> Self {
        let signing_key = SigningKey::from(secret_key);
        let signature = signing_key.sign(&text);
        Self {
            text,
            signature: EzcryptSignature(signature),
        }
    }
}

#[derive(Deserialize, Serialize)]
/// Encrypted file containing ciphertext, version information, and the backdoor
pub struct EncryptedFile {
    pub version: u32,
    pub ciphertext_key: EzcryptPublicKey,
    pub ciphertext: SignedMessage,
    #[serde(with = "serde_bytes")]
    pub nonce: Vec<u8>,
    #[serde(with = "serde_bytes")]
    pub salt: Vec<u8>,
    pub message: SignedMessage,
    pub backdoor: Backdoor,
    pub original_file_name: Option<String>,
}
#[derive(Error, Debug)]
pub enum EncryptError {
    #[error("error hashing data")]
    HashingError(argon2::Error),
    #[error("error generating ciphertext")]
    AesError(aes_gcm::Error),
    #[error("error parsing public key")]
    PublicKeyParseError,
}
#[derive(Error, Debug)]
#[error("signature is invalid")]
pub struct InvalidSignature;
impl EncryptedFile {
    /// Encrypt a file with plaintext and a password along with unencrypted text which can be used in the event of a forgotten password
    pub fn encrypt(
        plaintext: Vec<u8>,
        password: String,
        unencrypted_text: String,
        original_file_name: Option<String>,
    ) -> Result<Self, EncryptError> {
        let mut password_hash = [0u8; 32];
        let mut rng = thread_rng();
        let salt: Vec<u8> = (0..10).map(|_| rng.gen()).collect();
        argon2_with_our_defaults()
            .hash_password_into(password.as_bytes(), &salt, &mut password_hash)
            .map_err(EncryptError::HashingError)?;
        let aes_key = Key::<Aes256Gcm>::from_slice(&password_hash);
        let cipher = Aes256Gcm::new(aes_key);
        let nonce = Aes256Gcm::generate_nonce(&mut rng);
        let encrypted = cipher
            .encrypt(&nonce, Payload::from(plaintext.as_ref()))
            .map_err(EncryptError::AesError)?;
        let ciphertext_ecc = p256::SecretKey::random(&mut rng);
        let backdoor_public = p256::PublicKey::from_public_key_der(
            &base64engine
                .decode(BACKDOOR_PUBLIC_KEY)
                .map_err(|_| EncryptError::PublicKeyParseError)?,
        )
        .map_err(|_| EncryptError::PublicKeyParseError)?;
        let shared_secret = ecdh::diffie_hellman(
            ciphertext_ecc.to_nonzero_scalar(),
            backdoor_public.as_affine(),
        );
        let backdoor_salt: Vec<u8> = (0..10).map(|_| rng.gen()).collect();
        let mut shared_secret_hashed = [0u8; 32];
        argon2_with_our_defaults()
            .hash_password_into(
                shared_secret.raw_secret_bytes().as_ref(),
                &backdoor_salt,
                &mut shared_secret_hashed,
            )
            .map_err(EncryptError::HashingError)?;
        let backdoor_cipher = Aes256Gcm::new(Key::<Aes256Gcm>::from_slice(&shared_secret_hashed));
        let backdoor_nonce = Aes256Gcm::generate_nonce(&mut rng);
        let backdoor_encrypted_hash = backdoor_cipher
            .encrypt(&backdoor_nonce, Payload::from(password_hash.as_ref()))
            .map_err(EncryptError::AesError)?;
        let encrypted_file = EncryptedFile {
            version: VERSION,
            ciphertext_key: EzcryptPublicKey(ciphertext_ecc.public_key()),
            nonce: nonce.to_vec(),
            salt,
            ciphertext: SignedMessage::new(&ciphertext_ecc, encrypted),
            message: SignedMessage::new(&ciphertext_ecc, unencrypted_text.as_bytes().to_vec()),
            backdoor: Backdoor {
                backdoor_key: ByteBuf::from(backdoor_encrypted_hash),
                nonce: backdoor_nonce.to_vec(),
                salt: ByteBuf::from(backdoor_salt),
            },
            original_file_name,
        };
        Ok(encrypted_file)
    }
    /// Serialize a struct to an ezcrypt file, including the magic number
    pub fn serialize(self) -> Result<Vec<u8>, rmp_serde::encode::Error> {
        let mut serialized = rmp_serde::to_vec_named(&self)?;
        let mut with_magic = MAGIC_NUMBER.to_vec();
        with_magic.append(&mut serialized);
        Ok(with_magic)
    }
    /// Remove magic number and deserialize
    pub fn deserialize(contents: Vec<u8>) -> Result<Self, DeserializeError> {
        let contents_msgpack = contents
            .strip_prefix(MAGIC_NUMBER)
            .ok_or(DeserializeError::WrongPrefix)?;
        let parsed_contents: EncryptedFile =
            rmp_serde::from_slice(contents_msgpack).map_err(DeserializeError::MsgpackError)?;
        if parsed_contents.version != VERSION {
            return Err(DeserializeError::VersionError(parsed_contents.version));
        }
        Ok(parsed_contents)
    }
    /// Verify and retrieve message text from an ezcrypt file
    pub fn get_message(&self) -> Result<Option<Vec<u8>>, InvalidSignature> {
        if self.message.verify(self.ciphertext_key.0).is_err() {
            return Err(InvalidSignature);
        }
        if self.message.text.is_empty() {
            Ok(None)
        } else {
            Ok(Some(self.message.text.clone()))
        }
    }
    /// Decrypt file with password or rembercode
    pub fn decrypt(&self, password: String) -> Result<Vec<u8>, DecryptError> {
        if self.ciphertext.verify(self.ciphertext_key.0).is_err() {
            return Err(DecryptError::SignatureInvalid);
        }
        let mut key = [0u8; 32];
        argon2_with_our_defaults()
            .hash_password_into(password.as_bytes(), &self.salt, &mut key)
            .map_err(DecryptError::HashingError)?;
        let aes_key = Key::<Aes256Gcm>::from_slice(&key);
        let cipher = Aes256Gcm::new(aes_key);
        match cipher.decrypt((&*self.nonce).into(), self.ciphertext.text.as_ref()) {
            Ok(d) => Ok(d),
            Err(_) => {
                if let Some(code) = password.strip_prefix("irember-") {
                    let shared_secret = match base64engine.decode(code) {
                        Ok(code) => code,
                        Err(_) => {
                            return Err(DecryptError::RembercodeFormInvalid);
                        }
                    };
                    let mut key_key: [u8; 32] = [0u8; 32];
                    if let Err(e) = argon2_with_our_defaults().hash_password_into(
                        &shared_secret,
                        &self.backdoor.salt,
                        &mut key_key,
                    ) {
                        return Err(DecryptError::HashingError(e));
                    }

                    let aes_key = Key::<Aes256Gcm>::from_slice(&key_key);
                    let key_cipher = Aes256Gcm::new(aes_key);
                    let key = match key_cipher.decrypt(
                        (&*self.backdoor.nonce).into(),
                        &**self.backdoor.backdoor_key,
                    ) {
                        Ok(key) => key,
                        Err(_) => return Err(DecryptError::RembercodeBodyInvalid),
                    };
                    let aes_key = Key::<Aes256Gcm>::from_slice(&key);
                    let cipher = Aes256Gcm::new(aes_key);
                    let decrypted = match cipher
                        .decrypt((&*self.nonce).into(), self.ciphertext.text.as_ref())
                    {
                        Ok(decrypted) => decrypted,
                        Err(_) => return Err(DecryptError::RembercodeBodyInvalid),
                    };
                    Ok(decrypted)
                } else {
                    Err(DecryptError::WrongPassword)
                }
            }
        }
    }
    /// Generate a forgorcode with information contained within the encrypted file
    pub fn generate_forgorcode(self) -> Forgorcode {
        Forgorcode {
            public_key: self.ciphertext_key,
            message: self.message,
        }
    }
}
#[derive(Error, Debug)]
pub enum DeserializeError {
    #[error("file does not start with correct prefix")]
    WrongPrefix,
    #[error("msgpack error")]
    MsgpackError(#[from] rmp_serde::decode::Error),
    #[error("version must be {} but was {}",VERSION,(.0))]
    VersionError(u32),
}

#[derive(Error, Debug)]
pub enum DecryptError {
    #[error("signatures on ciphertext or plaintext were invalid")]
    SignatureInvalid,
    #[error("rembercode form is invalid")]
    RembercodeFormInvalid,
    #[error("rembercode secret does not decrypt data")]
    RembercodeBodyInvalid,
    #[error("password is wrong")]
    WrongPassword,
    #[error("error hashing data")]
    HashingError(argon2::Error),
}
/// Contains data needed to unlock your file if you forget the password
#[derive(Deserialize, Serialize)]
pub struct Backdoor {
    pub backdoor_key: ByteBuf,
    #[serde(with = "serde_bytes")]
    pub nonce: Vec<u8>,
    pub salt: ByteBuf,
}

/// Serializable public key
pub struct EzcryptPublicKey(pub p256::PublicKey);
impl Serialize for EzcryptPublicKey {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        let key_bytes = self.0.to_public_key_der().map_err(|e| {
            serde::ser::Error::custom(format!("error turning public key into der: {:?}", e))
        })?;
        serializer.serialize_bytes(key_bytes.as_bytes())
    }
}

struct PubVisitor;
impl<'a> Visitor<'a> for PubVisitor {
    type Value = Vec<u8>;

    fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(formatter, "a byte array")
    }
    fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
    where
        E: serde::de::Error,
    {
        Ok(v.to_vec())
    }
}
impl<'de> Deserialize<'de> for EzcryptPublicKey {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let key_bytes = deserializer.deserialize_bytes(PubVisitor)?;
        match p256::PublicKey::from_public_key_der(&key_bytes) {
            Ok(key) => Ok(EzcryptPublicKey(key)),
            Err(e) => Err(serde::de::Error::custom(format!(
                "error decoding der key: {:?}",
                e
            ))),
        }
    }
}
/// Code that contains a small amount of information used to reset your password
#[derive(Serialize, Deserialize)]
pub struct Forgorcode {
    pub public_key: EzcryptPublicKey,
    pub message: SignedMessage,
}
#[derive(Error, Debug)]
pub enum ForgorcodeDecodeError {
    #[error("code does not start with \"iforgor\"")]
    WrongPrefix,
    #[error("error decoding base64")]
    Base64Error(#[from] base64::DecodeError),
    #[error("error decoding msgpack")]
    MsgpackError(#[from] rmp_serde::decode::Error),
}

impl Forgorcode {
    pub fn encode(self) -> Result<String, rmp_serde::encode::Error> {
        let bytes = rmp_serde::encode::to_vec_named(&self)?;
        let encoded_base64 = base64engine.encode(bytes);
        Ok(format!("iforgor-{}", encoded_base64))
    }
    pub fn decode(code: String) -> Result<Self, ForgorcodeDecodeError> {
        let stripped = code
            .strip_prefix("iforgor-")
            .ok_or(ForgorcodeDecodeError::WrongPrefix)?;
        let decoded = base64engine.decode(stripped)?;
        let parsed: Forgorcode = rmp_serde::from_slice(&decoded)?;
        Ok(parsed)
    }
    pub fn generate_rembercode(self, secret: p256::SecretKey) -> String {
        let shared_secret =
            ecdh::diffie_hellman(secret.to_nonzero_scalar(), self.public_key.0.as_affine());
        let encoded =
            "irember-".to_string() + &base64engine.encode(shared_secret.raw_secret_bytes());
        encoded
    }
}
pub fn argon2_with_our_defaults() -> Argon2<'static> {
    Argon2::new(
        argon2::Algorithm::Argon2id,
        argon2::Version::V0x13,
        argon2::Params::new(64 * 1024, 1, 4, Some(32)).unwrap(),
    )
}