age 0.12.1

//! The "encrypted age identity file" identity type.

use std::{cell::Cell, io};

use crate::{fl, scrypt, Callbacks, DecryptError, Decryptor, EncryptError, IdentityFile};

/// An encrypted age identity file.
///
/// This type can be explicitly decrypted to obtain an [`IdentityFile`]. If you want a
/// type that can be used directly as an identity and caches the decryption result
/// internally, use [`Identity`].
pub struct EncryptedIdentity<R: io::Read, C: Callbacks> {
    decryptor: Decryptor<R>,
    max_work_factor: Option<u8>,
    callbacks: C,
}

impl<R: io::Read, C: Callbacks> EncryptedIdentity<R, C> {
    /// Parses an encrypted identity from an input containing valid UTF-8.
    ///
    /// Returns `Ok(None)` if the input contains an age ciphertext that is not encrypted
    /// to a passphrase.
    pub(crate) fn from_buffer(
        data: R,
        callbacks: C,
        max_work_factor: Option<u8>,
    ) -> Result<Option<Self>, DecryptError> {
        let decryptor = Decryptor::new(data)?;
        Ok(Self::new(decryptor, callbacks, max_work_factor))
    }

    /// Constructs a new encrypted identity from a [`Decryptor`].
    ///
    /// Returns `Ok(None)` if the input contains an age ciphertext that is not encrypted
    /// to a passphrase.
    pub fn new(decryptor: Decryptor<R>, callbacks: C, max_work_factor: Option<u8>) -> Option<Self> {
        decryptor.is_scrypt().then_some(EncryptedIdentity {
            decryptor,
            max_work_factor,
            callbacks,
        })
    }

    /// Decrypts this encrypted identity.
    ///
    /// The provided filename (if any) will be used in the passphrase request message.
    pub fn decrypt(self, filename: Option<&str>) -> Result<IdentityFile<C>, DecryptError> {
        let passphrase = match self.callbacks.request_passphrase(&fl!(
            "encrypted-passphrase-prompt",
            filename = filename.unwrap_or_default()
        )) {
            Some(passphrase) => passphrase,
            None => Err(DecryptError::KeyDecryptionFailed)?,
        };

        let mut identity = scrypt::Identity::new(passphrase);
        if let Some(max_work_factor) = self.max_work_factor {
            identity.set_max_work_factor(max_work_factor);
        }

        self.decryptor
            .decrypt(Some(&identity as _).into_iter())
            .map_err(|e| {
                if matches!(e, DecryptError::DecryptionFailed) {
                    DecryptError::KeyDecryptionFailed
                } else {
                    e
                }
            })
            .and_then(|stream| {
                let file = IdentityFile::from_buffer(io::BufReader::new(stream))?
                    .with_callbacks(self.callbacks);
                Ok(file)
            })
    }
}

/// The state of the encrypted age identity.
enum IdentityState<R: io::Read, C: Callbacks> {
    Encrypted(EncryptedIdentity<R, C>),
    Decrypted(IdentityFile<C>),

    /// The file was not correctly encrypted, or did not contain age identities. We cache
    /// this error in case the caller tries to use this identity again. The `Option` is to
    /// enable implementing `Default` so we can use `Cell::take`, but we don't ever allow
    /// the `None` case to persist.
    Poisoned(Option<DecryptError>),
}

impl<R: io::Read, C: Callbacks> Default for IdentityState<R, C> {
    fn default() -> Self {
        Self::Poisoned(None)
    }
}

impl<R: io::Read, C: Callbacks> IdentityState<R, C> {
    /// Decrypts this encrypted identity if necessary.
    ///
    /// Returns the (possibly cached) identities, and a boolean marking if the identities
    /// were not cached (and we just asked the user for a passphrase).
    fn decrypt(self, filename: Option<&str>) -> Result<(IdentityFile<C>, bool), DecryptError> {
        match self {
            Self::Encrypted(encrypted) => encrypted.decrypt(filename).map(|file| (file, true)),
            Self::Decrypted(identity_file) => Ok((identity_file, false)),
            // `IdentityState::decrypt` is only ever called with `Some`.
            Self::Poisoned(e) => Err(e.unwrap()),
        }
    }
}

/// An encrypted age identity file.
///
/// This type can be used directly as an identity and caches the decryption result
/// internally. If you want a type that can be explicitly decrypted to obtain an
/// [`IdentityFile`], use [`EncryptedIdentity`].
pub struct Identity<R: io::Read, C: Callbacks> {
    state: Cell<IdentityState<R, C>>,
    filename: Option<String>,
}

