keyfork/keyfork-shard/src/lib.rs

use std::{
    io::{stdin, stdout, Write},
    str::FromStr,
};

use aes_gcm::{
    aead::{Aead, AeadCore, OsRng},
    Aes256Gcm, KeyInit,
};
use hkdf::Hkdf;
use sha2::Sha256;
use keyfork_mnemonic_util::{Mnemonic, Wordlist};
use keyfork_prompt::{qrencode, Message as PromptMessage, PromptManager};
use sharks::{Share, Sharks};
use x25519_dalek::{EphemeralSecret, PublicKey};

#[cfg(feature = "openpgp")]
pub mod openpgp;

#[derive(thiserror::Error, Debug)]
pub enum SharksError {
    #[error("Error creating share: {0}")]
    Share(String),

    #[error("Error combining shares: {0}")]
    CombineShare(String),
}

#[derive(thiserror::Error, Debug)]
#[error("Mnemonic did not store enough data")]
pub struct InvalidMnemonicData;

/// Decrypt hunk version 1:
/// 1 byte: Version
/// 1 byte: Threshold
/// Data: &[u8]
pub(crate) const HUNK_VERSION: u8 = 1;
pub(crate) const HUNK_OFFSET: usize = 2;

/// # Panics
///
/// The function may panic if it is given payloads generated using a version of Keyfork that is
/// incompatible with the currently running version.
pub fn remote_decrypt(w: &mut impl Write) -> Result<(), Box<dyn std::error::Error>> {
    let mut pm = PromptManager::new(stdin(), stdout())?;
    let wordlist = Wordlist::default();

    let mut iter_count = None;
    let mut shares = vec![];

    let mut threshold = 0;

    while iter_count.is_none() || iter_count.is_some_and(|i| i > 0) {
        let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
        let nonce_mnemonic =
            unsafe { Mnemonic::from_raw_entropy(nonce.as_slice(), Default::default()) };
        let our_key = EphemeralSecret::random();
        let key_mnemonic =
            Mnemonic::from_entropy(PublicKey::from(&our_key).as_bytes(), Default::default())?;
        let combined_mnemonic = format!("{nonce_mnemonic} {key_mnemonic}");
        pm.prompt_message(&PromptMessage::Text(format!(
            "Our words: {combined_mnemonic}"
        )))?;

        if let Ok(qrcode) = qrencode::qrencode(&combined_mnemonic) {
            pm.prompt_message(&PromptMessage::Data(qrcode))?;
        }

        let their_words = pm.prompt_wordlist("Their words: ", &wordlist)?;

        let mut pubkey_words = their_words.split_whitespace().take(24).peekable();
        let mut payload_words = their_words.split_whitespace().skip(24).take(48).peekable();
        let mut pubkey_mnemonic = String::new();
        let mut payload_mnemonic = String::new();
        while let Some(word) = pubkey_words.next() {
            pubkey_mnemonic.push_str(word);
            if pubkey_words.peek().is_some() {
                pubkey_mnemonic.push(' ');
            }
        }
        while let Some(word) = payload_words.next() {
            payload_mnemonic.push_str(word);
            if payload_words.peek().is_some() {
                payload_mnemonic.push(' ');
            }
        }

        let their_key = Mnemonic::from_str(&pubkey_mnemonic)?.entropy();
        let their_key: [u8; 32] = their_key.try_into().map_err(|_| InvalidMnemonicData)?;

        let shared_secret = our_key
            .diffie_hellman(&PublicKey::from(their_key))
            .to_bytes();
        let hkdf = Hkdf::<Sha256>::new(None, &shared_secret);
        let mut hkdf_output = [0u8; 256 / 8];
        hkdf.expand(&[], &mut hkdf_output)?;
        let shared_key =
            Aes256Gcm::new_from_slice(&hkdf_output)?;

        let payload = Mnemonic::from_str(&payload_mnemonic)?.entropy();
        let payload =
            shared_key.decrypt(&nonce, &payload[..payload[payload.len() - 1] as usize])?;
        assert_eq!(HUNK_VERSION, payload[0], "Incompatible hunk version");

        match &mut iter_count {
            Some(n) => {
                // Must be > 0 to start loop, can't go lower
                *n -= 1;
            }
            None => {
                // NOTE: Should always be >= 1, < 256 due to Shamir constraints
                threshold = payload[1];
                let _ = iter_count.insert(threshold - 1);
            }
        }

        shares.push(payload[HUNK_OFFSET..].to_vec());
    }

    let shares = shares
        .into_iter()
        .map(|s| Share::try_from(s.as_slice()))
        .collect::<Result<Vec<_>, &str>>()
        .map_err(|e| SharksError::Share(e.to_string()))?;
    let secret = Sharks(threshold)
        .recover(&shares)
        .map_err(|e| SharksError::CombineShare(e.to_string()))?;

    /*
     * Verification would take up too much size, mnemonic would be very large
    let userid = UserID::from("keyfork-sss");
    let kdr = DerivationRequest::new(
        DerivationAlgorithm::Ed25519,
        &DerivationPath::from_str("m/7366512'/0'")?,
    )
    .derive_with_master_seed(secret.to_vec())?;
    let derived_cert = keyfork_derive_openpgp::derive(
        kdr,
        &[KeyFlags::empty().set_certification().set_signing()],
        userid,
    )?;

    // NOTE: Signatures on certs will be different. Compare fingerprints instead.
    let derived_fp = derived_cert.fingerprint();
    let expected_fp = root_cert.fingerprint();
    if derived_fp != expected_fp {
        return Err(Error::InvalidSecret(derived_fp, expected_fp));
    }
    */

    w.write_all(&secret)?;

    Ok(())
}