pmcp_code_mode/policy/

types.rs

//! Domain types for policy evaluation.
//!
//! These types represent the entities used in Cedar policy evaluation.
//! They are pure domain types with no AWS SDK dependency.

use crate::graphql::GraphQLQueryInfo;
use std::collections::HashSet;

/// Server configuration for policy evaluation.
///
/// Uses unified attribute names that match the Cedar schema:
/// - `allow_write`, `allow_delete`, `allow_admin` (unified action flags)
/// - `blocked_operations`, `allowed_operations` (unified operation lists)
#[derive(Debug, Clone)]
pub struct ServerConfigEntity {
    /// Server ID
    pub server_id: String,

    /// Server type (e.g., "graphql")
    pub server_type: String,

    /// Whether write operations (mutations) are allowed
    pub allow_write: bool,

    /// Whether delete operations are allowed
    pub allow_delete: bool,

    /// Whether admin operations (introspection) are allowed
    pub allow_admin: bool,

    /// Allowed operation names (allowlist mode)
    pub allowed_operations: HashSet<String>,

    /// Blocked operation names (blocklist mode)
    pub blocked_operations: HashSet<String>,

    /// Maximum query depth
    pub max_depth: u32,

    /// Maximum field count
    pub max_field_count: u32,

    /// Maximum estimated cost
    pub max_cost: u32,

    /// Maximum API calls (for compatibility with unified schema)
    pub max_api_calls: u32,

    /// Fields that should be blocked
    pub blocked_fields: HashSet<String>,

    /// Allowed sensitive data categories
    pub allowed_sensitive_categories: HashSet<String>,
}

impl Default for ServerConfigEntity {
    fn default() -> Self {
        Self {
            server_id: "unknown".to_string(),
            server_type: "graphql".to_string(),
            allow_write: false,
            allow_delete: false,
            allow_admin: false,
            allowed_operations: HashSet::new(),
            blocked_operations: HashSet::new(),
            max_depth: 10,
            max_field_count: 100,
            max_cost: 1000,
            max_api_calls: 50,
            blocked_fields: HashSet::new(),
            allowed_sensitive_categories: HashSet::new(),
        }
    }
}

/// Operation entity for policy evaluation.
#[derive(Debug, Clone)]
pub struct OperationEntity {
    /// Unique ID for this operation
    pub id: String,

    /// Operation type: "query", "mutation", or "subscription"
    pub operation_type: String,

    /// Operation name (if provided)
    pub operation_name: String,

    /// Root fields accessed
    pub root_fields: HashSet<String>,

    /// Types accessed
    pub accessed_types: HashSet<String>,

    /// Fields accessed (Type.field format)
    pub accessed_fields: HashSet<String>,

    /// Query nesting depth
    pub depth: u32,

    /// Total field count
    pub field_count: u32,

    /// Estimated query cost
    pub estimated_cost: u32,

    /// Whether introspection is used
    pub has_introspection: bool,

    /// Whether sensitive data is accessed
    pub accesses_sensitive_data: bool,

    /// Sensitive data categories accessed
    pub sensitive_categories: HashSet<String>,
}

impl OperationEntity {
    /// Create from GraphQL query info.
    pub fn from_query_info(query_info: &GraphQLQueryInfo) -> Self {
        use crate::graphql::GraphQLOperationType;

        let operation_type = match query_info.operation_type {
            GraphQLOperationType::Query => "query",
            GraphQLOperationType::Mutation => "mutation",
            GraphQLOperationType::Subscription => "subscription",
        };

        Self {
            id: query_info
                .operation_name
                .clone()
                .unwrap_or_else(|| "anonymous".to_string()),
            operation_type: operation_type.to_string(),
            operation_name: query_info.operation_name.clone().unwrap_or_default(),
            root_fields: query_info.root_fields.iter().cloned().collect(),
            accessed_types: query_info.types_accessed.iter().cloned().collect(),
            accessed_fields: query_info.fields_accessed.iter().cloned().collect(),
            depth: query_info.max_depth as u32,
            field_count: query_info.fields_accessed.len() as u32,
            estimated_cost: query_info.fields_accessed.len() as u32,
            has_introspection: query_info.has_introspection,
            accesses_sensitive_data: false,
            sensitive_categories: HashSet::new(),
        }
    }
}

