anubis-rage 1.4.0

#![forbid(unsafe_code)]

use anubis_age::{
    armor::{ArmoredReader, ArmoredWriter, Format},
    cli_common::{file_io, read_identities, read_recipients, StdinGuard, UiCallbacks},
    pqc::mldsa::SigningKey,
    pqc::signed,
    secrecy::ExposeSecret,
    Identity,
};
use clap::{CommandFactory, Parser};
use i18n_embed::DesktopLanguageRequester;

use std::fs;
use std::io::{self, Read, Write};
use std::path::Path;

use log::warn;
mod cli;
use cli::AgeOptions;

mod error;

mod i18n;

#[macro_export]
macro_rules! fl {
    ($message_id:literal) => {{
        i18n_embed_fl::fl!($crate::i18n::LANGUAGE_LOADER, $message_id)
    }};

    ($message_id:literal, $($args:expr),* $(,)?) => {{
        i18n_embed_fl::fl!($crate::i18n::LANGUAGE_LOADER, $message_id, $($args), *)
    }};
}

macro_rules! warning {
    ($warning_id:literal) => {{
        eprintln!("{}", fl!("warning-msg", warning = fl!($warning_id)));
    }};
}

fn set_up_io(
    input: Option<String>,
    output: Option<String>,
    output_format: file_io::OutputFormat,
) -> io::Result<(file_io::InputReader, file_io::OutputWriter)> {
    let input = file_io::InputReader::new(input)?;

    // Create an output to the user-requested location.
    let output =
        file_io::OutputWriter::new(output, true, output_format, 0o666, input.is_terminal())?;

    Ok((input, output))
}

type ReadCheckerMatchCase = (&'static [u8], Box<dyn FnOnce() -> io::Result<()>>);
type ReadCheckerMatcher = Option<(&'static [u8], usize, Box<dyn FnOnce() -> io::Result<()>>)>;

/// A wrapper around a reader that checks it for various prefixes.
struct ReadChecker<R: io::Read, const N: usize> {
    inner: R,
    matches: [ReadCheckerMatcher; N],
}

impl<R: io::Read, const N: usize> ReadChecker<R, N> {
    fn new(inner: R, matches: [ReadCheckerMatchCase; N]) -> Self {
        Self {
            inner,
            matches: matches.map(|(prefix, on_match)| Some((prefix, 0, on_match))),
        }
    }
}

impl<R: io::Read, const N: usize> io::Read for ReadChecker<R, N> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let read = self.inner.read(buf)?;
        let data = &buf[..read];

        for matcher in &mut self.matches {
            if let Some((prefix, start, on_match)) = matcher.take() {
                let to_check = &prefix[start..];
                if to_check.len() > data.len() {
                    // We haven't read enough data to verify a full match; check for a
                    // partial match, and update the matched counter so we keep checking.
                    if to_check.starts_with(data) {
                        *matcher = Some((prefix, start + data.len(), on_match));
                    }
                } else if data.starts_with(to_check) {
                    on_match()?;
                    // Don't set matched so we stop checking.
                }
            }
        }

        Ok(read)
    }
}

