secretumvault/src/core/seal.rs

use serde::{Deserialize, Serialize};
use sharks::{Share, Sharks};

use crate::config::SealConfig;
use crate::crypto::CryptoBackend;
use crate::error::{CryptoError, CryptoResult, Result, VaultError};
use crate::storage::StorageBackend;

/// Master key used to encrypt all secrets in the vault
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MasterKey {
    pub key_data: Vec<u8>,
}

impl MasterKey {
    /// Generate a new random 32-byte master key
    pub async fn generate(crypto: &dyn CryptoBackend) -> CryptoResult<Self> {
        let key_data = crypto.random_bytes(32).await?;
        Ok(Self { key_data })
    }
}

/// State of the vault seal mechanism
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SealState {
    /// Vault is sealed (master key is split across shares)
    Sealed,
    /// Vault is unsealed (master key is reconstructed and available)
    Unsealed,
}

/// Seal mechanism using Shamir Secret Sharing
#[derive(Debug)]
pub struct SealMechanism {
    threshold: u8,
    shares_count: usize,
    state: SealState,
    master_key: Option<MasterKey>,
}

impl SealMechanism {
    /// Create a new seal mechanism from configuration
    pub fn new(config: &SealConfig) -> Result<Self> {
        // Validate threshold and shares
        if config.shamir.threshold == 0 || config.shamir.shares == 0 {
            return Err(VaultError::crypto(
                "Threshold and shares must be greater than 0".to_string(),
            ));
        }

        if config.shamir.threshold > config.shamir.shares {
            return Err(VaultError::crypto(
                "Threshold cannot be greater than shares count".to_string(),
            ));
        }

        if config.shamir.threshold > 255 {
            return Err(VaultError::crypto(
                "Threshold cannot exceed 255".to_string(),
            ));
        }

        Ok(Self {
            threshold: config.shamir.threshold as u8,
            shares_count: config.shamir.shares,
            state: SealState::Sealed,
            master_key: None,
        })
    }

    /// Get current seal state
    pub fn state(&self) -> SealState {
        self.state
    }

    /// Check if vault is sealed
    pub fn is_sealed(&self) -> bool {
        self.state == SealState::Sealed
    }

    /// Initialize the vault with a new master key (first initialization)
    pub async fn init(
        &mut self,
        crypto: &dyn CryptoBackend,
        storage: &dyn StorageBackend,
    ) -> Result<SealInitResult> {
        if self.state == SealState::Unsealed {
            return Err(VaultError::crypto(
                "Vault is already initialized".to_string(),
            ));
        }

        // Generate new master key
        let master_key = MasterKey::generate(crypto).await?;

        // Split into Shamir shares
        let shares = self.split_into_shares(&master_key)?;

        // Store shares (would be distributed to operators in production)
        let share_storage = ShareStorage {
            shares: shares.iter().map(|s| s.to_vec()).collect(),
            threshold: self.threshold,
            shares_count: self.shares_count as u8,
        };

        let share_json =
            serde_json::to_string(&share_storage).map_err(|e| VaultError::crypto(e.to_string()))?;

        storage
            .store_secret(
                "sys/seal/shares",
                &crate::storage::EncryptedData {
                    ciphertext: share_json.as_bytes().to_vec(),
                    nonce: vec![],
                    algorithm: "plain".to_string(),
                },
            )
            .await
            .map_err(|e| VaultError::storage(e.to_string()))?;

        self.master_key = Some(master_key);
        self.state = SealState::Unsealed;

        Ok(SealInitResult {
            shares: shares.iter().map(hex::encode).collect(),
            threshold: self.threshold,
        })
    }

    /// Unseal the vault using provided shares
    pub fn unseal(&mut self, shares_data: &[&[u8]]) -> CryptoResult<()> {
        use std::convert::TryFrom;

        if shares_data.len() < self.threshold as usize {
            return Err(CryptoError::InvalidAlgorithm(format!(
                "Need at least {} shares to unseal, got {}",
                self.threshold,
                shares_data.len()
            )));
        }

        // Parse shares from byte slices
        let shares: std::result::Result<Vec<Share>, _> = shares_data
            .iter()
            .map(|data| Share::try_from(*data).map_err(|_| "Invalid share format"))
            .collect();

        let shares = shares
            .map_err(|_| CryptoError::InvalidAlgorithm("Failed to parse shares".to_string()))?;

        // Reconstruct master key from shares
        let sharks = Sharks(self.threshold);
        let reconstructed = sharks.recover(shares.as_slice()).map_err(|e| {
            CryptoError::InvalidAlgorithm(format!("Failed to reconstruct secret: {:?}", e))
        })?;

        self.master_key = Some(MasterKey {
            key_data: reconstructed,
        });
        self.state = SealState::Unsealed;

        Ok(())
    }

