keyfork/keyfork-shard/src/bin/keyfork-shard-decrypt-openpgp.rs

use std::{
    env,
    fs::File,
    io::{stdin, stdout},
    path::{Path, PathBuf},
    process::ExitCode,
    str::FromStr,
};

use aes_gcm::{aead::Aead, aes::cipher::consts::U12, Aes256Gcm, KeyInit, Nonce};
use x25519_dalek::{EphemeralSecret, PublicKey};

use keyfork_mnemonic_util::{Mnemonic, Wordlist};
use keyfork_prompt::PromptManager;
use keyfork_shard::openpgp::{decrypt_one, discover_certs, openpgp::Cert, parse_messages};

type Result<T, E = Box<dyn std::error::Error>> = std::result::Result<T, E>;

fn validate<'a>(
    messages_file: impl AsRef<Path>,
    key_discovery: impl Into<Option<&'a str>>,
) -> Result<(File, Vec<Cert>)> {
    let key_discovery = key_discovery.into().map(PathBuf::from);
    key_discovery.as_ref().map(std::fs::metadata).transpose()?;

    // Load certs from path
    let certs = key_discovery
        .map(discover_certs)
        .transpose()?
        .unwrap_or(vec![]);

    Ok((File::open(messages_file)?, certs))
}

fn run() -> Result<()> {
    let mut args = env::args();
    let program_name = args.next().expect("program name");
    let args = args.collect::<Vec<_>>();
    let (messages_file, cert_list) = match args.as_slice() {
        [messages_file, key_discovery] => validate(messages_file, key_discovery.as_str())?,
        [messages_file] => validate(messages_file, None)?,
        _ => panic!("Usage: {program_name} messages_file [key_discovery]"),
    };

    let mut encrypted_messages = parse_messages(messages_file)?;

    let encrypted_metadata = encrypted_messages
        .pop_front()
        .expect("any pgp encrypted message");

    // Receive their key

    let mut pm = PromptManager::new(stdin(), stdout())?;
    let wordlist = Wordlist::default();
    let their_words = pm.prompt_wordlist("Their words: ", &wordlist)?;
    let mut nonce_words = their_words
        .split_whitespace()
        .take(9)
        .peekable();
    let mut pubkey_words = their_words
        .split_whitespace()
        .skip(9)
        .take(24)
        .peekable();
    let mut nonce_mnemonic = String::new();
    let mut pubkey_mnemonic = String::new();
    while let Some(word) = nonce_words.next() {
        nonce_mnemonic.push_str(word);
        if nonce_words.peek().is_some() {
            nonce_mnemonic.push(' ');
        }
    }
    while let Some(word) = pubkey_words.next() {
        pubkey_mnemonic.push_str(word);
        if pubkey_words.peek().is_some() {
            pubkey_mnemonic.push(' ');
        }
    }
    let their_key = Mnemonic::from_str(&pubkey_mnemonic)?.entropy();
    let their_key: [u8; 32] = their_key.try_into().expect("24 words");
    let their_nonce = Mnemonic::from_str(&nonce_mnemonic)?.entropy();
    let their_nonce = Nonce::<U12>::from_slice(&their_nonce);

    let our_key = EphemeralSecret::random();
    let our_mnemonic =
        Mnemonic::from_entropy(PublicKey::from(&our_key).as_bytes(), Default::default())?;
    // TODO: Encode using `qrencode -t ansiutf8 -m 2`
    pm.prompt_message(&format!("Our words: {our_mnemonic}"))?;

    let shared_secret = our_key
        .diffie_hellman(&PublicKey::from(their_key))
        .to_bytes();

    let share = decrypt_one(encrypted_messages.into(), &cert_list, encrypted_metadata)?;
    assert!(share.len() <= 65, "invalid share length (too long)");
    const LEN: u8 = 24 * 3;
    let mut encrypted_payload = [(LEN - share.len() as u8); LEN as usize];
    encrypted_payload[..share.len()].copy_from_slice(&share);

    let shared_key = Aes256Gcm::new_from_slice(&shared_secret)?;
    let bytes = shared_key.encrypt(their_nonce, share.as_slice()).unwrap();

    // NOTE: Padding length is less than u8::MAX because 24 * 4 < u8::MAX
    const ENC_LEN: u8 = 24 * 4;
    let mut out_bytes = [(ENC_LEN - bytes.len() as u8); ENC_LEN as usize];
    assert!(bytes.len() < out_bytes.len(), "encrypted payload larger than acceptable limit");
    out_bytes[..bytes.len()].clone_from_slice(&bytes);

    // safety: size of out_bytes is constant and always % 4 == 0
    let mnemonic = unsafe { Mnemonic::from_raw_entropy(&out_bytes, Default::default()) };

    pm.prompt_message(&format!("Our payload: {mnemonic}"))?;

    Ok(())
}

fn main() -> ExitCode {
    let result = run();
    if let Err(e) = result {
        eprintln!("Error: {e}");
        let mut source = e.source();
        while let Some(new_error) = source.take() {
            eprintln!("Source: {new_error}");
            source = new_error.source();
        }
        return ExitCode::FAILURE;
    }
    ExitCode::SUCCESS
}