fn encrypt(opts: AgeOptions) -> Result<(), error::EncryptError> {
    if opts.plugin_name.is_some() {
        return Err(error::EncryptError::PluginNameFlag);
    }

    if opts.passphrase {
        return Err(error::EncryptError::PassphraseDisabled);
    }

    // Require signing key
    let signing_key_path = opts
        .signing_key
        .ok_or(error::EncryptError::MissingSigningKey)?;

    // Load signing key
    let key_contents = fs::read_to_string(&signing_key_path)?;
    let signing_key: SigningKey = key_contents
        .lines()
        .find(|line| line.starts_with("ANUBIS-MLDSA-87-SECRET"))
        .ok_or_else(|| {
            error::EncryptError::SigningError("no signing key found in file".to_string())
        })?
        .parse()
        .map_err(|e| error::EncryptError::SigningError(format!("invalid signing key: {}", e)))?;

    let (format, output_format) = if opts.armor {
        (Format::AsciiArmor, file_io::OutputFormat::Text)
    } else {
        (Format::Binary, file_io::OutputFormat::Binary)
    };

    let (input, output) = set_up_io(opts.input.clone(), opts.output.clone(), output_format)?;

    let is_stdin = match input {
        file_io::InputReader::File(_) => false,
        file_io::InputReader::Stdin(_) => true,
    };
    let mut stdin_guard = StdinGuard::new(is_stdin);

    let is_stdout = match output {
        file_io::OutputWriter::File(..) => false,
        file_io::OutputWriter::Stdout(..) => true,
    };

    if opts.recipient.is_empty() && opts.recipients_file.is_empty() && opts.identity.is_empty() {
        return Err(error::EncryptError::Age(
            anubis_age::EncryptError::MissingRecipients,
        ));
    }

    let recipients = read_recipients(
        opts.recipient,
        opts.recipients_file,
        opts.identity,
        opts.max_work_factor,
        &mut stdin_guard,
    )?;

    let encryptor =
        anubis_age::Encryptor::with_recipients(recipients.iter().map(|r| r.as_ref() as _))?;

    // Encrypt to temporary buffer
    let mut encrypted_buffer = Vec::new();
    {
        let mut enc_output =
            encryptor.wrap_output(ArmoredWriter::wrap_output(&mut encrypted_buffer, format)?)?;

        const AGE_MAGIC: &[u8] = b"anubis-encryption.org/";
        const ARMORED_BEGIN_MARKER: &[u8] = b"-----BEGIN AGE ENCRYPTED FILE-----";
        let warn_double_encrypting = Box::new(|| {
            warning!("warn-double-encrypting");
            Ok(())
        });

        io::copy(
            &mut ReadChecker::new(
                input,
                [
                    (AGE_MAGIC, warn_double_encrypting.clone()),
                    (ARMORED_BEGIN_MARKER, warn_double_encrypting),
                ],
            ),
            &mut enc_output,
        )?;
        enc_output.finish().and_then(|armor| armor.finish())?;
    }

    // Sign the encrypted data
    let mut signed_output = Vec::new();
    signed::sign_encrypted_file(&encrypted_buffer, &signing_key, &mut signed_output)
        .map_err(|e| error::EncryptError::SigningError(e.to_string()))?;

    // Write signed output to final destination
    let map_io_errors = |e: io::Error| match e.kind() {
        io::ErrorKind::BrokenPipe => error::EncryptError::BrokenPipe {
            is_stdout,
            source: e,
        },
        _ => e.into(),
    };

    match output {
        file_io::OutputWriter::File(mut f) => {
            f.write(&signed_output).map_err(map_io_errors)?;
        }
        file_io::OutputWriter::Stdout(mut s) => {
            s.write(&signed_output).map_err(map_io_errors)?;
        }
    }

    eprintln!("File encrypted and signed with ML-DSA-87");
    Ok(())
}

fn write_output<R: io::Read, W: io::Write>(
    mut input: R,
    mut output: W,
) -> Result<(), error::DecryptError> {
    io::copy(&mut input, &mut output)?;

    Ok(())
}

