anubis_age/
identity.rs

1use std::fs::File;
2use std::io;
3
4use crate::{
5    pqc::mlkem, Callbacks, DecryptError, EncryptError, IdentityFileConvertError, NoCallbacks,
6};
7
8#[cfg(feature = "cli-common")]
9use crate::cli_common::file_io::InputReader;
10
11/// The supported kinds of identities within an [`IdentityFile`].
12#[derive(Clone)]
13enum IdentityFileEntry {
14    /// The ML-KEM-1024 identity type.
15    MlKem1024(mlkem::Identity),
16}
17
18impl IdentityFileEntry {
19    pub(crate) fn into_identity(
20        self,
21        callbacks: impl Callbacks,
22    ) -> Result<Box<dyn crate::Identity + Send + Sync>, DecryptError> {
23        match self {
24            IdentityFileEntry::MlKem1024(i) => Ok(Box::new(i)),
25        }
26    }
27}
28
29/// A list of identities that has been parsed from some input file.
30pub struct IdentityFile<C: Callbacks> {
31    filename: Option<String>,
32    identities: Vec<IdentityFileEntry>,
33    pub(crate) callbacks: C,
34}
35
36impl IdentityFile<NoCallbacks> {
37    /// Parses one or more identities from a file containing valid UTF-8.
38    pub fn from_file(filename: String) -> io::Result<Self> {
39        File::open(&filename)
40            .map(io::BufReader::new)
41            .and_then(|data| IdentityFile::parse_identities(Some(filename), data))
42    }
43
44    /// Parses one or more identities from a buffered input containing valid UTF-8.
45    pub fn from_buffer<R: io::BufRead>(data: R) -> io::Result<Self> {
46        Self::parse_identities(None, data)
47    }
48
49    /// Parses one or more identities from an [`InputReader`];
50    #[cfg(feature = "cli-common")]
51    pub fn from_input_reader(reader: InputReader) -> io::Result<Self> {
52        let filename = reader.filename().map(String::from);
53        Self::parse_identities(filename, io::BufReader::new(reader))
54    }
55
56    fn parse_identities<R: io::BufRead>(filename: Option<String>, data: R) -> io::Result<Self> {
57        let mut identities = vec![];
58
59        for (line_number, line) in data.lines().enumerate() {
60            let line = line?;
61            if line.is_empty() || line.starts_with('#') {
62                continue;
63            }
64
65            if let Ok(identity) = line.parse::<mlkem::Identity>() {
66                identities.push(IdentityFileEntry::MlKem1024(identity));
67            } else {
68                // Return a line number in place of the line, so we don't leak the file
69                // contents in error messages.
70                return Err(io::Error::new(
71                    io::ErrorKind::InvalidData,
72                    if let Some(filename) = filename {
73                        format!(
74                            "identity file {} contains non-identity data on line {}",
75                            filename,
76                            line_number + 1
77                        )
78                    } else {
79                        format!(
80                            "identity file contains non-identity data on line {}",
81                            line_number + 1
82                        )
83                    },
84                ));
85            }
86        }
87
88        Ok(IdentityFile {
89            filename,
90            identities,
91            callbacks: NoCallbacks,
92        })
93    }
94}
95
96impl<C: Callbacks> IdentityFile<C> {
97    /// Returns the file name this identity list originated from, if any.
98    pub(crate) fn filename(&self) -> Option<&str> {
99        self.filename.as_deref()
100    }
101
102    /// Sets the provided callbacks on this identity file, so that if this is an encrypted
103    /// identity, it can potentially be decrypted.
104    pub fn with_callbacks<D: Callbacks>(self, callbacks: D) -> IdentityFile<D> {
105        IdentityFile {
106            filename: self.filename,
107            identities: self.identities,
108            callbacks,
109        }
110    }
111
112    /// Writes a recipients file containing the recipients corresponding to the identities
113    /// in this file.
114    ///
115    /// Returns an error if this file is empty, or if it contains plugin identities (which
116    /// can only be converted by the plugin binary itself).
117    pub fn write_recipients_file<W: io::Write>(
118        &self,
119        mut output: W,
120    ) -> Result<(), IdentityFileConvertError> {
121        if self.identities.is_empty() {
122            return Err(IdentityFileConvertError::NoIdentities {
123                filename: self.filename.clone(),
124            });
125        }
126
127        for identity in &self.identities {
128            if let IdentityFileEntry::MlKem1024(sk) = identity {
129                writeln!(output, "{}", sk.to_public())
130                    .map_err(IdentityFileConvertError::FailedToWriteOutput)?;
131            }
132        }
133
134        Ok(())
135    }
136
137    /// Returns recipients for the identities in this file.
138    ///
139    /// Plugin identities will be merged into one [`Recipient`] per unique plugin.
140    ///
141    /// [`Recipient`]: crate::Recipient
142    pub fn to_recipients(&self) -> Result<Vec<Box<dyn crate::Recipient + Send>>, EncryptError> {
143        let mut recipients = RecipientsAccumulator::new();
144        recipients.with_identities_ref(self);
145        recipients.build()
146    }
147
148    /// Returns the identities in this file.
149    pub(crate) fn to_identities(
150        &self,
151    ) -> impl Iterator<Item = Result<Box<dyn crate::Identity + Send + Sync>, DecryptError>> + '_
152    {
153        self.identities
154            .iter()
155            .map(|entry| entry.clone().into_identity(self.callbacks.clone()))
156    }
157
158    /// Returns the identities in this file.
159    pub fn into_identities(
160        self,
161    ) -> Result<Vec<Box<dyn crate::Identity + Send + Sync>>, DecryptError> {
162        self.identities
163            .into_iter()
164            .map(|entry| entry.into_identity(self.callbacks.clone()))
165            .collect()
166    }
167}
168
169pub(crate) struct RecipientsAccumulator {
170    recipients: Vec<Box<dyn crate::Recipient + Send>>,
171}
172
173impl RecipientsAccumulator {
174    pub(crate) fn new() -> Self {
175        Self { recipients: vec![] }
176    }
177
178    #[cfg(feature = "cli-common")]
179    pub(crate) fn push(&mut self, recipient: Box<dyn crate::Recipient + Send>) {
180        self.recipients.push(recipient);
181    }
182
183    #[cfg(feature = "cli-common")]
184    pub(crate) fn with_identities<C: Callbacks>(&mut self, identity_file: IdentityFile<C>) {
185        for entry in identity_file.identities {
186            if let IdentityFileEntry::MlKem1024(i) = entry {
187                self.recipients.push(Box::new(i.to_public()));
188            }
189        }
190    }
191
192    pub(crate) fn with_identities_ref<C: Callbacks>(&mut self, identity_file: &IdentityFile<C>) {
193        for entry in &identity_file.identities {
194            if let IdentityFileEntry::MlKem1024(i) = entry {
195                self.recipients.push(Box::new(i.to_public()));
196            }
197        }
198    }
199
200    pub(crate) fn build(self) -> Result<Vec<Box<dyn crate::Recipient + Send>>, EncryptError> {
201        Ok(self.recipients)
202    }
203}
204
205#[cfg(test)]
206pub(crate) mod tests {
207    use anubis_core::secrecy::ExposeSecret;
208    use std::io::BufReader;
209
210    use super::{IdentityFile, IdentityFileEntry};
211
212    #[test]
213    fn mlkem_identity_roundtrip() {
214        let identity = crate::pqc::mlkem::Identity::generate();
215        let encoded = identity.to_string();
216        let file = IdentityFile::from_buffer(BufReader::new(encoded.expose_secret().as_bytes()))
217            .expect("parse mlkem identity");
218        assert!(matches!(
219            file.identities[0],
220            IdentityFileEntry::MlKem1024(_)
221        ));
222    }
223}