/// Authorization decision from policy evaluation.
#[derive(Debug, Clone)]
pub struct AuthorizationDecision {
    /// Whether the operation is allowed
    pub allowed: bool,

    /// Policy IDs that determined the decision
    pub determining_policies: Vec<String>,

    /// Error messages (if any)
    pub errors: Vec<String>,
}

/// Script entity for policy evaluation (OpenAPI Code Mode).
///
/// Unlike GraphQL's single Operation entity, OpenAPI Code Mode validates
/// JavaScript scripts that can contain multiple API calls with loops and logic.
#[cfg(feature = "openapi-code-mode")]
#[derive(Debug, Clone)]
pub struct ScriptEntity {
    /// Unique ID for this script validation
    pub id: String,

    /// Script type: "read_only", "mixed", or "write_only"
    pub script_type: String,

    /// Whether script contains any write operations (POST/PUT/PATCH/DELETE)
    pub has_writes: bool,

    /// Whether script contains DELETE operations
    pub has_deletes: bool,

    /// Total number of API calls in the script
    pub total_api_calls: u32,

    /// Number of GET calls
    pub read_calls: u32,

    /// Number of POST/PUT/PATCH calls
    pub write_calls: u32,

    /// Number of DELETE calls
    pub delete_calls: u32,

    /// Set of all paths accessed
    pub accessed_paths: HashSet<String>,

    /// Set of all HTTP methods used
    pub accessed_methods: HashSet<String>,

    /// Normalized path patterns (IDs replaced with *)
    pub path_patterns: HashSet<String>,

    /// Called operations in "METHOD:pathPattern" format for allowlist/blocklist matching
    pub called_operations: HashSet<String>,

    /// Maximum loop iterations (from .slice() bounds)
    pub loop_iterations: u32,

    /// Maximum nesting depth in the AST
    pub nesting_depth: u32,

    /// Script length in characters
    pub script_length: u32,

    /// Whether script accesses sensitive paths (/admin, /internal, etc.)
    pub accesses_sensitive_path: bool,

    /// Whether script has an unbounded loop
    pub has_unbounded_loop: bool,

    /// Whether script uses dynamic path interpolation
    pub has_dynamic_path: bool,

    /// Whether script has a @returns output declaration
    pub has_output_declaration: bool,

    /// Fields declared in the @returns annotation
    pub output_fields: HashSet<String>,

    /// Whether script uses spread operators in output (potential field leakage)
    pub has_spread_in_output: bool,
}

#[cfg(feature = "openapi-code-mode")]
impl ScriptEntity {
    /// Build from JavaScript code analysis.
    pub fn from_javascript_info(
        info: &crate::javascript::JavaScriptCodeInfo,
        sensitive_patterns: &[String],
    ) -> Self {
        use crate::javascript::HttpMethod;

        let mut accessed_paths = HashSet::new();
        let mut accessed_methods = HashSet::new();
        let mut path_patterns = HashSet::new();
        let mut called_operations = HashSet::new();
        let mut read_calls = 0u32;
        let mut write_calls = 0u32;
        let mut delete_calls = 0u32;
        let mut has_dynamic_path = false;
        let mut accesses_sensitive_path = false;

        for api_call in &info.api_calls {
            accessed_paths.insert(api_call.path.clone());
            let method_str = format!("{:?}", api_call.method).to_uppercase();
            accessed_methods.insert(method_str.clone());

            // Normalize path to pattern
            let pattern = normalize_path_to_pattern(&api_call.path);
            path_patterns.insert(pattern.clone());

            // Build called operation string
            called_operations.insert(format!("{}:{}", method_str, pattern));

            // Count by method type
            match api_call.method {
                HttpMethod::Get | HttpMethod::Head | HttpMethod::Options => read_calls += 1,
                HttpMethod::Delete => delete_calls += 1,
                _ => write_calls += 1,
            }

            // Track dynamic paths
            if api_call.is_dynamic_path {
                has_dynamic_path = true;
            }

            // Check for sensitive path access
            let path_lower = api_call.path.to_lowercase();
            for pattern in sensitive_patterns {
                if path_lower.contains(&pattern.to_lowercase()) {
                    accesses_sensitive_path = true;
                    break;
                }
            }
        }

        // Determine script type
        let has_writes = write_calls > 0 || delete_calls > 0;
        let has_reads = read_calls > 0;
        let script_type = match (has_reads, has_writes) {
            (true, false) => "read_only",
            (false, true) => "write_only",
            (true, true) => "mixed",
            (false, false) => "empty",
        };

        Self {
            id: info
                .api_calls
                .first()
                .map(|c| format!("{}:{}", format!("{:?}", c.method).to_uppercase(), c.path))
                .unwrap_or_else(|| "script".to_string()),
            script_type: script_type.to_string(),
            has_writes,
            has_deletes: delete_calls > 0,
            total_api_calls: info.api_calls.len() as u32,
            read_calls,
            write_calls,
            delete_calls,
            accessed_paths,
            accessed_methods,
            path_patterns,
            called_operations,
            loop_iterations: 0,
            nesting_depth: info.max_depth as u32,
            script_length: 0,
            accesses_sensitive_path,
            has_unbounded_loop: !info.all_loops_bounded && info.loop_count > 0,
            has_dynamic_path,
            has_output_declaration: info.output_declaration.has_declaration,
            output_fields: info.output_declaration.declared_fields.clone(),
            has_spread_in_output: info.output_declaration.has_spread_risk
                || info.has_output_spread_risk,
        }
    }