    /// Seal the vault (clear master key from memory)
    pub fn seal(&mut self) {
        self.master_key = None;
        self.state = SealState::Sealed;
    }

    /// Get the master key (panics if vault is sealed)
    pub fn master_key(&self) -> Result<&MasterKey> {
        self.master_key
            .as_ref()
            .ok_or_else(|| VaultError::crypto("Vault is sealed".to_string()))
    }

    /// Split master key into Shamir shares
    fn split_into_shares(&self, key: &MasterKey) -> CryptoResult<Vec<Vec<u8>>> {
        let sharks = Sharks(self.threshold);
        let dealer = sharks.dealer(&key.key_data);

        let shares: Vec<Vec<u8>> = dealer
            .take(self.shares_count)
            .map(|share| Vec::<u8>::from(&share))
            .collect();

        if shares.len() != self.shares_count {
            return Err(CryptoError::InvalidAlgorithm(
                "Failed to generate required number of shares".to_string(),
            ));
        }

        Ok(shares)
    }
}

/// Result of vault initialization
#[derive(Debug, Serialize)]
pub struct SealInitResult {
    pub shares: Vec<String>, // Hex-encoded shares
    pub threshold: u8,
}

/// Storage structure for vault shares
#[derive(Debug, Serialize, Deserialize)]
struct ShareStorage {
    shares: Vec<Vec<u8>>,
    threshold: u8,
    shares_count: u8,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_seal_mechanism_creation() {
        let shamir_config = crate::config::ShamirSealConfig {
            threshold: 3,
            shares: 5,
        };
        let config = SealConfig {
            seal_type: "shamir".to_string(),
            shamir: shamir_config,
            auto_unseal: Default::default(),
        };
        let seal = SealMechanism::new(&config).expect("Failed to create seal");
        assert_eq!(seal.state(), SealState::Sealed);
        assert!(seal.is_sealed());
    }

    #[test]
    fn test_invalid_threshold() {
        let shamir_config = crate::config::ShamirSealConfig {
            threshold: 5,
            shares: 3, // threshold > shares
        };
        let config = SealConfig {
            seal_type: "shamir".to_string(),
            shamir: shamir_config,
            auto_unseal: Default::default(),
        };
        assert!(SealMechanism::new(&config).is_err());
    }

    #[test]
    fn test_shamir_reconstruct() {
        let shamir_config = crate::config::ShamirSealConfig {
            threshold: 2,
            shares: 3,
        };
        let config = SealConfig {
            seal_type: "shamir".to_string(),
            shamir: shamir_config,
            auto_unseal: Default::default(),
        };
        let seal = SealMechanism::new(&config).expect("Failed to create seal");
        let key = MasterKey {
            key_data: vec![42u8; 32],
        };

        let shares = seal.split_into_shares(&key).expect("Failed to split");
        assert_eq!(shares.len(), 3);

        // Test reconstruction with threshold shares
        let mut seal2 = SealMechanism::new(&config).expect("Failed to create seal");
        let share_refs: Vec<&[u8]> = vec![&shares[0], &shares[1]];
        seal2.unseal(&share_refs).expect("Failed to unseal");
        assert!(!seal2.is_sealed());
        assert_eq!(seal2.master_key().unwrap().key_data, key.key_data);
    }

    #[test]
    fn test_seal_unseal_cycle() {
        let shamir_config = crate::config::ShamirSealConfig {
            threshold: 2,
            shares: 3,
        };
        let config = SealConfig {
            seal_type: "shamir".to_string(),
            shamir: shamir_config,
            auto_unseal: Default::default(),
        };
        let mut seal = SealMechanism::new(&config).expect("Failed to create seal");

        // Initially sealed
        assert!(seal.is_sealed());

        // Seal again (should be no-op)
        seal.seal();
        assert!(seal.is_sealed());
    }
}