mrapids 0.1.31

//! Policy model types for agent access control

use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap};
use std::fmt;

/// Data classification level for API operations
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
pub enum DataClassification {
    #[serde(rename = "public")]
    Public,
    #[serde(rename = "internal")]
    Internal,
    #[serde(rename = "confidential")]
    Confidential,
    #[serde(rename = "regulated")]
    Regulated,
}

impl fmt::Display for DataClassification {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DataClassification::Public => write!(f, "public"),
            DataClassification::Internal => write!(f, "internal"),
            DataClassification::Confidential => write!(f, "confidential"),
            DataClassification::Regulated => write!(f, "regulated"),
        }
    }
}

fn default_classification() -> DataClassification {
    DataClassification::Internal
}

/// Policy configuration set
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct PolicySet {
    /// Policy version (e.g., "1.0")
    pub version: String,

    /// Metadata about the policy
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub metadata: Option<PolicyMetadata>,

    /// Default settings that apply globally
    pub defaults: PolicyDefaults,

    /// List of policy rules evaluated in order
    pub rules: Vec<PolicyRule>,

    /// Per-operation data classification labels (BTreeMap for deterministic ordering)
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    pub classifications: BTreeMap<String, DataClassification>,

    /// Credential scoping rules
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub credential_scopes: Vec<CredentialScope>,
}

/// Credential scope — restricts what an auth profile can do
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct CredentialScope {
    /// Auth profile or API key name this scope applies to
    pub profile: String,

    /// Allowed HTTP methods for this credential (empty = all)
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub allowed_methods: Vec<String>,

    /// Allowed tags for this credential (empty = all)
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub allowed_tags: Vec<String>,

    /// Operations that require human approval with this credential
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub requires_approval: Vec<String>,
}

/// Policy metadata
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct PolicyMetadata {
    pub name: String,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub author: Option<String>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub created: Option<String>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub modified: Option<String>,
}

/// Default policy settings
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct PolicyDefaults {
    /// HTTP methods allowed by default
    #[serde(default = "default_allow_methods")]
    pub allow_methods: Vec<String>,

    /// Whether to deny external references by default
    #[serde(default = "default_true")]
    pub deny_external_refs: bool,

    /// Whether to require authentication by default
    #[serde(default = "default_true")]
    pub require_auth: bool,

    /// Default audit level
    #[serde(default = "default_audit_level")]
    pub audit_level: String,

    /// Default data classification for operations not explicitly classified
    #[serde(default = "default_classification")]
    pub default_classification: DataClassification,

    /// Read-only mode: block all POST/PUT/PATCH/DELETE globally
    #[serde(default)]
    pub read_only: bool,

    /// Max API calls per MCP session (0 = unlimited)
    #[serde(default)]
    pub max_calls_per_session: u32,

    /// Max API calls per minute (0 = unlimited)
    #[serde(default)]
    pub max_calls_per_minute: u32,

    /// Warn when confidential data will be sent to external LLM
    #[serde(default)]
    pub warn_on_confidential_to_llm: bool,

    /// Block regulated data from being sent to external LLM
    #[serde(default)]
    pub block_regulated_to_llm: bool,
}

/// Individual policy rule
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct PolicyRule {
    /// Rule name for identification
    pub name: String,

    /// Human-readable description
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,

    /// Pattern to match (glob-style, e.g., "api.github.com/*")
    pub pattern: String,

    /// Conditions that must be met for rule to apply
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub conditions: Option<Vec<PolicyCondition>>,

    /// Actions to allow
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub allow: Option<PolicyAction>,

    /// Actions to deny (takes precedence over allow)
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub deny: Option<PolicyAction>,

    /// Audit configuration for this rule
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub audit: Option<AuditConfig>,

    /// Custom explanation for denials
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub explain: Option<String>,
}

/// Policy condition
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct PolicyCondition {
    /// Required auth profile
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub auth_profile: Option<String>,

    /// Time window constraint (e.g., "business_hours", "weekdays")
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub time_window: Option<String>,

    /// Source IP constraint
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub source_ip: Option<String>,

    /// Environment constraint (e.g., "production", "staging")
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub environment: Option<String>,
}