    /// Get the policy action for this script using unified action model.
    pub fn action(&self) -> &'static str {
        match self.script_type.as_str() {
            "read_only" | "empty" => "Read",
            "write_only" | "mixed" => {
                if self.has_deletes {
                    "Delete"
                } else {
                    "Write"
                }
            }
            _ => "Read",
        }
    }
}

/// Normalize a path to a pattern by replacing numeric/UUID segments with *.
#[cfg(feature = "openapi-code-mode")]
pub fn normalize_path_to_pattern(path: &str) -> String {
    path.split('/')
        .map(|segment| {
            if segment.chars().all(|c| c.is_ascii_digit())
                || (segment.len() == 36 && segment.contains('-'))
            {
                "*"
            } else {
                segment
            }
        })
        .collect::<Vec<_>>()
        .join("/")
}

/// Normalize an operation string to the canonical "METHOD:/path" format.
#[cfg(feature = "openapi-code-mode")]
pub fn normalize_operation_format(op: &str) -> String {
    let trimmed = op.trim();

    let (method, path) = if let Some(idx) = trimmed.find(':') {
        let potential_path = trimmed[idx + 1..].trim();
        if potential_path.starts_with('/') {
            let method = trimmed[..idx].trim();
            (method, potential_path)
        } else {
            return trimmed.to_string();
        }
    } else if let Some(idx) = trimmed.find(' ') {
        let method = trimmed[..idx].trim();
        let path = trimmed[idx + 1..].trim();
        (method, path)
    } else {
        return trimmed.to_string();
    };

    let method_upper = method.to_uppercase();

    let normalized_path = path
        .split('/')
        .map(|segment| {
            if segment.starts_with('{') && segment.ends_with('}') {
                "*"
            } else if segment.starts_with(':') {
                "*"
            } else if segment.chars().all(|c| c.is_ascii_digit()) {
                "*"
            } else if segment.len() == 36 && segment.contains('-') {
                "*"
            } else {
                segment
            }
        })
        .collect::<Vec<_>>()
        .join("/");

    format!("{}:{}", method_upper, normalized_path)
}

/// Server configuration for OpenAPI Code Mode.
#[cfg(feature = "openapi-code-mode")]
#[derive(Debug, Clone)]
pub struct OpenAPIServerEntity {
    pub server_id: String,
    pub server_type: String,

    // Unified action flags
    pub allow_write: bool,
    pub allow_delete: bool,
    pub allow_admin: bool,

    // Write mode: "allow_all", "deny_all", "allowlist", "blocklist"
    pub write_mode: String,

    // Unified limits
    pub max_depth: u32,
    pub max_cost: u32,
    pub max_api_calls: u32,

    // OpenAPI-specific limits
    pub max_loop_iterations: u32,
    pub max_script_length: u32,
    pub max_nesting_depth: u32,
    pub execution_timeout_seconds: u32,

