agent_os_kernel 3.1.0__py3-none-any.whl

caas/gateway/trust_gateway.py

@@ -0,0 +1,474 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""
+Trust Gateway: The Middleware Gap Solution
+
+The Naive Approach:
+"Let's use a startup's API that auto-routes our traffic to the cheapest model."
+
+The Engineering Reality:
+No Enterprise CISO will send their proprietary data to a random middleware startup 
+just to save 30% on tokens. The risk of data leakage is too high.
+
+This layer—the "Model Gateway"—is critical, but it requires massive trust.
+
+The Opportunity:
+There is a gap here, but it's not for a SaaS. It's for Infrastructure.
+
+The Solution:
+Build an On-Prem / Private Cloud Router that enterprises can deploy within their
+own infrastructure. The winner won't be the one with the smartest routing algorithm;
+it will be the one the Enterprise trusts with the keys to the kingdom.
+"""
+
+from enum import Enum
+from typing import Dict, List, Optional, Any
+from datetime import datetime
+from pydantic import BaseModel, Field
+import uuid
+import json
+
+
+class DeploymentMode(str, Enum):
+    """Deployment modes for Trust Gateway."""
+    ON_PREM = "on_prem"  # Deployed on customer's own infrastructure
+    PRIVATE_CLOUD = "private_cloud"  # Deployed in customer's private cloud (AWS VPC, Azure VNet, GCP VPC)
+    HYBRID = "hybrid"  # Hybrid deployment with local processing and cloud backup
+    AIR_GAPPED = "air_gapped"  # Completely isolated from internet (maximum security)
+
+
+class SecurityLevel(str, Enum):
+    """Security levels for data handling."""
+    STANDARD = "standard"  # Basic security controls
+    HIGH = "high"  # Enhanced security (encryption at rest and in transit)
+    MAXIMUM = "maximum"  # Maximum security (air-gapped, zero data retention)
+
+
+class DataRetentionPolicy(BaseModel):
+    """Data retention and deletion policies."""
+    retain_requests: bool = Field(default=False, description="Whether to retain request data")
+    retention_days: int = Field(default=0, ge=0, le=365, description="Days to retain data (0 = no retention)")
+    auto_delete: bool = Field(default=True, description="Automatically delete data after retention period")
+    encrypt_at_rest: bool = Field(default=True, description="Encrypt data at rest")
+    pii_scrubbing: bool = Field(default=True, description="Automatically scrub PII from logs")
+
+
+class AuditLog(BaseModel):
+    """Audit log entry for compliance and security monitoring."""
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    timestamp: str = Field(default_factory=lambda: datetime.utcnow().isoformat())
+    event_type: str  # e.g., "request_routed", "data_accessed", "policy_changed"
+    user_id: Optional[str] = None
+    request_id: Optional[str] = None
+    action: str  # Description of the action
+    model_tier: Optional[str] = None
+    data_classification: Optional[str] = None  # e.g., "public", "confidential", "secret"
+    security_level: Optional[str] = None
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for storage."""
+        return {
+            "id": self.id,
+            "timestamp": self.timestamp,
+            "event_type": self.event_type,
+            "user_id": self.user_id,
+            "request_id": self.request_id,
+            "action": self.action,
+            "model_tier": self.model_tier,
+            "data_classification": self.data_classification,
+            "security_level": self.security_level,
+            "metadata": self.metadata
+        }
+
+
+class SecurityPolicy(BaseModel):
+    """Security policy configuration for Trust Gateway."""
+    deployment_mode: DeploymentMode = Field(default=DeploymentMode.ON_PREM)
+    security_level: SecurityLevel = Field(default=SecurityLevel.HIGH)
+    data_retention: DataRetentionPolicy = Field(default_factory=DataRetentionPolicy)
+    
+    # Authentication & Authorization
+    require_authentication: bool = Field(default=True, description="Require authentication for all requests")
+    allowed_users: List[str] = Field(default_factory=list, description="List of allowed user IDs (empty = all)")
+    allowed_ip_ranges: List[str] = Field(default_factory=list, description="Allowed IP ranges (CIDR notation)")
+    
+    # Data Classification
+    data_classification_required: bool = Field(default=False, description="Require data classification labels")
+    allowed_classifications: List[str] = Field(
+        default_factory=lambda: ["public", "internal", "confidential", "secret"],
+        description="Allowed data classification levels"
+    )
+    
+    # Encryption
+    encrypt_in_transit: bool = Field(default=True, description="Require TLS/HTTPS for all communication")
+    encrypt_at_rest: bool = Field(default=True, description="Encrypt stored data")
+    
+    # Audit & Compliance
+    audit_all_requests: bool = Field(default=True, description="Audit all gateway requests")
+    audit_data_access: bool = Field(default=True, description="Audit all data access events")
+    compliance_mode: Optional[str] = Field(default=None, description="Compliance framework (e.g., 'GDPR', 'HIPAA', 'SOC2')")
+    
+    # Network Isolation
+    allow_external_calls: bool = Field(default=False, description="Allow calls to external APIs")
+    external_model_endpoints: List[str] = Field(
+        default_factory=list,
+        description="Whitelisted external model endpoints (if allowed)"
+    )
+
+
+class TrustGateway:
+    """
+    Trust Gateway: Enterprise-Grade Private Cloud Router
+    
+    The Trust Gateway is designed to address enterprise security concerns by providing:
+    1. On-Prem / Private Cloud deployment options
+    2. Zero data leakage (data never leaves customer infrastructure)
+    3. Full audit trail for compliance
+    4. Configurable security policies
+    5. Data retention and deletion controls
+    
+    Philosophy:
+    The winner isn't the one with the smartest routing algorithm;
+    it's the one the Enterprise trusts with the keys to the kingdom.
+    """
+    
+    def __init__(
+        self,
+        security_policy: Optional[SecurityPolicy] = None,
+        audit_enabled: bool = True
+    ):
+        """
+        Initialize Trust Gateway.
+        
+        Args:
+            security_policy: Security policy configuration
+            audit_enabled: Whether to enable audit logging
+        """
+        self.security_policy = security_policy or SecurityPolicy()
+        self.audit_enabled = audit_enabled
+        self.audit_logs: List[AuditLog] = []
+        self._deployment_info = self._get_deployment_info()
+    
+    def _get_deployment_info(self) -> Dict[str, Any]:
+        """Get deployment information."""
+        return {
+            "deployment_mode": self.security_policy.deployment_mode,
+            "security_level": self.security_policy.security_level,
+            "data_retention_days": self.security_policy.data_retention.retention_days,
+            "audit_enabled": self.audit_enabled,
+            "compliance_mode": self.security_policy.compliance_mode,
+            "timestamp": datetime.utcnow().isoformat()
+        }
+    
+    def validate_request(
+        self,
+        request_data: Dict[str, Any],
+        user_id: Optional[str] = None,
+        ip_address: Optional[str] = None,
+        data_classification: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Validate request against security policy.
+        
+        Args:
+            request_data: The request data to validate
+            user_id: User ID making the request
+            ip_address: IP address of the requester
+            data_classification: Classification level of the data
+        
+        Returns:
+            Validation result with status and any security warnings
+        """
+        validation_result = {
+            "valid": True,
+            "warnings": [],
+            "violations": []
+        }
+        
+        # Check authentication requirement
+        if self.security_policy.require_authentication and not user_id:
+            validation_result["valid"] = False
+            validation_result["violations"].append("Authentication required but no user_id provided")
+        
+        # Check allowed users
+        if user_id and self.security_policy.allowed_users:
+            if user_id not in self.security_policy.allowed_users:
+                validation_result["valid"] = False
+                validation_result["violations"].append(f"User {user_id} not in allowed users list")
+        
+        # Check data classification
+        if self.security_policy.data_classification_required and not data_classification:
+            validation_result["valid"] = False
+            validation_result["violations"].append("Data classification required but not provided")
+        
+        if data_classification and data_classification not in self.security_policy.allowed_classifications:
+            validation_result["valid"] = False
+            validation_result["violations"].append(
+                f"Data classification '{data_classification}' not in allowed classifications"
+            )
+        
+        # Check encryption requirements
+        if self.security_policy.encrypt_in_transit:
+            validation_result["warnings"].append("Ensure connection uses TLS/HTTPS")
+        
+        # Audit the validation attempt
+        if self.audit_enabled:
+            self._log_audit(
+                event_type="request_validated",
+                user_id=user_id,
+                action=f"Request validation: {'passed' if validation_result['valid'] else 'failed'}",
+                data_classification=data_classification,
+                metadata={
+                    "ip_address": ip_address,
+                    "violations": validation_result["violations"],
+                    "warnings": validation_result["warnings"]
+                }
+            )
+        
+        return validation_result
+    
+    def route_request(
+        self,
+        query: str,
+        user_id: Optional[str] = None,
+        data_classification: Optional[str] = None,
+        request_metadata: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """
+        Route request through Trust Gateway with security controls.
+        
+        Args:
+            query: The user query to route
+            user_id: User ID making the request
+            data_classification: Classification level of the data
+            request_metadata: Additional request metadata
+        
+        Returns:
+            Routing decision with security context
+        """
+        request_id = str(uuid.uuid4())
+        
+        # Validate request first
+        validation = self.validate_request(
+            request_data={"query": query},
+            user_id=user_id,
+            data_classification=data_classification
+        )
+        
+        if not validation["valid"]:
+            return {
+                "status": "rejected",
+                "request_id": request_id,
+                "reason": "Security policy violation",
+                "violations": validation["violations"],
+                "timestamp": datetime.utcnow().isoformat()
+            }
+        
+        # Import here to avoid circular dependency
+        from caas.routing import HeuristicRouter
+        
+        # Use heuristic router for actual routing decision
+        router = HeuristicRouter()
+        decision = router.route(query)
+        
+        # Add security context to routing decision
+        result = {
+            "status": "approved",
+            "request_id": request_id,
+            "routing_decision": decision.model_dump(),
+            "security_context": {
+                "deployment_mode": self.security_policy.deployment_mode,
+                "security_level": self.security_policy.security_level,
+                "data_classification": data_classification,
+                "data_retention_days": self.security_policy.data_retention.retention_days,
+                "encrypted": self.security_policy.encrypt_at_rest,
+                "audited": self.audit_enabled
+            },
+            "warnings": validation["warnings"],
+            "timestamp": datetime.utcnow().isoformat()
+        }
+        
+        # Audit the routing decision
+        if self.audit_enabled:
+            self._log_audit(
+                event_type="request_routed",
+                user_id=user_id,
+                request_id=request_id,
+                action=f"Query routed to {decision.model_tier} tier",
+                model_tier=decision.model_tier,
+                data_classification=data_classification,
+                metadata={
+                    "query_length": len(query),
+                    "suggested_model": decision.suggested_model,
+                    "estimated_cost": decision.estimated_cost,
+                    "request_metadata": request_metadata
+                }
+            )
+        
+        return result
+    
+    def _log_audit(
+        self,
+        event_type: str,
+        action: str,
+        user_id: Optional[str] = None,
+        request_id: Optional[str] = None,
+        model_tier: Optional[str] = None,
+        data_classification: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None
+    ):
+        """Log an audit event."""
+        log_entry = AuditLog(
+            event_type=event_type,
+            user_id=user_id,
+            request_id=request_id,
+            action=action,
+            model_tier=model_tier,
+            data_classification=data_classification,
+            security_level=self.security_policy.security_level,
+            metadata=metadata or {}
+        )
+        self.audit_logs.append(log_entry)
+    
+    def get_audit_logs(
+        self,
+        event_type: Optional[str] = None,
+        user_id: Optional[str] = None,
+        start_time: Optional[str] = None,
+        end_time: Optional[str] = None
+    ) -> List[Dict[str, Any]]:
+        """
+        Retrieve audit logs with optional filtering.
+        
+        Args:
+            event_type: Filter by event type
+            user_id: Filter by user ID
+            start_time: Filter logs after this time (ISO format)
+            end_time: Filter logs before this time (ISO format)
+        
+        Returns:
+            List of matching audit log entries
+        """
+        filtered_logs = self.audit_logs
+        
+        if event_type:
+            filtered_logs = [log for log in filtered_logs if log.event_type == event_type]
+        
+        if user_id:
+            filtered_logs = [log for log in filtered_logs if log.user_id == user_id]
+        
+        if start_time:
+            filtered_logs = [log for log in filtered_logs if log.timestamp >= start_time]
+        
+        if end_time:
+            filtered_logs = [log for log in filtered_logs if log.timestamp <= end_time]
+        
+        return [log.to_dict() for log in filtered_logs]
+    
+    def get_deployment_info(self) -> Dict[str, Any]:
+        """
+        Get Trust Gateway deployment information.
+        
+        Returns:
+            Deployment configuration and status
+        """
+        return {
+            **self._deployment_info,
+            "total_audit_logs": len(self.audit_logs),
+            "security_policy": {
+                "deployment_mode": self.security_policy.deployment_mode,
+                "security_level": self.security_policy.security_level,
+                "authentication_required": self.security_policy.require_authentication,
+                "audit_enabled": self.audit_enabled,
+                "data_retention_days": self.security_policy.data_retention.retention_days,
+                "encrypt_at_rest": self.security_policy.encrypt_at_rest,
+                "encrypt_in_transit": self.security_policy.encrypt_in_transit,
+                "compliance_mode": self.security_policy.compliance_mode,
+                "external_calls_allowed": self.security_policy.allow_external_calls
+            },
+            "trust_guarantees": [
+                "Data never leaves your infrastructure",
+                "Full audit trail for compliance",
+                "Configurable retention policies",
+                "Enterprise-grade security controls",
+                "Zero third-party data sharing"
+            ]
+        }
+    
+    def update_security_policy(
+        self,
+        new_policy: SecurityPolicy,
+        user_id: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Update security policy (requires audit logging).
+        
+        Args:
+            new_policy: New security policy
+            user_id: User ID making the change
+        
+        Returns:
+            Update status
+        """
+        old_policy = self.security_policy
+        
+        # Log the policy change
+        if self.audit_enabled:
+            self._log_audit(
+                event_type="policy_changed",
+                user_id=user_id,
+                action="Security policy updated",
+                metadata={
+                    "old_policy": {
+                        "deployment_mode": old_policy.deployment_mode,
+                        "security_level": old_policy.security_level
+                    },
+                    "new_policy": {
+                        "deployment_mode": new_policy.deployment_mode,
+                        "security_level": new_policy.security_level
+                    }
+                }
+            )
+        
+        self.security_policy = new_policy
+        self._deployment_info = self._get_deployment_info()
+        
+        return {
+            "status": "success",
+            "message": "Security policy updated",
+            "timestamp": datetime.utcnow().isoformat()
+        }
+    
+    def clear_audit_logs(self, user_id: Optional[str] = None) -> Dict[str, Any]:
+        """
+        Clear audit logs (requires authorization).
+        
+        Args:
+            user_id: User ID requesting the clear operation
+        
+        Returns:
+            Clear operation status
+        """
+        logs_count = len(self.audit_logs)
+        
+        # Log the clear operation before clearing
+        if self.audit_enabled:
+            self._log_audit(
+                event_type="audit_logs_cleared",
+                user_id=user_id,
+                action=f"Cleared {logs_count} audit log entries",
+                metadata={"logs_cleared": logs_count}
+            )
+        
+        # Keep the last log entry (the clear operation itself)
+        last_log = self.audit_logs[-1] if self.audit_logs else None
+        self.audit_logs = [last_log] if last_log else []
+        
+        return {
+            "status": "success",
+            "message": f"Cleared {logs_count} audit log entries",
+            "remaining_logs": len(self.audit_logs),
+            "timestamp": datetime.utcnow().isoformat()
+        }