keyfork/keyfork-shard/src/openpgp/smartcard.rs

use std::collections::HashSet;

use super::openpgp::{
    self,
    cert::Cert,
    packet::{PKESK, SKESK},
    parse::stream::{DecryptionHelper, MessageLayer, MessageStructure, VerificationHelper},
    Fingerprint,
};

use card_backend_pcsc::PcscBackend;
use openpgp_card_sequoia::{state::Open, Card};

#[derive(Clone, Debug)]
pub enum SmartcardFailure {
    SmartCardNotFound,
    SmartCardHasNoDecrypt,
}

impl std::fmt::Display for SmartcardFailure {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::SmartCardNotFound => f.write_str("No smart card backend was stored"),
            Self::SmartCardHasNoDecrypt => f.write_str("Selected smart card has no decrypt key"),
        }
    }
}

impl std::error::Error for SmartcardFailure {}

#[derive(Default)]
pub struct SmartcardManager {
    current_card: Option<Card<Open>>,
    root: Option<Cert>,
}

impl SmartcardManager {
    pub fn new() -> Self {
        Default::default()
    }

    // Sets the root cert, returning the old cert
    pub fn set_root_cert(&mut self, cert: impl Into<Option<Cert>>) -> Option<Cert> {
        let mut cert = cert.into();
        std::mem::swap(&mut self.root, &mut cert);
        cert
    }

    /// Utility function to prompt for a newline from standard input.
    pub fn prompt(&self, prompt: impl std::fmt::Display) -> std::io::Result<()> {
        eprint!("{prompt}: ");
        std::io::stdin().read_line(&mut String::new()).map(|_| ())
    }

    /// Utility function to obtain a prompt response from the command line.
    pub fn prompt_input(&self, prompt: impl std::fmt::Display) -> std::io::Result<String> {
        eprint!("{prompt}: ");
        let mut output = String::new();
        std::io::stdin().read_line(&mut output)?;
        Ok(output)
    }

    /*
    /// Return all [`Fingerprint`] for the currently accessible backends.
    ///
    /// NOTE: Only implemented for decryption keys.
    pub fn iter_fingerprints() -> impl Iterator<Item = Fingerprint> {
        PcscBackend::cards(None).into_iter().flat_map(|iter| {
            iter.filter_map(|backend| {
                let backend = backend.ok()?;
                let mut card = Card::<Open>::new(backend).ok()?;
                let transaction = card.transaction().ok()?;
                transaction
                    .fingerprints()
                    .ok()?
                    .decryption()
                    .map(|fp| Fingerprint::from_bytes(fp.as_bytes()))
            })
        })
    }
    */

    /// Load any backend.
    pub fn load_any_card(&mut self) -> Result<Fingerprint, Box<dyn std::error::Error>> {
        PcscBackend::cards(None)?
            .next()
            .transpose()?
            .ok_or(SmartcardFailure::SmartCardNotFound.into())
            .and_then(
                |backend| -> Result<Fingerprint, Box<dyn std::error::Error>> {
                    let mut card = Card::<Open>::new(backend)?;
                    let transaction = card.transaction()?;
                    let fingerprint = transaction
                        .fingerprints()?
                        .decryption()
                        .map(|fp| Fingerprint::from_bytes(fp.as_bytes()))
                        .ok_or(SmartcardFailure::SmartCardHasNoDecrypt)?;
                    drop(transaction);
                    self.current_card.replace(card);
                    Ok(fingerprint)
                },
            )
    }

    /// Load a backend if any [`Fingerprint`] has been matched by a currently active card.
    ///
    /// NOTE: Only implemented for decryption keys.
    pub fn load_any_fingerprint(
        &mut self,
        fingerprints: impl IntoIterator<Item = Fingerprint>,
    ) -> Result<Option<Fingerprint>, Box<dyn std::error::Error>> {
        // NOTE: This can't be HashSet::from_iter() because from_iter() requires a passed-in state
        // I do not want to provide.
        let mut requested_fingerprints = HashSet::new();
        requested_fingerprints.extend(fingerprints);

        let mut had_any_backend = false;

        while !had_any_backend {
            // Load all backends, confirm if any have any fingerprints
            for backend in PcscBackend::cards(None)? {
                had_any_backend = true;
                let backend = backend?;
                let mut card = Card::<Open>::new(backend)?;
                let transaction = card.transaction()?;
                let mut fingerprint = None;
                if let Some(fp) = transaction
                    .fingerprints()?
                    .decryption()
                    .map(|fp| Fingerprint::from_bytes(fp.as_bytes()))
                {
                    if requested_fingerprints.contains(&fp) {
                        fingerprint.replace(fp);
                    }
                }
                drop(transaction);
                if fingerprint.is_some() {
                    self.current_card.replace(card);
                    return Ok(fingerprint);
                }
            }

            eprintln!("No matching smartcard detected.");
            self.prompt("Please plug in a smart card and press enter")?;
        }

        Ok(None)
    }
}

impl VerificationHelper for &mut SmartcardManager {
    fn get_certs(&mut self, ids: &[openpgp::KeyHandle]) -> openpgp::Result<Vec<Cert>> {
        Ok(ids
            .iter()
            .flat_map(|kh| self.root.as_ref().filter(|cert| cert.key_handle() == *kh))
            .cloned()
            .collect())
    }

    fn check(&mut self, structure: MessageStructure) -> openpgp::Result<()> {
        for layer in structure.into_iter() {
            #[allow(unused_variables)]
            match layer {
                MessageLayer::Compression { algo } => {}
                MessageLayer::Encryption {
                    sym_algo,
                    aead_algo,
                } => {}
                MessageLayer::SignatureGroup { results } => {
                    for result in results {
                        if let Err(e) = result {
                            // FIXME: anyhow leak
                            return Err(anyhow::anyhow!(e.to_string()));
                        }
                    }
                }
            }
        }
        Ok(())
    }
}

impl DecryptionHelper for &mut SmartcardManager {
    fn decrypt<D>(
        &mut self,
        pkesks: &[PKESK],
        _skesks: &[SKESK],
        sym_algo: Option<openpgp::types::SymmetricAlgorithm>,
        mut decrypt: D,
    ) -> openpgp::Result<Option<Fingerprint>>
    where
        D: FnMut(openpgp::types::SymmetricAlgorithm, &openpgp::crypto::SessionKey) -> bool,
    {
        let mut card = self.current_card.take();
        let Some(card) = card.as_mut() else {
            return Err(SmartcardFailure::SmartCardNotFound.into());
        };

        let mut transaction = card.transaction()?;
        let fp = transaction
            .fingerprints()?
            .decryption()
            .map(|fp| Fingerprint::from_bytes(fp.as_bytes()));
        let pin = self.prompt_input("User PIN")?;
        let mut user = transaction.to_user_card(pin.as_str().trim())?;
        let mut decryptor =
            user.decryptor(&|| eprintln!("Touch confirmation needed for decryption"))?;
        for pkesk in pkesks {
            if pkesk
                .decrypt(&mut decryptor, sym_algo)
                .map(|(algo, sk)| decrypt(algo, &sk))
                .unwrap_or(false)
            {
                return Ok(fp);
            }
        }

        Err(SmartcardFailure::SmartCardNotFound.into())
    }
}