    // Unified operation lists
    pub allowed_operations: HashSet<String>,
    pub blocked_operations: HashSet<String>,

    // OpenAPI-specific method controls
    pub allowed_methods: HashSet<String>,
    pub blocked_methods: HashSet<String>,
    pub allowed_path_patterns: HashSet<String>,
    pub blocked_path_patterns: HashSet<String>,
    pub sensitive_path_patterns: HashSet<String>,

    // Auto-approval settings
    pub auto_approve_read_only: bool,
    pub max_api_calls_for_auto_approve: u32,

    // Field control (two-tier blocklist)
    pub internal_blocked_fields: HashSet<String>,
    pub output_blocked_fields: HashSet<String>,
    pub require_output_declaration: bool,
}

#[cfg(feature = "openapi-code-mode")]
impl Default for OpenAPIServerEntity {
    fn default() -> Self {
        Self {
            server_id: "unknown".to_string(),
            server_type: "openapi".to_string(),
            allow_write: false,
            allow_delete: false,
            allow_admin: false,
            write_mode: "deny_all".to_string(),
            max_depth: 10,
            max_cost: 1000,
            max_api_calls: 50,
            max_loop_iterations: 100,
            max_script_length: 10000,
            max_nesting_depth: 10,
            execution_timeout_seconds: 30,
            allowed_operations: HashSet::new(),
            blocked_operations: HashSet::new(),
            allowed_methods: HashSet::new(),
            blocked_methods: HashSet::new(),
            allowed_path_patterns: HashSet::new(),
            blocked_path_patterns: ["/admin".into(), "/internal".into()].into_iter().collect(),
            sensitive_path_patterns: ["/admin".into(), "/internal".into(), "/debug".into()]
                .into_iter()
                .collect(),
            auto_approve_read_only: true,
            max_api_calls_for_auto_approve: 10,
            internal_blocked_fields: HashSet::new(),
            output_blocked_fields: HashSet::new(),
            require_output_declaration: false,
        }
    }
}

/// Get the Cedar schema in JSON format.
///
/// Uses unified action model with Read/Write/Delete/Admin actions.
pub fn get_code_mode_schema_json() -> serde_json::Value {
    let applies_to = serde_json::json!({
        "principalTypes": ["Operation"],
        "resourceTypes": ["Server"],
        "context": {
            "type": "Record",
            "attributes": {
                "serverId": { "type": "String", "required": true },
                "serverType": { "type": "String", "required": true },
                "userId": { "type": "String", "required": false },
                "sessionId": { "type": "String", "required": false }
            }
        }
    });

    serde_json::json!({
        "CodeMode": {
            "entityTypes": {
                "Operation": {
                    "shape": {
                        "type": "Record",
                        "attributes": {
                            "operationType": { "type": "String", "required": true },
                            "operationName": { "type": "String", "required": true },
                            "rootFields": { "type": "Set", "element": { "type": "String" } },
                            "accessedTypes": { "type": "Set", "element": { "type": "String" } },
                            "accessedFields": { "type": "Set", "element": { "type": "String" } },
                            "depth": { "type": "Long", "required": true },
                            "fieldCount": { "type": "Long", "required": true },
                            "estimatedCost": { "type": "Long", "required": true },
                            "hasIntrospection": { "type": "Boolean", "required": true },
                            "accessesSensitiveData": { "type": "Boolean", "required": true },
                            "sensitiveCategories": { "type": "Set", "element": { "type": "String" } }
                        }
                    }
                },
                "Server": {
                    "shape": {
                        "type": "Record",
                        "attributes": {
                            "serverId": { "type": "String", "required": true },
                            "serverType": { "type": "String", "required": true },
                            "maxDepth": { "type": "Long", "required": true },
                            "maxCost": { "type": "Long", "required": true },
                            "maxApiCalls": { "type": "Long", "required": true },
                            "allowWrite": { "type": "Boolean", "required": true },
                            "allowDelete": { "type": "Boolean", "required": true },
                            "allowAdmin": { "type": "Boolean", "required": true },
                            "blockedOperations": { "type": "Set", "element": { "type": "String" } },
                            "allowedOperations": { "type": "Set", "element": { "type": "String" } },
                            "blockedFields": { "type": "Set", "element": { "type": "String" } }
                        }
                    }
                }
            },
            "actions": {
                "Read": { "appliesTo": applies_to },
                "Write": { "appliesTo": applies_to },
                "Delete": { "appliesTo": applies_to },
                "Admin": { "appliesTo": applies_to }
            }
        }
    })
}

