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

use std::{
    collections::{HashMap, HashSet},
    sync::{Arc, Mutex},
};

use keyfork_prompt::{
    validators::{PinValidator, Validator},
    Error as PromptError, Message, PromptHandler,
};

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

use anyhow::Context;
use card_backend_pcsc::PcscBackend;
use openpgp_card::{Error as CardError, StatusBytes};
use openpgp_card_sequoia::{state::Open, types::Error as SequoiaCardError, Card};

#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error("Smart card could not decrypt any matching PKESK packets")]
    SmartCardCouldNotDecrypt,

    #[error("No smart card was found")]
    SmartCardNotFound,

    #[error("Selected smart card has no decryption key")]
    SmartCardHasNoDecrypt,

    #[error("Smart card backend error: {0}")]
    SmartCardBackend(#[from] card_backend::SmartcardError),

    #[error("Smartcard password status unavailable: {0}")]
    PwStatusBytes(SequoiaCardError),

    #[error("Could not open smart card")]
    OpenSmartCard(SequoiaCardError),

    #[error("Could not initialize transaction")]
    Transaction(SequoiaCardError),

    #[error("Could not load fingerprints")]
    Fingerprints(SequoiaCardError),

    #[error("Invalid PIN entered too many times")]
    InvalidPIN,

    #[error("Prompt failed: {0}")]
    Prompt(#[from] PromptError),
}

pub type Result<T, E = Error> = std::result::Result<T, E>;

fn format_name(input: impl AsRef<str>) -> String {
    let mut n = input
        .as_ref()
        .split("<<")
        .take(2)
        .map(|s| s.replace('<', " "))
        .collect::<Vec<_>>();
    n.reverse();
    n.join(" ")
}

#[allow(clippy::module_name_repetitions)]
pub struct SmartcardManager<P: PromptHandler> {
    current_card: Option<Card<Open>>,
    root: Option<Cert>,
    pm: Arc<Mutex<P>>,
    pin_cache: HashMap<Fingerprint, String>,
}

impl<P: PromptHandler> SmartcardManager<P> {
    pub fn new(p: Arc<Mutex<P>>) -> Result<Self> {
        Ok(Self {
            current_card: None,
            root: None,
            pm: p,
            pin_cache: 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
    }

    /// Load any backend.
    pub fn load_any_card(&mut self) -> Result<Fingerprint> {
        let card_backend = loop {
            if let Some(c) = PcscBackend::cards(None)?.next().transpose()? {
                break c;
            }
            self.pm
                .lock()
                .unwrap()
                .prompt_message(Message::Text(
                    "No smart card was found. Please plug in a smart card and press enter"
                        .to_string(),
                ))?;
        };
        let mut card = Card::<Open>::new(card_backend).map_err(Error::OpenSmartCard)?;
        let transaction = card.transaction().map_err(Error::Transaction)?;
        let fingerprint = transaction
            .fingerprints()
            .map_err(Error::Fingerprints)?
            .decryption()
            .map(|fp| Fingerprint::from_bytes(fp.as_bytes()))
            .ok_or(Error::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>> {
        // 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).map_err(Error::OpenSmartCard)?;
                let transaction = card.transaction().map_err(Error::Transaction)?;
                let mut fingerprint = None;
                if let Some(fp) = transaction
                    .fingerprints()
                    .map_err(Error::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);
                }
            }

            self.pm
                .lock()
                .unwrap()
                .prompt_message(Message::Text(
                    "Please plug in a smart card and press enter".to_string(),
                ))?;
        }

        Ok(None)
    }
}

impl<P: PromptHandler> VerificationHelper for &mut SmartcardManager<P> {
    fn get_certs(&mut self, ids: &[openpgp::KeyHandle]) -> openpgp::Result<Vec<Cert>> {
        #[allow(clippy::flat_map_option)]
        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 {
            #[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!("Verification error: {}", e.to_string()));
                        }
                    }
                }
            }
        }
        Ok(())
    }
}

impl<P: PromptHandler> DecryptionHelper for &mut SmartcardManager<P> {
    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(Error::SmartCardNotFound.into());
        };

        let mut transaction = card
            .transaction()
            .context("Could not initialize transaction")?;
        let fp = transaction
            .fingerprints()
            .context("Could not load fingerprints")?
            .decryption()
            .map(|fp| Fingerprint::from_bytes(fp.as_bytes()))
            .ok_or(Error::SmartCardHasNoDecrypt)?;
        let cardholder_name = format_name(
            transaction
                .cardholder_name()
                .context("Could not load (optionally empty) cardholder name")?,
        );
        let card_id = transaction
            .application_identifier()
            .context("Could not load application identifier")?
            .ident();
        let mut pin = self.pin_cache.get(&fp).cloned();
        let pin_validator = PinValidator {
            min_length: Some(6),
            ..Default::default()
        }
        .to_fn();
        while transaction
            .pw_status_bytes()
            .map_err(Error::PwStatusBytes)?
            .err_count_pw1()
            > 0
            && pin.is_none()
        {
            transaction.reload_ard()?;
            let attempts = transaction
                .pw_status_bytes()
                .map_err(Error::PwStatusBytes)?
                .err_count_pw1();
            let rpea = "Remaining PIN entry attempts";
            let message = if cardholder_name.is_empty() {
                format!("Unlock card {card_id}\n{rpea}: {attempts}\n\nPIN: ")
            } else {
                format!("Unlock card {card_id} ({cardholder_name})\n{rpea}: {attempts}\n\nPIN: ")
            };
            let temp_pin = self
                .pm
                .lock()
                .unwrap()
                .prompt_validated_passphrase(&message, 3, &pin_validator)?;
            let verification_status = transaction.verify_user_pin(temp_pin.as_str().trim());
            match verification_status {
                #[allow(clippy::ignored_unit_patterns)]
                Ok(_) => {
                    self.pin_cache.insert(fp.clone(), temp_pin.clone());
                    pin.replace(temp_pin);
                }
                // NOTE: This should not be hit, because of the above validator.
                Err(CardError::CardStatus(StatusBytes::IncorrectParametersCommandDataField)) => {
                    self.pm
                        .lock()
                        .unwrap()
                        .prompt_message(Message::Text("Invalid PIN length entered.".to_string()))?;
                }
                Err(_) => {}
            }
        }
        let pin = pin.ok_or(Error::InvalidPIN)?;
        let mut user = transaction
            .to_user_card(pin.as_str().trim())
            .context("Could not load user smartcard from PIN")?;
        let mut decryptor = user
            .decryptor(&|| eprintln!("Touch confirmation needed for decryption"))
            .context("Could not decrypt using smartcard")?;
        for pkesk in pkesks {
            if pkesk
                .decrypt(&mut decryptor, sym_algo)
                .is_some_and(|(algo, sk)| decrypt(algo, &sk))
            {
                return Ok(Some(fp));
            }
        }

        Err(Error::SmartCardCouldNotDecrypt.into())
    }
}