impl<R: io::Read, C: Callbacks> Identity<R, C> {
    /// Parses an encrypted identity from an input containing valid UTF-8.
    ///
    /// `filename` is the path to the file that the input is reading from, if any.
    ///
    /// Returns `Ok(None)` if the input contains an age ciphertext that is not encrypted
    /// to a passphrase.
    pub fn from_buffer(
        data: R,
        filename: Option<String>,
        callbacks: C,
        max_work_factor: Option<u8>,
    ) -> Result<Option<Self>, DecryptError> {
        let encrypted = EncryptedIdentity::from_buffer(data, callbacks, max_work_factor)?;
        Ok(encrypted.map(|encrypted| Identity {
            state: Cell::new(IdentityState::Encrypted(encrypted)),
            filename,
        }))
    }

    /// Returns the recipients contained within this encrypted identity.
    ///
    /// If this encrypted identity has not been decrypted yet, calling this method will
    /// trigger a passphrase request.
    pub fn recipients(&self) -> Result<Vec<Box<dyn crate::Recipient + Send>>, EncryptError> {
        match self.state.take().decrypt(self.filename.as_deref()) {
            Ok((identity_file, _)) => {
                let recipients = identity_file.to_recipients();
                self.state.set(IdentityState::Decrypted(identity_file));
                recipients
            }
            Err(e) => {
                self.state.set(IdentityState::Poisoned(Some(e.clone())));
                Err(EncryptError::EncryptedIdentities(e))
            }
        }
    }

    /// Attempts to unwrap stanzas with the identities contained within this encrypted
    /// identity.
    ///
    /// We don't want to ask the user for the passphrase on every stanza decryption, and
    /// we don't want to store the entire encrypted age identity file in memory. Instead,
    /// the first time that an encrypted identity is decrypted with, we ask the caller for
    /// the passphrase, and perform validity checks on the decrypted data. We then cache
    /// the decrypted identities for subsequent calls.
    ///
    /// Because the `age::Identity` trait requires immutable references, this means that
    /// we need to use interior mutability here.
    fn unwrap_stanzas_base<F>(
        &self,
        filter: F,
    ) -> Option<Result<age_core::format::FileKey, DecryptError>>
    where
        F: Fn(
            Result<Box<dyn crate::Identity + Send + Sync>, DecryptError>,
        ) -> Option<Result<age_core::format::FileKey, DecryptError>>,
    {
        match self.state.take().decrypt(self.filename.as_deref()) {
            Ok((identity_file, requested_passphrase)) => {
                let result = identity_file.to_identities().find_map(filter);

                // If we requested a passphrase to decrypt, and none of the identities
                // matched, warn the user.
                if requested_passphrase && result.is_none() {
                    identity_file.callbacks.display_message(&fl!(
                        "encrypted-warn-no-match",
                        filename = self.filename.as_deref().unwrap_or_default()
                    ));
                }

                self.state.set(IdentityState::Decrypted(identity_file));
                result
            }
            Err(e) => {
                self.state.set(IdentityState::Poisoned(Some(e.clone())));
                Some(Err(e))
            }
        }
    }
}

impl<R: io::Read, C: Callbacks> crate::Identity for Identity<R, C> {
    fn unwrap_stanza(
        &self,
        stanza: &age_core::format::Stanza,
    ) -> Option<Result<age_core::format::FileKey, DecryptError>> {
        self.unwrap_stanzas_base(|identity| match identity {
            Ok(i) => i.unwrap_stanza(stanza),
            Err(e) => Some(Err(e)),
        })
    }

    fn unwrap_stanzas(
        &self,
        stanzas: &[age_core::format::Stanza],
    ) -> Option<Result<age_core::format::FileKey, DecryptError>> {
        self.unwrap_stanzas_base(|identity| match identity {
            Ok(i) => i.unwrap_stanzas(stanzas),
            Err(e) => Some(Err(e)),
        })
    }
}

#[cfg(test)]
mod tests {
    use std::sync::{Arc, Mutex};

    use age_core::secrecy::{ExposeSecret, SecretString};

    use super::Identity;
    use crate::{x25519, Callbacks, DecryptError, Identity as _, Recipient as _};

    #[cfg(feature = "armor")]
    use crate::armor::ArmoredReader;

    const TEST_ENCRYPTED_IDENTITY_PASSPHRASE: &str = "foobar";

