keyfork/keyfork-shard/src/lib.rs

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

use aes_gcm::{
    aead::{Aead, AeadCore, OsRng},
    Aes256Gcm, KeyInit,
};
use keyfork_mnemonic_util::{Mnemonic, Wordlist};
use keyfork_prompt::{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),
}

/// 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;

pub fn remote_decrypt() -> 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())?;
        pm.prompt_message(PromptMessage::Text(format!(
            "Our words: {nonce_mnemonic} {key_mnemonic}"
        )))?;

        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().expect("24 words");

        let shared_secret = our_key
            .diffie_hellman(&PublicKey::from(their_key))
            .to_bytes();
        let shared_key =
            Aes256Gcm::new_from_slice(&shared_secret).expect("Invalid length of constant key size");

        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));
    }
    */

    print!("{}", smex::encode(&secret));

    Ok(())
}