fn decrypt(opts: AgeOptions) -> Result<(), error::DecryptError> {
    if opts.armor {
        return Err(error::DecryptError::ArmorFlag);
    }
    if opts.passphrase {
        return Err(error::DecryptError::PassphraseFlag);
    }

    if !opts.recipient.is_empty() {
        return Err(error::DecryptError::RecipientFlag);
    }
    if !opts.recipients_file.is_empty() {
        return Err(error::DecryptError::RecipientsFileFlag);
    }

    if !(opts.identity.is_empty() || opts.plugin_name.is_none()) {
        return Err(error::DecryptError::MixedIdentityAndPluginName);
    }

    // Require verification key
    let verify_key_str = opts
        .verify_key
        .ok_or(error::DecryptError::MissingVerifyKey)?;

    // Parse verification key
    let verification_key = if let Ok(vk) = verify_key_str.parse() {
        vk
    } else {
        // Try to read as file
        let key_contents = fs::read_to_string(&verify_key_str)?;
        key_contents
            .lines()
            .find(|line| line.starts_with("anubis1mldsa87"))
            .ok_or_else(|| {
                error::DecryptError::VerificationError("no verification key found".to_string())
            })?
            .parse()
            .map_err(|e| {
                error::DecryptError::VerificationError(format!("invalid verification key: {}", e))
            })?
    };

    let (mut input, output) = set_up_io(opts.input, opts.output, file_io::OutputFormat::Unknown)?;

    let is_stdin = match input {
        file_io::InputReader::File(_) => false,
        file_io::InputReader::Stdin(_) => true,
    };
    let mut stdin_guard = StdinGuard::new(is_stdin);

    let stdin_identity = opts.identity.iter().any(|s| s == "-");
    if opts.plugin_name.is_some() {
        return Err(error::DecryptError::PluginNameFlag);
    }
    let identities = read_identities(opts.identity, opts.max_work_factor, &mut stdin_guard)?;

    // Read signed file into buffer
    let mut signed_data = Vec::new();
    input.read_to_end(&mut signed_data)?;

    // Verify signature and extract encrypted data
    let encrypted_data = signed::verify_and_extract(&signed_data, &verification_key)
        .map_err(|e| error::DecryptError::VerificationError(e.to_string()))?;

    eprintln!("Signature verified successfully");

    // CRLF_MANGLED_INTRO and UTF16_MANGLED_INTRO are the intro lines of the age format after
    // mangling by various versions of PowerShell redirection, truncated to the length of the
    // correct intro line. See https://github.com/FiloSottile/age/issues/290 for more info.
    const CRLF_MANGLED_INTRO: &[u8] = b"anubis-encryption.org/v1\r";
    const UTF16_MANGLED_INTRO: &[u8] =
        b"\xff\xfea\x00g\x00e\x00-\x00e\x00n\x00c\x00r\x00y\x00p\x00";
    let err_powershell_corruption = Box::new(|| {
        Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            error::DetectedPowerShellCorruptionError,
        ))
    });

    let encrypted_reader = ReadChecker::new(
        &encrypted_data[..],
        [
            (CRLF_MANGLED_INTRO, err_powershell_corruption.clone()),
            (UTF16_MANGLED_INTRO, err_powershell_corruption),
        ],
    );

    let decryptor = anubis_age::Decryptor::new_buffered(ArmoredReader::new(encrypted_reader))?;

    if identities.is_empty() {
        return Err(error::DecryptError::MissingIdentities { stdin_identity });
    }

    decryptor
        .decrypt(identities.iter().map(|i| i.as_ref() as &dyn Identity))
        .map_err(|e| e.into())
        .and_then(|input| write_output(input, output))
}

fn main() -> Result<(), error::Error> {
    use std::env::args;

    env_logger::builder()
        .format_timestamp(None)
        .filter_level(log::LevelFilter::Off)
        .parse_default_env()
        .init();

    let requested_languages = DesktopLanguageRequester::requested_languages();
    i18n::load_languages(&requested_languages);
    anubis_age::localizer()
        .select(&requested_languages)
        .unwrap();

    // If you are piping input with no other args, this will not allow
    // it.
    if console::user_attended() && args().len() == 1 {
        AgeOptions::command()
            .print_help()
            .map_err(error::EncryptError::Io)?;
        return Ok(());
    }

    let opts = AgeOptions::parse();

    if opts.encrypt && opts.decrypt {
        return Err(error::Error::MixedEncryptAndDecrypt);
    }
    if !(opts.identity.is_empty() || opts.encrypt || opts.decrypt) {
        return Err(error::Error::IdentityFlagAmbiguous);
    }

    if let (Some(in_file), Some(out_file)) = (&opts.input, &opts.output) {
        // Check that the given filenames do not correspond to the same file.
        let in_path = Path::new(&in_file);
        let out_path = Path::new(&out_file);
        match (in_path.canonicalize(), out_path.canonicalize()) {
            (Ok(in_abs), Ok(out_abs)) if in_abs == out_abs => {
                return Err(error::Error::SameInputAndOutput(out_file.clone()));
            }
            _ => (),
        }
    }

    if opts.decrypt {
        decrypt(opts).map_err(error::Error::from)
    } else {
        encrypt(opts).map_err(error::Error::from)
    }
}