/// Get the Cedar schema for OpenAPI Code Mode in JSON format.
#[cfg(feature = "openapi-code-mode")]
pub fn get_openapi_code_mode_schema_json() -> serde_json::Value {
    let applies_to = serde_json::json!({
        "principalTypes": ["Script"],
        "resourceTypes": ["Server"],
        "context": {
            "type": "Record",
            "attributes": {
                "serverId": { "type": "String", "required": true },
                "serverType": { "type": "String", "required": true },
                "userId": { "type": "String", "required": false },
                "sessionId": { "type": "String", "required": false }
            }
        }
    });

    serde_json::json!({
        "CodeMode": {
            "entityTypes": {
                "Script": {
                    "shape": {
                        "type": "Record",
                        "attributes": {
                            "scriptType": { "type": "String", "required": true },
                            "hasWrites": { "type": "Boolean", "required": true },
                            "hasDeletes": { "type": "Boolean", "required": true },
                            "totalApiCalls": { "type": "Long", "required": true },
                            "readCalls": { "type": "Long", "required": true },
                            "writeCalls": { "type": "Long", "required": true },
                            "deleteCalls": { "type": "Long", "required": true },
                            "accessedPaths": { "type": "Set", "element": { "type": "String" } },
                            "accessedMethods": { "type": "Set", "element": { "type": "String" } },
                            "pathPatterns": { "type": "Set", "element": { "type": "String" } },
                            "calledOperations": { "type": "Set", "element": { "type": "String" } },
                            "loopIterations": { "type": "Long", "required": true },
                            "nestingDepth": { "type": "Long", "required": true },
                            "scriptLength": { "type": "Long", "required": true },
                            "accessesSensitivePath": { "type": "Boolean", "required": true },
                            "hasUnboundedLoop": { "type": "Boolean", "required": true },
                            "hasDynamicPath": { "type": "Boolean", "required": true },
                            "outputFields": { "type": "Set", "element": { "type": "String" } },
                            "hasOutputDeclaration": { "type": "Boolean", "required": true },
                            "hasSpreadInOutput": { "type": "Boolean", "required": true }
                        }
                    }
                },
                "Server": {
                    "shape": {
                        "type": "Record",
                        "attributes": {
                            "serverId": { "type": "String", "required": true },
                            "serverType": { "type": "String", "required": true },
                            "maxDepth": { "type": "Long", "required": true },
                            "maxCost": { "type": "Long", "required": true },
                            "maxApiCalls": { "type": "Long", "required": true },
                            "allowWrite": { "type": "Boolean", "required": true },
                            "allowDelete": { "type": "Boolean", "required": true },
                            "allowAdmin": { "type": "Boolean", "required": true },
                            "blockedOperations": { "type": "Set", "element": { "type": "String" } },
                            "allowedOperations": { "type": "Set", "element": { "type": "String" } },
                            "blockedFields": { "type": "Set", "element": { "type": "String" } },
                            "maxLoopIterations": { "type": "Long", "required": true },
                            "maxScriptLength": { "type": "Long", "required": true },
                            "maxNestingDepth": { "type": "Long", "required": true },
                            "executionTimeoutSeconds": { "type": "Long", "required": true },
                            "allowedMethods": { "type": "Set", "element": { "type": "String" } },
                            "blockedMethods": { "type": "Set", "element": { "type": "String" } },
                            "allowedPathPatterns": { "type": "Set", "element": { "type": "String" } },
                            "blockedPathPatterns": { "type": "Set", "element": { "type": "String" } },
                            "sensitivePathPatterns": { "type": "Set", "element": { "type": "String" } },
                            "autoApproveReadOnly": { "type": "Boolean", "required": true },
                            "maxApiCallsForAutoApprove": { "type": "Long", "required": true },
                            "internalBlockedFields": { "type": "Set", "element": { "type": "String" } },
                            "outputBlockedFields": { "type": "Set", "element": { "type": "String" } },
                            "requireOutputDeclaration": { "type": "Boolean", "required": true }
                        }
                    }
                }
            },
            "actions": {
                "Read": { "appliesTo": applies_to },
                "Write": { "appliesTo": applies_to },
                "Delete": { "appliesTo": applies_to },
                "Admin": { "appliesTo": applies_to }
            }
        }
    })
}