/// Policy action specification
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct PolicyAction {
    /// Specific operations (supports wildcards, e.g., ["get*", "list*"])
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub operations: Option<Vec<String>>,

    /// HTTP methods
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub methods: Option<Vec<String>>,

    /// Allow/deny all operations
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub all: Option<bool>,

    /// Tags to match
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tags: Option<Vec<String>>,
}

/// Audit configuration
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct AuditConfig {
    /// Audit level: "none", "basic", "detailed"
    #[serde(default = "default_audit_level")]
    pub level: String,

    /// Include request body in audit
    #[serde(default)]
    pub include_body: bool,

    /// Include response in audit
    #[serde(default)]
    pub include_response: bool,
}

// Default value functions
fn default_allow_methods() -> Vec<String> {
    vec!["GET".to_string(), "HEAD".to_string(), "OPTIONS".to_string()]
}

fn default_true() -> bool {
    true
}

fn default_audit_level() -> String {
    "basic".to_string()
}

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

    #[test]
    fn test_policy_defaults() {
        let defaults = PolicyDefaults {
            allow_methods: default_allow_methods(),
            deny_external_refs: true,
            require_auth: true,
            audit_level: "basic".to_string(),
            default_classification: DataClassification::Internal,
            read_only: false,
            max_calls_per_session: 0,
            max_calls_per_minute: 0,
            warn_on_confidential_to_llm: false,
            block_regulated_to_llm: false,
        };

        assert_eq!(defaults.allow_methods, vec!["GET", "HEAD", "OPTIONS"]);
        assert!(defaults.deny_external_refs);
        assert!(defaults.require_auth);
    }

    #[test]
    fn test_data_flow_defaults() {
        // Verify data flow fields default to false for backward compatibility
        let yaml = r#"
version: "1.0"
defaults:
  allow_methods: ["GET"]
  require_auth: false
  audit_level: "basic"
rules: []
"#;
        let policy: PolicySet = serde_yaml::from_str(yaml).unwrap();
        assert_eq!(policy.defaults.warn_on_confidential_to_llm, false);
        assert_eq!(policy.defaults.block_regulated_to_llm, false);
    }

    #[test]
    fn test_policy_rule_serialization() {
        let rule = PolicyRule {
            name: "test-rule".to_string(),
            description: Some("Test rule".to_string()),
            pattern: "*.example.com/*".to_string(),
            conditions: None,
            allow: Some(PolicyAction {
                methods: Some(vec!["GET".to_string()]),
                operations: None,
                all: None,
                tags: None,
            }),
            deny: None,
            audit: None,
            explain: None,
        };

        let json = serde_json::to_string_pretty(&rule).unwrap();
        assert!(json.contains("test-rule"));
        assert!(json.contains("*.example.com/*"));
    }

    #[test]
    fn test_backward_compat_old_policy_yaml() {
        // Old policy files don't have read_only, max_calls_per_session, max_calls_per_minute
        // They should parse correctly with serde defaults (false, 0, 0)
        let yaml = r#"
version: "1.0"
defaults:
  allow_methods: ["GET"]
  deny_external_refs: true
  require_auth: true
  audit_level: "basic"
rules: []
"#;
        let policy: PolicySet = serde_yaml::from_str(yaml).unwrap();
        assert_eq!(policy.defaults.read_only, false);
        assert_eq!(policy.defaults.max_calls_per_session, 0);
        assert_eq!(policy.defaults.max_calls_per_minute, 0);
        assert_eq!(policy.defaults.allow_methods, vec!["GET"]);
        assert!(policy.defaults.require_auth);
    }

    #[test]
    fn test_new_policy_fields_yaml() {
        let yaml = r#"
version: "1.0"
defaults:
  allow_methods: ["GET"]
  require_auth: false
  audit_level: "basic"
  read_only: true
  max_calls_per_session: 100
  max_calls_per_minute: 10
rules: []
"#;
        let policy: PolicySet = serde_yaml::from_str(yaml).unwrap();
        assert!(policy.defaults.read_only);
        assert_eq!(policy.defaults.max_calls_per_session, 100);
        assert_eq!(policy.defaults.max_calls_per_minute, 10);
    }

    #[test]
    fn test_classification_serde() {
        let yaml = r#"
version: "1.0"
defaults:
  allow_methods: ["GET"]
  require_auth: false
  audit_level: "basic"
  default_classification: confidential
rules: []
classifications:
  "getPortfolio": confidential
  "getStockPrice": public
  "getThesis": internal
  "*admin*": regulated
"#;
        let policy: PolicySet = serde_yaml::from_str(yaml).unwrap();
        assert_eq!(
            policy.defaults.default_classification,
            DataClassification::Confidential
        );
        assert_eq!(
            policy.classifications.get("getPortfolio"),
            Some(&DataClassification::Confidential)
        );
        assert_eq!(
            policy.classifications.get("getStockPrice"),
            Some(&DataClassification::Public)
        );
        assert_eq!(
            policy.classifications.get("getThesis"),
            Some(&DataClassification::Internal)
        );
        assert_eq!(
            policy.classifications.get("*admin*"),
            Some(&DataClassification::Regulated)
        );

        // Round-trip: serialize and deserialize
        let serialized = serde_yaml::to_string(&policy).unwrap();
        let roundtripped: PolicySet = serde_yaml::from_str(&serialized).unwrap();
        assert_eq!(
            roundtripped.classifications.get("getPortfolio"),
            Some(&DataClassification::Confidential)
        );
        assert_eq!(
            roundtripped.defaults.default_classification,
            DataClassification::Confidential
        );
    }

    #[test]
    fn test_backward_compat_no_classifications() {
        // Old policy without classifications or default_classification fields
        let yaml = r#"
version: "1.0"
defaults:
  allow_methods: ["GET"]
  deny_external_refs: true
  require_auth: true
  audit_level: "basic"
rules: []
"#;
        let policy: PolicySet = serde_yaml::from_str(yaml).unwrap();
        assert!(policy.classifications.is_empty());
        assert_eq!(
            policy.defaults.default_classification,
            DataClassification::Internal
        );
        assert!(policy.credential_scopes.is_empty());
    }

    #[test]
    fn test_credential_scope_serde() {
        let yaml = r#"
version: "1.0"
defaults:
  allow_methods: ["GET"]
  require_auth: false
  audit_level: "basic"
rules: []
credential_scopes:
  - profile: "agent-readonly"
    allowed_methods: ["GET"]
    allowed_tags: ["portfolio", "stocks"]
  - profile: "agent-full"
    allowed_methods: ["GET", "POST"]
    requires_approval: ["DELETE", "PATCH"]
"#;
        let policy: PolicySet = serde_yaml::from_str(yaml).unwrap();
        assert_eq!(policy.credential_scopes.len(), 2);

        let readonly = &policy.credential_scopes[0];
        assert_eq!(readonly.profile, "agent-readonly");
        assert_eq!(readonly.allowed_methods, vec!["GET"]);
        assert_eq!(readonly.allowed_tags, vec!["portfolio", "stocks"]);
        assert!(readonly.requires_approval.is_empty());

        let full = &policy.credential_scopes[1];
        assert_eq!(full.profile, "agent-full");
        assert_eq!(full.allowed_methods, vec!["GET", "POST"]);
        assert!(full.allowed_tags.is_empty());
        assert_eq!(full.requires_approval, vec!["DELETE", "PATCH"]);

        // Round-trip
        let serialized = serde_yaml::to_string(&policy).unwrap();
        let roundtripped: PolicySet = serde_yaml::from_str(&serialized).unwrap();
        assert_eq!(roundtripped.credential_scopes.len(), 2);
        assert_eq!(roundtripped.credential_scopes[0].profile, "agent-readonly");
        assert_eq!(
            roundtripped.credential_scopes[1].requires_approval,
            vec!["DELETE", "PATCH"]
        );
    }
}