diff --git a/src/main.rs b/src/main.rs index 4999bf9..095e287 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,7 +19,8 @@ use std::{ io, }; use x_wing::{ - Ciphertext, DecapsulationKey, EncapsulationKey, DECAPSULATION_KEY_SIZE, ENCAPSULATION_KEY_SIZE, + Ciphertext, DecapsulationKey, EncapsulationKey, CIPHERTEXT_SIZE, DECAPSULATION_KEY_SIZE, + ENCAPSULATION_KEY_SIZE, }; const PLUGIN_NAME: &str = "xwing"; @@ -44,6 +45,23 @@ struct RecipientPlugin { recipients: Vec, } +impl RecipientPlugin { + fn wrap_file_key(&self, file_key: FileKey) -> Vec { + self.recipients + .iter() + .map(|recipient| { + let (ct, ss) = recipient.encapsulate(&mut OsRng).unwrap(); + let wrapped_key = aead_encrypt(&ss, file_key.expose_secret()); + Stanza { + tag: PLUGIN_NAME.to_string(), + args: vec![BASE64_STANDARD.encode(ct.as_bytes())], + body: wrapped_key, + } + }) + .collect() + } +} + impl RecipientPluginV1 for RecipientPlugin { fn add_recipient( &mut self, @@ -110,20 +128,7 @@ impl RecipientPluginV1 for RecipientPlugin { ) -> io::Result>, Vec>> { Ok(Ok(file_keys .into_iter() - .map(|file_key| { - self.recipients - .iter() - .map(|recipient| { - let (ct, ss) = recipient.encapsulate(&mut OsRng).unwrap(); - let wrapped_key = aead_encrypt(&ss, file_key.expose_secret()); - Stanza { - tag: PLUGIN_NAME.to_string(), - args: vec![BASE64_STANDARD.encode(ct.as_bytes())], - body: wrapped_key, - } - }) - .collect() - }) + .map(|file_key| self.wrap_file_key(file_key)) .collect())) } @@ -137,6 +142,107 @@ struct IdentityPlugin { identities: Vec, } +impl IdentityPlugin { + fn decrypt_stanzas( + &self, + file_index: usize, + stanzas: Vec<(usize, Stanza)>, + ) -> Result> { + let mut file_key = None; + let mut errors = Vec::new(); + + for (stanza_index, stanza) in stanzas { + let arg = match stanza.args.first() { + Some(arg) => arg, + None => { + errors.push(identity::Error::Stanza { + file_index, + stanza_index, + message: "Stanza is missing arguments".to_string(), + }); + continue; + } + }; + + let ct = match BASE64_STANDARD.decode(arg) { + Ok(ct) => ct, + Err(_) => { + errors.push(identity::Error::Stanza { + file_index, + stanza_index, + message: "Malformed base64".to_string(), + }); + continue; + } + }; + + let ct: [u8; CIPHERTEXT_SIZE] = match ct.try_into() { + Ok(ct) => ct, + Err(_) => { + errors.push(identity::Error::Stanza { + file_index, + stanza_index, + message: "Malformed ciphertext".to_string(), + }); + continue; + } + }; + + let ct = Ciphertext::from(&ct); + + let ss = match self + .identities + .iter() + .filter_map(|key| key.decapsulate(&ct).ok()) + .next() + { + Some(ss) => ss, + None => { + errors.push(identity::Error::Stanza { + file_index, + stanza_index, + message: "No identity found that can decrypt the file".to_string(), + }); + continue; + } + }; + + let unwrapped_file_key = match aead_decrypt(&ss, FILE_KEY_BYTES, &stanza.body) { + Ok(file_key) => FileKey::new(Box::new(file_key.try_into().unwrap_or_else(|_| { + panic!( + "aead_decrypt returned a plaintext with a different size as {}", + FILE_KEY_BYTES + ) + }))), + Err(e) => { + errors.push(identity::Error::Stanza { + file_index, + stanza_index, + message: e.to_string(), + }); + continue; + } + }; + + file_key = Some(unwrapped_file_key); + } + + if !errors.is_empty() { + return Err(errors); + } + + if let Some(file_key) = file_key { + return Ok(file_key); + } + + Err(vec![identity::Error::Stanza { + file_index, + stanza_index: 0, + message: "No stanzas found to be handeled by this plugin".to_string(), + }]) + } +} + impl IdentityPluginV1 for IdentityPlugin { fn add_identity( &mut self, @@ -153,7 +259,7 @@ impl IdentityPluginV1 for IdentityPlugin { let bytes: [u8; DECAPSULATION_KEY_SIZE] = match bytes.try_into() { Ok(x) => x, - _ => { + Err(_) => { return Err(identity::Error::Identity { index, message: "Invalid identity".to_string(), @@ -171,43 +277,20 @@ impl IdentityPluginV1 for IdentityPlugin { files: Vec>, mut _callbacks: impl Callbacks, ) -> io::Result>>> { - Ok(files - .into_iter() - .map(|file| try_decrypt_file_key(&self.identities, file)) - .enumerate() - .map(|(file_index, file_key)| { - file_key.ok_or(vec![identity::Error::Stanza { - file_index, - stanza_index: 1, - message: "Invalid stanzas".to_string(), - }]) - }) - .enumerate() - .collect()) - } -} + let mut ret = HashMap::default(); -fn try_decrypt_file_key(keys: &Vec, stanzas: Vec) -> Option { - stanzas - .iter() - .filter_map(|stanza| try_decrypt_stanza(keys, stanza)) - .next() -} + for (file_index, stanzas) in files.into_iter().enumerate() { + let x_wing_stanzas = stanzas + .into_iter() + .enumerate() + .filter(|(_, stanza)| stanza.tag == PLUGIN_NAME) + .collect(); -fn try_decrypt_stanza(keys: &Vec, stanza: &Stanza) -> Option { - let ct = BASE64_STANDARD.decode(stanza.args.first()?).ok()?; - let ct = Ciphertext::from(&ct.try_into().ok()?); - - for key in keys { - let ss = key.decapsulate(&ct).unwrap(); - let file_key = aead_decrypt(&ss, FILE_KEY_BYTES, &stanza.body); - - if let Ok(file_key) = file_key { - return Some(FileKey::new(Box::new(file_key.try_into().ok()?))); + ret.insert(file_index, self.decrypt_stanzas(file_index, x_wing_stanzas)); } - } - None + Ok(ret) + } } #[derive(Debug, Parser)] @@ -225,8 +308,7 @@ fn main() -> io::Result<()> { return Ok(()); } - // Here you can assume the binary is being run directly by a user, - // and perform administrative tasks like generating keys. + // Here you can assume the binary is being run directly by a user, and perform administrative tasks like generating keys. let (sk, pk) = x_wing::generate_key_pair_from_os_rng(); print_new_identity(PLUGIN_NAME, sk.as_bytes(), &pk.as_bytes());