    const TEST_ENCRYPTED_IDENTITY: &str = "-----BEGIN AGE ENCRYPTED FILE-----
YWdlLWVuY3J5cHRpb24ub3JnL3YxCi0+IHNjcnlwdCBza2I4R0t6L2NLT2s4cGlI
TTRGRjFRIDEwCnVodTdORmZjcCtjRmdnYU54bm8rZEJ5NWlrVHZLY1hyRzZEN2JE
ZVpwWnMKLS0tIEZTcDlSL3oyRC9NQ3JZa3ZvUzNaNlk4bnhBSUdJRTFrMmE4QzMr
UVNETlkK34fdtpwZz+qQaGuirGHEdodVe4JvnSG3ANQpWhkDcsRzoe/+OuHXNdnv
zhBhaKdthstzGXbd2yJbLrTH1A3YbWO+/3zTIZENzKU9XbibLLQ4M/TXwKMzoObY
oiMf5/+8GiQVREtHmm24wsc/479cVwnGVTdH7DL+wANmyf6S9Vc2FYQmXjLDxsJ0
LMF6Cpgcg09C2gg4pcb4TFUWmDuxnZrfggrptOtyzC8O8aRuKPZqCGnzoWNOWl86
fOrxrKTj7xCdNS3+OrCdnBC8Z9cKDxjCGWW3fkjLsYha0Jo=
-----END AGE ENCRYPTED FILE-----
";

    const TEST_RECIPIENT: &str = "age1ysxuaeqlk7xd8uqsh8lsnfwt9jzzjlqf49ruhpjrrj5yatlcuf7qke4pqe";

    #[derive(Clone)]
    struct MockCallbacks(Arc<Mutex<Option<Option<&'static str>>>>);

    impl MockCallbacks {
        fn new(passphrase: Option<&'static str>) -> Self {
            MockCallbacks(Arc::new(Mutex::new(Some(passphrase))))
        }
    }

    impl Callbacks for MockCallbacks {
        fn display_message(&self, _: &str) {
            unimplemented!()
        }

        fn confirm(&self, _: &str, _: &str, _: Option<&str>) -> Option<bool> {
            unimplemented!()
        }

        fn request_public_string(&self, _: &str) -> Option<String> {
            unimplemented!()
        }

        /// This intentionally panics if called twice.
        fn request_passphrase(&self, _: &str) -> Option<SecretString> {
            self.0
                .lock()
                .unwrap()
                .take()
                .expect("passphrase is only input once")
                .to_owned()
                .map(SecretString::from)
        }
    }

    #[test]
    #[cfg(feature = "armor")]
    fn round_trip() {
        use age_core::format::FileKey;

        let pk: x25519::Recipient = TEST_RECIPIENT.parse().unwrap();
        let file_key = FileKey::new(Box::new([12; 16]));
        let (wrapped, labels) = pk.wrap_file_key(&file_key).unwrap();
        assert!(labels.is_empty());

        // Unwrapping with the wrong passphrase fails.
        {
            let buf = ArmoredReader::new(TEST_ENCRYPTED_IDENTITY.as_bytes());
            let identity = Identity::from_buffer(
                buf,
                None,
                MockCallbacks::new(Some("wrong passphrase")),
                None,
            )
            .unwrap()
            .unwrap();

            if let Err(e) = identity.unwrap_stanzas(&wrapped).unwrap() {
                assert!(matches!(e, DecryptError::KeyDecryptionFailed));
            } else {
                panic!("Should have failed");
            }
        }

        // Unwrapping fails if we cannot obtain a passphrase.
        {
            let buf = ArmoredReader::new(TEST_ENCRYPTED_IDENTITY.as_bytes());
            let identity = Identity::from_buffer(buf, None, MockCallbacks::new(None), None)
                .unwrap()
                .unwrap();

            if let Err(e) = identity.unwrap_stanzas(&wrapped).unwrap() {
                assert!(matches!(e, DecryptError::KeyDecryptionFailed));
            } else {
                panic!("Should have failed");
            }
        }

        let buf = ArmoredReader::new(TEST_ENCRYPTED_IDENTITY.as_bytes());
        let identity = Identity::from_buffer(
            buf,
            None,
            MockCallbacks::new(Some(TEST_ENCRYPTED_IDENTITY_PASSPHRASE)),
            None,
        )
        .unwrap()
        .unwrap();
        let unwrapped = identity.unwrap_stanzas(&wrapped);
        assert_eq!(
            unwrapped.unwrap().unwrap().expose_secret(),
            file_key.expose_secret()
        );

        // Unwrapping a second time doesn't re-decrypt.
        identity.unwrap_stanzas(&wrapped);
    }
}