use std::{cell::Cell, io};
use crate::{
decryptor::PassphraseDecryptor, fl, Callbacks, DecryptError, Decryptor, EncryptError,
IdentityFile, IdentityFileEntry,
};
enum IdentityState<R: io::Read> {
Encrypted {
decryptor: PassphraseDecryptor<R>,
max_work_factor: Option<u8>,
},
Decrypted(Vec<IdentityFileEntry>),
Poisoned(Option<DecryptError>),
}
impl<R: io::Read> Default for IdentityState<R> {
fn default() -> Self {
Self::Poisoned(None)
}
}
impl<R: io::Read> IdentityState<R> {
fn decrypt<C: Callbacks>(
self,
filename: Option<&str>,
callbacks: C,
) -> Result<(Vec<IdentityFileEntry>, bool), DecryptError> {
match self {
Self::Encrypted {
decryptor,
max_work_factor,
} => {
let passphrase = match callbacks.request_passphrase(&fl!(
"encrypted-passphrase-prompt",
filename = filename.unwrap_or_default()
)) {
Some(passphrase) => passphrase,
None => todo!(),
};
decryptor
.decrypt(&passphrase, max_work_factor)
.map_err(|e| {
if matches!(e, DecryptError::DecryptionFailed) {
DecryptError::KeyDecryptionFailed
} else {
e
}
})
.and_then(|stream| {
let file = IdentityFile::from_buffer(io::BufReader::new(stream))?;
Ok((file.into_identities(), true))
})
}
Self::Decrypted(identities) => Ok((identities, false)),
Self::Poisoned(e) => Err(e.unwrap()),
}
}
}
pub struct Identity<R: io::Read, C: Callbacks> {
state: Cell<IdentityState<R>>,
filename: Option<String>,
callbacks: C,
}
impl<R: io::Read, C: Callbacks> Identity<R, C> {
pub fn from_buffer(
data: R,
filename: Option<String>,
callbacks: C,
max_work_factor: Option<u8>,
) -> Result<Option<Self>, DecryptError> {
match Decryptor::new(data)? {
Decryptor::Recipients(_) => Ok(None),
Decryptor::Passphrase(decryptor) => Ok(Some(Identity {
state: Cell::new(IdentityState::Encrypted {
decryptor,
max_work_factor,
}),
filename,
callbacks,
})),
}
}
pub fn recipients(&self) -> Result<Vec<Box<dyn crate::Recipient + Send>>, EncryptError> {
match self
.state
.take()
.decrypt(self.filename.as_deref(), self.callbacks.clone())
{
Ok((identities, _)) => {
let recipients = identities
.iter()
.map(|entry| entry.to_recipient(self.callbacks.clone()))
.collect::<Result<Vec<_>, _>>();
self.state.set(IdentityState::Decrypted(identities));
recipients
}
Err(e) => {
self.state.set(IdentityState::Poisoned(Some(e.clone())));
Err(EncryptError::EncryptedIdentities(e))
}
}
}
fn unwrap_stanzas_base<F>(
&self,
filter: F,
) -> Option<Result<age_core::format::FileKey, DecryptError>>
where
F: Fn(
Result<Box<dyn crate::Identity>, DecryptError>,
) -> Option<Result<age_core::format::FileKey, DecryptError>>,
{
match self
.state
.take()
.decrypt(self.filename.as_deref(), self.callbacks.clone())
{
Ok((identities, requested_passphrase)) => {
let result = identities
.iter()
.map(|entry| entry.clone().into_identity(self.callbacks.clone()))
.find_map(filter);
if requested_passphrase && result.is_none() {
self.callbacks.display_message(&fl!(
"encrypted-warn-no-match",
filename = self.filename.as_deref().unwrap_or_default()
));
}
self.state.set(IdentityState::Decrypted(identities));
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<&'static str>>>);
impl MockCallbacks {
fn new(passphrase: &'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!()
}
fn request_passphrase(&self, _: &str) -> Option<SecretString> {
Some(SecretString::new(
self.0.lock().unwrap().take().unwrap().to_owned(),
))
}
}
#[test]
#[cfg(feature = "armor")]
fn round_trip() {
let pk: x25519::Recipient = TEST_RECIPIENT.parse().unwrap();
let file_key = [12; 16].into();
let wrapped = pk.wrap_file_key(&file_key).unwrap();
{
let buf = ArmoredReader::new(TEST_ENCRYPTED_IDENTITY.as_bytes());
let identity =
Identity::from_buffer(buf, None, MockCallbacks::new("wrong passphrase"), 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(TEST_ENCRYPTED_IDENTITY_PASSPHRASE),
None,
)
.unwrap()
.unwrap();
let unwrapped = identity.unwrap_stanzas(&wrapped);
assert_eq!(
unwrapped.unwrap().unwrap().expose_secret(),
file_key.expose_secret()
);
identity.unwrap_stanzas(&wrapped);
}
}