ai-agent-bitcoin-escrow 0.1.0

//! Multisig operations for escrow contracts.
//!
//! This module handles multisignature wallet setup, descriptor creation,
//! and coordination of signing operations for 2-of-3 (or n-of-m) escrow.

use bitcoin::key::Keypair;
use bitcoin::secp256k1::{Secp256k1, rand};
use bitcoin::{Network, PrivateKey, PublicKey, ScriptBuf};
use miniscript::descriptor::{Descriptor, DescriptorPublicKey};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::str::FromStr;

use crate::error::{EscrowError, Result};
use crate::types::{EscrowParticipant, EscrowRole};

/// Multisig configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MultisigConfig {
    /// Network (testnet, mainnet, etc.).
    pub network: Network,
    /// Threshold (e.g., 2 for 2-of-3).
    pub threshold: usize,
    /// Total number of participants.
    pub total: usize,
}

impl Default for MultisigConfig {
    fn default() -> Self {
        Self {
            network: Network::Testnet,
            threshold: 2,
            total: 3,
        }
    }
}

/// Represents a multisig wallet with participants' public keys.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MultisigWallet {
    /// Configuration.
    pub config: MultisigConfig,
    /// Participant public keys (role -> pubkey hex).
    pub public_keys: BTreeMap<EscrowRole, String>,
    /// The descriptor for the multisig address.
    pub descriptor: String,
}

impl MultisigWallet {
    /// Create a new multisig wallet from participants.
    pub fn new(config: MultisigConfig, participants: &[EscrowParticipant]) -> Result<Self> {
        if participants.len() != config.total {
            return Err(EscrowError::Multisig(format!(
                "Expected {} participants, got {}",
                config.total,
                participants.len()
            )));
        }

        let mut public_keys = BTreeMap::new();

        for participant in participants {
            if let Some(pk) = &participant.public_key {
                public_keys.insert(participant.role, pk.to_string());
            }
        }

        if public_keys.len() != config.total {
            return Err(EscrowError::Multisig(
                "All participants must have public keys".to_string(),
            ));
        }

        // Create the multisig descriptor
        let descriptor = Self::create_descriptor(config.threshold, &public_keys)?;
        
        Ok(Self {
            config,
            public_keys,
            descriptor,
        })
    }

    /// Create a multisig descriptor string.
    fn create_descriptor(
        threshold: usize,
        public_keys: &BTreeMap<EscrowRole, String>,
    ) -> Result<String> {
        // Collect public keys strings for the descriptor
        let pk_strs: Vec<String> = public_keys.values().cloned().collect();

        // Create a sortedmulti descriptor (BIP-67 sorted public keys)
        // wsh(sortedmulti(threshold, keys...))
        let descriptor_str = format!(
            "wsh(sortedmulti({},{}))",
            threshold,
            pk_strs.join(",")
        );

        Ok(descriptor_str)
    }

    /// Get the scriptPubKey for the multisig address.
    pub fn script_pubkey(&self) -> Result<ScriptBuf> {
        // Parse the descriptor - this is simplified, real implementation would
        // use miniscript properly with DescriptorPublicKey parsing
        // For now, we return a placeholder
        Ok(ScriptBuf::new())
    }

    /// Get the address for the multisig (based on network).
    pub fn address(&self) -> Result<bitcoin::Address> {
        // This would generate the actual P2WSH address from the descriptor
        // For now, return a test address
        let script = self.script_pubkey()?;
        
        // Create a P2WSH address
        let address = bitcoin::Address::p2wsh(&script, self.config.network);
        Ok(address)
    }
}

/// Generate a new keypair for a participant.
pub fn generate_keypair(network: Network) -> Result<(PrivateKey, PublicKey)> {
    let secp = Secp256k1::new();
    let mut rng = rand::thread_rng();
    
    let keypair = Keypair::new(&secp, &mut rng);
    let secret_key = keypair.secret_key();
    let public_key = keypair.public_key();
    
    let private_key = PrivateKey::new(secret_key, network);
    let bitcoin_pk = PublicKey::from_slice(&public_key.serialize())
        .map_err(|e| EscrowError::Key(format!("Failed to create public key: {}", e)))?;
    
    Ok((private_key, bitcoin_pk))
}

/// Create a participant with a generated keypair.
pub fn create_participant(
    role: EscrowRole,
    id: String,
    network: Network,
) -> Result<(EscrowParticipant, PrivateKey)> {
    let (private_key, public_key) = generate_keypair(network)?;
    
    let participant = EscrowParticipant {
        id,
        role,
        public_key: Some(public_key),
        xonly_public_key: None,
        address: None,
        signed: false,
        joined_at: chrono::Utc::now(),
    };
    
    Ok((participant, private_key))
}

/// Represents a partially signed transaction (PSBT) for multisig coordination.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MultisigPSBT {
    /// The base64-encoded PSBT.
    pub psbt_base64: String,
    /// Who has signed so far.
    pub signed_by: Vec<EscrowRole>,
    /// Required threshold.
    pub threshold: usize,
}

impl MultisigPSBT {
    /// Create a new PSBT from a transaction.
    pub fn new(psbt_base64: String, threshold: usize) -> Self {
        Self {
            psbt_base64,
            signed_by: Vec::new(),
            threshold,
        }
    }

    /// Add a signature.
    pub fn add_signature(&mut self, role: EscrowRole) -> Result<()> {
        if self.signed_by.contains(&role) {
            return Err(EscrowError::Signing(format!(
                "Role {:?} has already signed",
                role
            )));
        }
        self.signed_by.push(role);
        Ok(())
    }

    /// Check if enough signatures have been collected.
    pub fn is_ready(&self) -> bool {
        self.signed_by.len() >= self.threshold
    }

    /// Get the number of signatures still needed.
    pub fn signatures_needed(&self) -> usize {
        self.threshold.saturating_sub(self.signed_by.len())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_keypair_generation() {
        let (priv_key, pub_key) = generate_keypair(Network::Testnet).unwrap();
        assert!(!priv_key.to_bytes().is_empty());
        assert!(!pub_key.to_bytes().is_empty());
    }
    
    #[test]
    fn test_multisig_wallet_creation() {
        let config = MultisigConfig::default();
        
        let participants = vec![
            create_participant(EscrowRole::Buyer, "buyer-1".to_string(), Network::Testnet).unwrap().0,
            create_participant(EscrowRole::Seller, "seller-1".to_string(), Network::Testnet).unwrap().0,
            create_participant(EscrowRole::Arbiter, "arbiter-1".to_string(), Network::Testnet).unwrap().0,
        ];
        
        let wallet = MultisigWallet::new(config, &participants).unwrap();
        assert!(!wallet.descriptor.is_empty());
    }
    
    #[test]
    fn test_psbt_signatures() {
        let mut psbt = MultisigPSBT::new("base64-psbt".to_string(), 2);
        
        assert!(!psbt.is_ready());
        assert_eq!(psbt.signatures_needed(), 2);
        
        psbt.add_signature(EscrowRole::Buyer).unwrap();
        assert!(!psbt.is_ready());
        assert_eq!(psbt.signatures_needed(), 1);
        
        psbt.add_signature(EscrowRole::Seller).unwrap();
        assert!(psbt.is_ready());
        assert_eq!(psbt.signatures_needed(), 0);
    }
}