uvb-mrvb 0.2.1

//! Core types for MRVB assertion signing and verification.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// MRVB signing mode configuration.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[derive(Default)]
pub enum MrvbMode {
    /// Classical algorithms only (Ed25519) - maximum compatibility
    #[default]
    ClassicalOnly,
    /// Post-quantum only (Dilithium3) - for future-proofing / internal systems
    #[cfg(feature = "pqc")]
    PqcOnly,
    /// Hybrid: both classical + PQC signatures - recommended for production
    #[cfg(feature = "hybrid")]
    Hybrid,
}

impl std::fmt::Display for MrvbMode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            MrvbMode::ClassicalOnly => write!(f, "ClassicalOnly"),
            #[cfg(feature = "pqc")]
            MrvbMode::PqcOnly => write!(f, "PqcOnly"),
            #[cfg(feature = "hybrid")]
            MrvbMode::Hybrid => write!(f, "Hybrid"),
        }
    }
}

/// MRVB configuration for signing operations.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MrvbConfig {
    /// Current signing mode
    pub mode: MrvbMode,
    /// Active keyset ID (for key rotation tracking)
    pub keyset_id: String,
}

impl Default for MrvbConfig {
    fn default() -> Self {
        Self {
            mode: MrvbMode::default(),
            keyset_id: "default".to_string(),
        }
    }
}

/// Classical keypair (e.g., Ed25519).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClassicalKeyPair {
    pub key_id: String,
    /// Algorithm identifier (e.g., "ed25519")
    pub algorithm: String,
    /// Private key bytes, base64 encoded for JSON serialization
    #[serde(with = "base64_bytes")]
    pub private_key: Vec<u8>,
    /// Public key bytes, base64 encoded for JSON serialization
    #[serde(with = "base64_bytes")]
    pub public_key: Vec<u8>,
}

/// Post-quantum keypair (e.g., Dilithium3).
#[cfg(feature = "pqc")]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PqcKeyPair {
    pub key_id: String,
    /// Algorithm identifier (e.g., "dilithium3")
    pub algorithm: String,
    /// Private key bytes, base64 encoded
    #[serde(with = "base64_bytes")]
    pub private_key: Vec<u8>,
    /// Public key bytes, base64 encoded
    #[serde(with = "base64_bytes")]
    pub public_key: Vec<u8>,
}

/// Stub for non-pqc builds
#[cfg(not(feature = "pqc"))]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PqcKeyPair {
    pub key_id: String,
    pub algorithm: String,
    #[serde(with = "base64_bytes")]
    pub private_key: Vec<u8>,
    #[serde(with = "base64_bytes")]
    pub public_key: Vec<u8>,
}

/// Combined keyset for hybrid signing/verification.
///
/// Matches the design from the MRVB+KMS pack.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KeyPairSet {
    pub keyset_id: String,
    /// Classical keypair (required for ClassicalOnly and Hybrid modes)
    pub classical: Option<ClassicalKeyPair>,
    /// PQC keypair (required for PqcOnly and Hybrid modes)
    pub pqc: Option<PqcKeyPair>,
    /// Creation timestamp
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub created_at: Option<chrono::DateTime<chrono::Utc>>,
    /// Rotation timestamp (when this keyset should be rotated)
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub rotate_after: Option<chrono::DateTime<chrono::Utc>>,
}

impl KeyPairSet {
    /// Check if this keyset supports the given signing mode.
    pub fn supports_mode(&self, mode: MrvbMode) -> bool {
        match mode {
            MrvbMode::ClassicalOnly => self.classical.is_some(),
            #[cfg(feature = "pqc")]
            MrvbMode::PqcOnly => self.pqc.is_some(),
            #[cfg(feature = "hybrid")]
            MrvbMode::Hybrid => self.classical.is_some() && self.pqc.is_some(),
        }
    }
}

/// Result of a hybrid signature operation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HybridSignature {
    /// Classical signature bytes (base64 encoded)
    #[serde(
        default,
        skip_serializing_if = "Option::is_none",
        with = "option_base64_bytes"
    )]
    pub classical_sig: Option<Vec<u8>>,
    /// PQC signature bytes (base64 encoded)
    #[serde(
        default,
        skip_serializing_if = "Option::is_none",
        with = "option_base64_bytes"
    )]
    pub pqc_sig: Option<Vec<u8>>,
    /// Classical algorithm used
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub alg_classical: Option<String>,
    /// PQC algorithm used
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub alg_pqc: Option<String>,
    /// Keyset ID used for signing
    pub keyset_id: String,
    /// Signing mode used
    pub mode: MrvbMode,
}

