logos-execution-zone/wallet/src/signing.rs

use anyhow::Result;
use keycard_wallet::{KeycardWallet, python_path};
use nssa::{AccountId, PrivateKey, PublicKey, Signature};
use pyo3::Python;

use crate::{WalletCore, cli::CliAccountMention};

/// Groups transaction signers by type to minimise Python GIL acquisition.
///
/// Local signers are signed in pure Rust; all keycard signers share a single Python session
/// with one `connect` / `close_session` pair.
#[derive(Default)]
pub struct SigningGroups {
    local: Vec<(AccountId, PrivateKey)>,
    keycard: Vec<(AccountId, String)>,
}

impl SigningGroups {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Add a sender. Keycard paths are queued for the hardware session; local accounts
    /// have their signing key resolved eagerly. Errors if no key is found.
    pub fn add_sender(
        &mut self,
        mention: &CliAccountMention,
        account_id: AccountId,
        wallet_core: &WalletCore,
    ) -> Result<()> {
        if let CliAccountMention::KeyPath(path) = mention {
            self.keycard.push((account_id, path.clone()));
            return Ok(());
        }
        let key = wallet_core
            .storage()
            .key_chain()
            .pub_account_signing_key(account_id)
            .ok_or_else(|| anyhow::anyhow!("signing key not found for account {account_id}"))?
            .clone();
        self.local.push((account_id, key));
        Ok(())
    }

    /// Add a recipient. Same as [`add_sender`] but silently skips accounts with no local
    /// key and no keycard path — they are foreign and require neither a signature nor a nonce.
    pub fn add_recipient(
        &mut self,
        mention: &CliAccountMention,
        account_id: AccountId,
        wallet_core: &WalletCore,
    ) -> Result<()> {
        if let CliAccountMention::KeyPath(path) = mention {
            self.keycard.push((account_id, path.clone()));
            return Ok(());
        }
        if let Some(key) = wallet_core
            .storage()
            .key_chain()
            .pub_account_signing_key(account_id)
        {
            self.local.push((account_id, key.clone()));
        }
        Ok(())
    }

    /// Returns `true` when a PIN is required (at least one keycard signer is present).
    #[must_use]
    pub const fn needs_pin(&self) -> bool {
        !self.keycard.is_empty()
    }

    /// Account IDs that require a nonce (every non-foreign signer).
    #[must_use]
    pub fn signing_ids(&self) -> Vec<AccountId> {
        self.local
            .iter()
            .map(|(id, _)| *id)
            .chain(self.keycard.iter().map(|(id, _)| *id))
            .collect()
    }

    /// Sign `hash` for every account in the group.
    ///
    /// Local accounts are signed in pure Rust. Keycard accounts share one Python session.
    pub fn sign_all(&self, hash: &[u8; 32], pin: &str) -> Result<Vec<(Signature, PublicKey)>> {
        let mut sigs: Vec<(Signature, PublicKey)> = self
            .local
            .iter()
            .map(|(_, key)| {
                (
                    Signature::new(key, hash),
                    PublicKey::new_from_private_key(key),
                )
            })
            .collect();

        if !self.keycard.is_empty() {
            pyo3::Python::with_gil(|py| -> pyo3::PyResult<()> {
                python_path::add_python_path(py)?;
                let wallet = KeycardWallet::new(py)?;
                wallet.connect(py, pin)?;
                for (_, path) in &self.keycard {
                    sigs.push(wallet.sign_message_for_path(py, path, hash)?);
                }
                drop(wallet.close_session(py));
                Ok(())
            })
            .map_err(anyhow::Error::from)?;
        }

        Ok(sigs)
    }
}

/// Lazily opens and reuses a single Keycard session for all keycard signers in one transaction.
pub struct KeycardSessionContext {
    pin: String,
    wallet: Option<KeycardWallet>,
}

impl KeycardSessionContext {
    pub fn new(pin: impl Into<String>) -> Self {
        Self {
            pin: pin.into(),
            wallet: None,
        }
    }

    pub fn get_or_connect<'py>(
        &'py mut self,
        py: Python<'py>,
    ) -> pyo3::PyResult<&'py KeycardWallet> {
        if self.wallet.is_none() {
            python_path::add_python_path(py)?;
            let wallet = KeycardWallet::new(py)?;
            wallet.connect(py, &self.pin)?;
            self.wallet = Some(wallet);
        }
        Ok(self.wallet.as_ref().unwrap())
    }

    pub fn close(self, py: Python<'_>) {
        if let Some(w) = self.wallet {
            drop(w.close_session(py));
        }
    }
}