/// Get baseline Cedar policies for OpenAPI Code Mode.
#[cfg(feature = "openapi-code-mode")]
pub fn get_openapi_baseline_policies() -> Vec<(&'static str, &'static str, &'static str)> {
    vec![
        (
            "permit_reads",
            "Permit all read operations (GET scripts)",
            r#"permit(principal, action == CodeMode::Action::"Read", resource);"#,
        ),
        (
            "permit_writes",
            "Permit write operations (when enabled)",
            r#"permit(principal, action == CodeMode::Action::"Write", resource) when { resource.allowWrite == true };"#,
        ),
        (
            "permit_deletes",
            "Permit delete operations (when enabled)",
            r#"permit(principal, action == CodeMode::Action::"Delete", resource) when { resource.allowDelete == true };"#,
        ),
        (
            "forbid_sensitive_paths",
            "Block scripts accessing sensitive paths",
            r#"forbid(principal, action, resource) when { principal.accessesSensitivePath == true };"#,
        ),
        (
            "forbid_unbounded_loops",
            "Block scripts with unbounded loops",
            r#"forbid(principal, action, resource) when { principal.hasUnboundedLoop == true };"#,
        ),
        (
            "forbid_excessive_api_calls",
            "Enforce API call limit",
            r#"forbid(principal, action, resource) when { principal.totalApiCalls > resource.maxApiCalls };"#,
        ),
        (
            "forbid_excessive_nesting",
            "Enforce nesting depth limit",
            r#"forbid(principal, action, resource) when { principal.nestingDepth > resource.maxNestingDepth };"#,
        ),
        (
            "forbid_output_blocked_fields",
            "Block scripts that return output-blocked fields",
            r#"forbid(principal, action, resource) when { principal.outputFields.containsAny(resource.outputBlockedFields) };"#,
        ),
        (
            "forbid_spread_without_declaration",
            "Block scripts with spread in output when output declaration is required",
            r#"forbid(principal, action, resource) when { principal.hasSpreadInOutput == true && resource.requireOutputDeclaration == true };"#,
        ),
        (
            "forbid_missing_output_declaration",
            "Block scripts without output declaration when required",
            r#"forbid(principal, action, resource) when { principal.hasOutputDeclaration == false && resource.requireOutputDeclaration == true };"#,
        ),
    ]
}

/// Get the baseline Cedar policies.
pub fn get_baseline_policies() -> Vec<(&'static str, &'static str, &'static str)> {
    vec![
        (
            "permit_reads",
            "Permit all read operations (queries)",
            r#"permit(principal, action == CodeMode::Action::"Read", resource);"#,
        ),
        (
            "permit_writes",
            "Permit write operations (when enabled)",
            r#"permit(principal, action == CodeMode::Action::"Write", resource) when { resource.allowWrite == true };"#,
        ),
        (
            "permit_deletes",
            "Permit delete operations (when enabled)",
            r#"permit(principal, action == CodeMode::Action::"Delete", resource) when { resource.allowDelete == true };"#,
        ),
        (
            "permit_admin",
            "Permit admin operations (when enabled)",
            r#"permit(principal, action == CodeMode::Action::"Admin", resource) when { resource.allowAdmin == true };"#,
        ),
        (
            "forbid_blocked_operations",
            "Block operations in blocklist",
            r#"forbid(principal, action, resource) when { resource.blockedOperations.contains(principal.operationName) };"#,
        ),
        (
            "forbid_blocked_fields",
            "Block access to blocked fields",
            r#"forbid(principal, action, resource) when { resource.blockedFields.containsAny(principal.accessedFields) };"#,
        ),
        (
            "forbid_excessive_depth",
            "Enforce maximum query depth",
            r#"forbid(principal, action, resource) when { principal.depth > resource.maxDepth };"#,
        ),
        (
            "forbid_excessive_cost",
            "Enforce maximum query cost",
            r#"forbid(principal, action, resource) when { principal.estimatedCost > resource.maxCost };"#,
        ),
    ]
}