anubis-age 1.4.0

use std::fs::File;
use std::io;

use crate::{
    pqc::mlkem, Callbacks, DecryptError, EncryptError, IdentityFileConvertError, NoCallbacks,
};

#[cfg(feature = "cli-common")]
use crate::cli_common::file_io::InputReader;

/// The supported kinds of identities within an [`IdentityFile`].
#[derive(Clone)]
enum IdentityFileEntry {
    /// The ML-KEM-1024 identity type.
    MlKem1024(mlkem::Identity),
}

impl IdentityFileEntry {
    pub(crate) fn into_identity(
        self,
        callbacks: impl Callbacks,
    ) -> Result<Box<dyn crate::Identity + Send + Sync>, DecryptError> {
        match self {
            IdentityFileEntry::MlKem1024(i) => Ok(Box::new(i)),
        }
    }
}

/// A list of identities that has been parsed from some input file.
pub struct IdentityFile<C: Callbacks> {
    filename: Option<String>,
    identities: Vec<IdentityFileEntry>,
    pub(crate) callbacks: C,
}

impl IdentityFile<NoCallbacks> {
    /// Parses one or more identities from a file containing valid UTF-8.
    pub fn from_file(filename: String) -> io::Result<Self> {
        File::open(&filename)
            .map(io::BufReader::new)
            .and_then(|data| IdentityFile::parse_identities(Some(filename), data))
    }

    /// Parses one or more identities from a buffered input containing valid UTF-8.
    pub fn from_buffer<R: io::BufRead>(data: R) -> io::Result<Self> {
        Self::parse_identities(None, data)
    }

    /// Parses one or more identities from an [`InputReader`];
    #[cfg(feature = "cli-common")]
    pub fn from_input_reader(reader: InputReader) -> io::Result<Self> {
        let filename = reader.filename().map(String::from);
        Self::parse_identities(filename, io::BufReader::new(reader))
    }

    fn parse_identities<R: io::BufRead>(filename: Option<String>, data: R) -> io::Result<Self> {
        let mut identities = vec![];

        for (line_number, line) in data.lines().enumerate() {
            let line = line?;
            if line.is_empty() || line.starts_with('#') {
                continue;
            }

            if let Ok(identity) = line.parse::<mlkem::Identity>() {
                identities.push(IdentityFileEntry::MlKem1024(identity));
            } else {
                // Return a line number in place of the line, so we don't leak the file
                // contents in error messages.
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    if let Some(filename) = filename {
                        format!(
                            "identity file {} contains non-identity data on line {}",
                            filename,
                            line_number + 1
                        )
                    } else {
                        format!(
                            "identity file contains non-identity data on line {}",
                            line_number + 1
                        )
                    },
                ));
            }
        }

        Ok(IdentityFile {
            filename,
            identities,
            callbacks: NoCallbacks,
        })
    }
}

impl<C: Callbacks> IdentityFile<C> {
    /// Returns the file name this identity list originated from, if any.
    pub(crate) fn filename(&self) -> Option<&str> {
        self.filename.as_deref()
    }

    /// Sets the provided callbacks on this identity file, so that if this is an encrypted
    /// identity, it can potentially be decrypted.
    pub fn with_callbacks<D: Callbacks>(self, callbacks: D) -> IdentityFile<D> {
        IdentityFile {
            filename: self.filename,
            identities: self.identities,
            callbacks,
        }
    }

    /// Writes a recipients file containing the recipients corresponding to the identities
    /// in this file.
    ///
    /// Returns an error if this file is empty, or if it contains plugin identities (which
    /// can only be converted by the plugin binary itself).
    pub fn write_recipients_file<W: io::Write>(
        &self,
        mut output: W,
    ) -> Result<(), IdentityFileConvertError> {
        if self.identities.is_empty() {
            return Err(IdentityFileConvertError::NoIdentities {
                filename: self.filename.clone(),
            });
        }

        for identity in &self.identities {
            if let IdentityFileEntry::MlKem1024(sk) = identity {
                writeln!(output, "{}", sk.to_public())
                    .map_err(IdentityFileConvertError::FailedToWriteOutput)?;
            }
        }

        Ok(())
    }

    /// Returns recipients for the identities in this file.
    ///
    /// Plugin identities will be merged into one [`Recipient`] per unique plugin.
    ///
    /// [`Recipient`]: crate::Recipient
    pub fn to_recipients(&self) -> Result<Vec<Box<dyn crate::Recipient + Send>>, EncryptError> {
        let mut recipients = RecipientsAccumulator::new();
        recipients.with_identities_ref(self);
        recipients.build()
    }

    /// Returns the identities in this file.
    pub(crate) fn to_identities(
        &self,
    ) -> impl Iterator<Item = Result<Box<dyn crate::Identity + Send + Sync>, DecryptError>> + '_
    {
        self.identities
            .iter()
            .map(|entry| entry.clone().into_identity(self.callbacks.clone()))
    }

    /// Returns the identities in this file.
    pub fn into_identities(
        self,
    ) -> Result<Vec<Box<dyn crate::Identity + Send + Sync>>, DecryptError> {
        self.identities
            .into_iter()
            .map(|entry| entry.into_identity(self.callbacks.clone()))
            .collect()
    }
}

pub(crate) struct RecipientsAccumulator {
    recipients: Vec<Box<dyn crate::Recipient + Send>>,
}

impl RecipientsAccumulator {
    pub(crate) fn new() -> Self {
        Self { recipients: vec![] }
    }

    #[cfg(feature = "cli-common")]
    pub(crate) fn push(&mut self, recipient: Box<dyn crate::Recipient + Send>) {
        self.recipients.push(recipient);
    }

    #[cfg(feature = "cli-common")]
    pub(crate) fn with_identities<C: Callbacks>(&mut self, identity_file: IdentityFile<C>) {
        for entry in identity_file.identities {
            if let IdentityFileEntry::MlKem1024(i) = entry {
                self.recipients.push(Box::new(i.to_public()));
            }
        }
    }

    pub(crate) fn with_identities_ref<C: Callbacks>(&mut self, identity_file: &IdentityFile<C>) {
        for entry in &identity_file.identities {
            if let IdentityFileEntry::MlKem1024(i) = entry {
                self.recipients.push(Box::new(i.to_public()));
            }
        }
    }

    pub(crate) fn build(self) -> Result<Vec<Box<dyn crate::Recipient + Send>>, EncryptError> {
        Ok(self.recipients)
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use anubis_core::secrecy::ExposeSecret;
    use std::io::BufReader;

    use super::{IdentityFile, IdentityFileEntry};

    #[test]
    fn mlkem_identity_roundtrip() {
        let identity = crate::pqc::mlkem::Identity::generate();
        let encoded = identity.to_string();
        let file = IdentityFile::from_buffer(BufReader::new(encoded.expose_secret().as_bytes()))
            .expect("parse mlkem identity");
        assert!(matches!(
            file.identities[0],
            IdentityFileEntry::MlKem1024(_)
        ));
    }
}