impl HybridSignature {
    /// Check if this signature has at least one valid component.
    pub fn has_any_signature(&self) -> bool {
        self.classical_sig.is_some() || self.pqc_sig.is_some()
    }

    /// Check if this signature has both components (for hybrid mode).
    #[cfg(feature = "hybrid")]
    pub fn has_both_signatures(&self) -> bool {
        self.classical_sig.is_some() && self.pqc_sig.is_some()
    }
}

/// Assertion claims structure for MRVB verification tokens.
///
/// This is similar to JWT claims but specialized for MRVB verification flows.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct AssertionClaims {
    /// Unique session identifier
    pub session_id: String,
    /// User identifier (optional, may be omitted for privacy)
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub user_id: Option<String>,
    /// Rail used for verification (e.g., "email", "sms", "webauthn")
    pub rail: String,
    /// Verification confidence level (e.g., "low", "medium", "high")
    pub verification_level: String,
    /// Issuance timestamp
    #[serde(with = "chrono::serde::ts_seconds")]
    pub issued_at: chrono::DateTime<chrono::Utc>,
    /// Expiration timestamp
    #[serde(with = "chrono::serde::ts_seconds")]
    pub expires_at: chrono::DateTime<chrono::Utc>,
    /// Additional metadata (extensible)
    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
    pub metadata: HashMap<String, serde_json::Value>,
}

impl AssertionClaims {
    /// Check if the assertion has expired.
    pub fn is_expired(&self) -> bool {
        chrono::Utc::now() > self.expires_at
    }

    /// Check if the assertion is valid (not expired and issued in the past).
    pub fn is_valid(&self) -> bool {
        let now = chrono::Utc::now();
        now >= self.issued_at && now <= self.expires_at
    }
}

/// A signed MRVB assertion token.
///
/// This combines the assertion claims with a cryptographic signature.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SignedAssertion {
    /// The assertion claims (payload)
    pub claims: AssertionClaims,
    /// The cryptographic signature
    pub signature: HybridSignature,
    /// Token format version (for future compatibility)
    #[serde(default = "default_version")]
    pub version: String,
}

fn default_version() -> String {
    "1.0".to_string()
}

impl SignedAssertion {
    /// Serialize to JSON for transmission or storage.
    pub fn to_json(&self) -> Result<String, serde_json::Error> {
        serde_json::to_string(self)
    }

    /// Deserialize from JSON.
    pub fn from_json(json: &str) -> Result<Self, serde_json::Error> {
        serde_json::from_str(json)
    }

    /// Serialize to compact JSON (no whitespace).
    pub fn to_compact_json(&self) -> Result<String, serde_json::Error> {
        serde_json::to_string(self)
    }
}

// ============================================================================
// Base64 serde helpers for clean JSON serialization of byte arrays
// ============================================================================

mod base64_bytes {
    use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
    use serde::{Deserialize, Deserializer, Serializer};

    pub fn serialize<S>(bytes: &Vec<u8>, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.serialize_str(&BASE64.encode(bytes))
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
    where
        D: Deserializer<'de>,
    {
        let s = String::deserialize(deserializer)?;
        BASE64.decode(&s).map_err(serde::de::Error::custom)
    }
}

mod option_base64_bytes {
    use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
    use serde::{Deserialize, Deserializer, Serializer};

    pub fn serialize<S>(bytes: &Option<Vec<u8>>, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        match bytes {
            Some(b) => serializer.serialize_some(&BASE64.encode(b)),
            None => serializer.serialize_none(),
        }
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<Option<Vec<u8>>, D::Error>
    where
        D: Deserializer<'de>,
    {
        let opt: Option<String> = Option::deserialize(deserializer)?;
        match opt {
            Some(s) => BASE64
                .decode(&s)
                .map(Some)
                .map_err(serde::de::Error::custom),
            None => Ok(None),
        }
    }
}