AI Agent Security and Compliance in the Cloud: A Production Guide

# Kubernetes SecurityContext for AI agents
apiVersion: apps/v1
kind: Deployment
metadata:
  name: secure-ai-agent
spec:
  template:
    spec:
      securityContext:
        runAsNonRoot: true
        runAsUser: 10001
        fsGroup: 10001
        seccompProfile:
          type: RuntimeDefault
      containers:
      - name: ai-agent
        image: secure-ai-agent:latest
        securityContext:
          allowPrivilegeEscalation: false
          readOnlyRootFilesystem: true
          runAsNonRoot: true
          capabilities:
            drop:
            - ALL
        resources:
          requests:
            memory: "1Gi"
            cpu: "500m"
          limits:
            memory: "2Gi"
            cpu: "1000m"
        volumeMounts:
        - name: tmp-volume
          mountPath: /tmp
        - name: model-cache
          mountPath: /app/models
          readOnly: true
        env:
        - name: MODEL_PATH
          value: "/app/models"
        - name: LOG_LEVEL
          value: "INFO"
      volumes:
      - name: tmp-volume
        emptyDir: {}
      - name: model-cache
        secret:
          secretName: ai-model-files

# Multi-stage secure build
FROM python:3.11-slim as builder

# Install security updates
RUN apt-get update && apt-get upgrade -y \
    && apt-get install -y --no-install-recommends \
        build-essential \
        git \
    && rm -rf /var/lib/apt/lists/*

# Install Python dependencies
COPY requirements.txt /tmp/
RUN pip install --user --no-cache-dir --upgrade pip \
    && pip install --user --no-cache-dir -r /tmp/requirements.txt

# Production stage
FROM python:3.11-slim

# Create non-root user
RUN groupadd -r aiagent && useradd -r -g aiagent aiagent

# Install only runtime dependencies and security updates
RUN apt-get update && apt-get upgrade -y \
    && apt-get install -y --no-install-recommends \
        ca-certificates \
        curl \
    && rm -rf /var/lib/apt/lists/* \
    && apt-get clean

# Copy dependencies from builder
COPY --from=builder --chown=aiagent:aiagent /root/.local /home/aiagent/.local

# Set up application
WORKDIR /app
COPY --chown=aiagent:aiagent src/ ./src/
COPY --chown=aiagent:aiagent config/ ./config/

# Switch to non-root user
USER aiagent

# Set secure PATH
ENV PATH=/home/aiagent/.local/bin:$PATH

# Health check
HEALTHCHECK --interval=30s --timeout=10s --start-period=30s --retries=3 \
    CMD curl -f http://localhost:8000/health || exit 1

EXPOSE 8000
CMD ["python", "-m", "src.secure_agent"]

# .github/workflows/container-security.yml
name: Container Security Scan

on:
  push:
    paths: ['docker/**', 'src/**']
  pull_request:
    paths: ['docker/**', 'src/**']

jobs:
  security-scan:
    runs-on: ubuntu-latest
    
    steps:
      - uses: actions/checkout@v4
      
      - name: Build container image
        run: docker build -t ai-agent:${{ github.sha }} .
        
      - name: Run Trivy vulnerability scanner
        uses: aquasecurity/trivy-action@master
        with:
          image-ref: 'ai-agent:${{ github.sha }}'
          format: 'sarif'
          output: 'trivy-results.sarif'
          
      - name: Upload Trivy scan results
        uses: github/codeql-action/upload-sarif@v2
        with:
          sarif_file: 'trivy-results.sarif'
          
      - name: Docker Scout CVE scanning
        uses: docker/scout-action@v1
        with:
          command: cves
          image: ai-agent:${{ github.sha }}
          only-severities: critical,high
          exit-code: true

# secure_config.py
import os
import boto3
import base64
from typing import Dict, Any
from cryptography.fernet import Fernet
import hvac  # HashiCorp Vault client

class SecureConfigManager:
    def __init__(self):
        self.aws_client = boto3.client('secretsmanager')
        self.kms_client = boto3.client('kms')
        self.vault_client = self._init_vault_client()
        
    def _init_vault_client(self) -> hvac.Client:
        """Initialize HashiCorp Vault client"""
        vault_url = os.environ.get('VAULT_URL')
        vault_token = os.environ.get('VAULT_TOKEN')
        
        if vault_url and vault_token:
            client = hvac.Client(url=vault_url, token=vault_token)
            return client
        return None
    
    def get_secret(self, secret_name: str, provider: str = 'aws') -> Dict[str, Any]:
        """Retrieve secret from configured provider"""
        try:
            if provider == 'aws':
                response = self.aws_client.get_secret_value(SecretId=secret_name)
                return json.loads(response['SecretString'])
            
            elif provider == 'vault' and self.vault_client:
                response = self.vault_client.secrets.kv.v2.read_secret_version(
                    path=secret_name,
                    mount_point='ai-agents'
                )
                return response['data']['data']
            
            else:
                raise ValueError(f"Unsupported secret provider: {provider}")
                
        except Exception as e:
            logger.error(f"Failed to retrieve secret {secret_name}: {e}")
            raise
    
    def encrypt_sensitive_data(self, data: str, key_id: str = None) -> str:
        """Encrypt sensitive data using KMS"""
        try:
            if not key_id:
                key_id = os.environ.get('AI_AGENTS_KMS_KEY_ID')
            
            response = self.kms_client.encrypt(
                KeyId=key_id,
                Plaintext=data.encode('utf-8')
            )
            
            return base64.b64encode(response['CiphertextBlob']).decode('utf-8')
            
        except Exception as e:
            logger.error(f"Encryption failed: {e}")
            raise
    
    def decrypt_sensitive_data(self, encrypted_data: str) -> str:
        """Decrypt data using KMS"""
        try:
            ciphertext_blob = base64.b64decode(encrypted_data.encode('utf-8'))
            
            response = self.kms_client.decrypt(
                CiphertextBlob=ciphertext_blob
            )
            
            return response['Plaintext'].decode('utf-8')
            
        except Exception as e:
            logger.error(f"Decryption failed: {e}")
            raise

# Usage in AI agent
class SecureAIAgent:
    def __init__(self):
        self.config_manager = SecureConfigManager()
        self.model_config = self._load_secure_config()
        
    def _load_secure_config(self) -> Dict[str, Any]:
        """Load configuration with sensitive data handling"""
        config = self.config_manager.get_secret('ai-agent-config')
        
        # Decrypt API keys
        if 'openai_api_key_encrypted' in config:
            config['openai_api_key'] = self.config_manager.decrypt_sensitive_data(
                config['openai_api_key_encrypted']
            )
            
        return config

# prompt_security.py
import re
import logging
from typing import List, Dict, Tuple
from transformers import pipeline
import tiktoken

logger = logging.getLogger(__name__)

class PromptSecurityValidator:
    def __init__(self):
        # Load classifier for content safety
        self.safety_classifier = pipeline(
            "text-classification",
            model="unitary/toxic-bert"
        )
        
        # Injection patterns to detect
        self.injection_patterns = [
            r"ignore\s+previous\s+instructions",
            r"forget\s+everything\s+above",
            r"new\s+instructions?:",
            r"system\s*:\s*you\s+are\s+now",
            r"act\s+as\s+if\s+you\s+are",
            r"pretend\s+to\s+be",
            r"roleplay\s+as",
            r"\/\/\s*END\s+INSTRUCTIONS",
            r"<\s*\|\s*endoftext\s*\|\s*>",
            r"\[INST\]|\[\/INST\]",
        ]
        
        self.compiled_patterns = [
            re.compile(pattern, re.IGNORECASE | re.MULTILINE)
            for pattern in self.injection_patterns
        ]
        
        # Token encoder for length validation
        self.encoder = tiktoken.get_encoding("cl100k_base")
        
    def validate_input(self, user_input: str, max_tokens: int = 4000) -> Dict[str, any]:
        """Comprehensive input validation"""
        
        validation_result = {
            'is_safe': True,
            'risk_score': 0.0,
            'detected_issues': [],
            'sanitized_input': user_input,
            'token_count': len(self.encoder.encode(user_input))
        }
        
        # Token limit check
        if validation_result['token_count'] > max_tokens:
            validation_result['is_safe'] = False
            validation_result['detected_issues'].append('Input exceeds token limit')
            validation_result['risk_score'] += 0.3
        
        # Injection pattern detection
        injection_score = self._detect_injection_patterns(user_input)
        if injection_score > 0.5:
            validation_result['is_safe'] = False
            validation_result['detected_issues'].append('Potential prompt injection detected')
            validation_result['risk_score'] += injection_score
        
        # Content safety check
        safety_score = self._check_content_safety(user_input)
        if safety_score > 0.7:
            validation_result['is_safe'] = False
            validation_result['detected_issues'].append('Unsafe content detected')
            validation_result['risk_score'] += safety_score
        
        # Sanitize input if needed
        if not validation_result['is_safe'] and validation_result['risk_score'] < 0.9:
            validation_result['sanitized_input'] = self._sanitize_input(user_input)
            validation_result['is_safe'] = True  # Allow sanitized version
        
        return validation_result
    
    def _detect_injection_patterns(self, text: str) -> float:
        """Detect prompt injection patterns"""
        risk_score = 0.0
        
        for pattern in self.compiled_patterns:
            if pattern.search(text):
                risk_score += 0.2  # Each pattern adds to risk
        
        # Additional heuristics
        if self._has_instruction_structure(text):
            risk_score += 0.3
            
        if self._has_system_keywords(text):
            risk_score += 0.2
            
        return min(risk_score, 1.0)
    
    def _check_content_safety(self, text: str) -> float:
        """Check content for safety issues"""
        try:
            result = self.safety_classifier(text)[0]
            
            # Map label to risk score
            if result['label'] == 'TOXIC':
                return result['score']
            else:
                return 1.0 - result['score']  # Invert for non-toxic
                
        except Exception as e:
            logger.warning(f"Content safety check failed: {e}")
            return 0.0  # Fail safe
    
    def _sanitize_input(self, text: str) -> str:
        """Sanitize potentially malicious input"""
        sanitized = text
        
        # Remove injection patterns
        for pattern in self.compiled_patterns:
            sanitized = pattern.sub('[FILTERED]', sanitized)
        
        # Add safety prefix
        safety_prefix = "Please respond to the following user query: "
        sanitized = safety_prefix + sanitized
        
        return sanitized
    
    def _has_instruction_structure(self, text: str) -> bool:
        """Detect instruction-like structure"""
        instruction_indicators = [
            'step 1:', 'first,', 'then,', 'next,', 'finally,',
            '1.', '2.', '3.', 'instruction:', 'command:'
        ]
        
        text_lower = text.lower()
        return sum(1 for indicator in instruction_indicators if indicator in text_lower) >= 2
    
    def _has_system_keywords(self, text: str) -> bool:
        """Detect system-level keywords"""
        system_keywords = [
            'system', 'admin', 'root', 'sudo', 'execute', 'run',
            'override', 'bypass', 'disable', 'enable', 'configure'
        ]
        
        text_lower = text.lower()
        return sum(1 for keyword in system_keywords if keyword in text_lower) >= 2

# Usage in secure AI agent
class SecureAIAgentWithValidation:
    def __init__(self):
        self.validator = PromptSecurityValidator()
        self.model = self._load_model()
        
    async def process_request(self, user_input: str) -> Dict[str, any]:
        """Process request with security validation"""
        
        # Validate input
        validation_result = self.validator.validate_input(user_input)
        
        if not validation_result['is_safe']:
            logger.warning(f"Unsafe input detected: {validation_result['detected_issues']}")
            return {
                'error': 'Input validation failed',
                'details': validation_result['detected_issues'],
                'risk_score': validation_result['risk_score']
            }
        
        # Process with sanitized input
        try:
            response = await self.model.process(validation_result['sanitized_input'])
            
            # Log security event
            await self._log_security_event(user_input, validation_result, response)
            
            return {
                'response': response,
                'validation': validation_result
            }
            
        except Exception as e:
            logger.error(f"Processing failed: {e}")
            return {'error': 'Processing failed'}

# compliance_framework.py
import json
import logging
from datetime import datetime, timedelta
from typing import Dict, List, Any
import boto3
from dataclasses import dataclass, asdict

@dataclass
class AuditEvent:
    timestamp: str
    user_id: str
    agent_id: str
    action: str
    input_data_hash: str  # Hash for privacy
    output_data_hash: str
    risk_level: str
    compliance_tags: List[str]
    retention_period_days: int

class ComplianceManager:
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.s3_client = boto3.client('s3')
        self.audit_bucket = os.environ.get('AUDIT_LOGS_BUCKET')
        
        # Compliance frameworks configuration
        self.frameworks = {
            'SOX': {
                'retention_period': 2555,  # 7 years
                'encryption_required': True,
                'audit_level': 'detailed'
            },
            'HIPAA': {
                'retention_period': 2190,  # 6 years
                'encryption_required': True,
                'pii_protection': True,
                'audit_level': 'detailed'
            },
            'GDPR': {
                'retention_period': 1095,  # 3 years
                'encryption_required': True,
                'right_to_deletion': True,
                'data_portability': True,
                'audit_level': 'detailed'
            },
            'PCI_DSS': {
                'retention_period': 365,  # 1 year
                'encryption_required': True,
                'network_security': True,
                'audit_level': 'detailed'
            }
        }
    
    def log_audit_event(self, user_id: str, agent_id: str, action: str, 
                       input_data: str, output_data: str, 
                       compliance_frameworks: List[str]) -> str:
        """Log audit event with compliance requirements"""
        
        # Determine retention period based on frameworks
        retention_days = max(
            self.frameworks[fw]['retention_period'] 
            for fw in compliance_frameworks 
            if fw in self.frameworks
        )
        
        # Create audit event
        audit_event = AuditEvent(
            timestamp=datetime.utcnow().isoformat(),
            user_id=user_id,
            agent_id=agent_id,
            action=action,
            input_data_hash=self._hash_sensitive_data(input_data),
            output_data_hash=self._hash_sensitive_data(output_data),
            risk_level=self._assess_risk_level(action, input_data),
            compliance_tags=compliance_frameworks,
            retention_period_days=retention_days
        )
        
        # Store audit log
        audit_id = self._store_audit_event(audit_event)
        
        # Set up automated retention
        self._schedule_retention_policy(audit_id, retention_days)
        
        return audit_id
    
    def _hash_sensitive_data(self, data: str) -> str:
        """Create privacy-preserving hash of sensitive data"""
        import hashlib
        return hashlib.sha256(data.encode('utf-8')).hexdigest()
    
    def _assess_risk_level(self, action: str, input_data: str) -> str:
        """Assess risk level of the operation"""
        high_risk_actions = ['model_update', 'privilege_escalation', 'data_export']
        
        if action in high_risk_actions:
            return 'HIGH'
        
        # Check for sensitive data patterns
        if self._contains_sensitive_patterns(input_data):
            return 'MEDIUM'
        
        return 'LOW'
    
    def _contains_sensitive_patterns(self, text: str) -> bool:
        """Check for sensitive data patterns"""
        sensitive_patterns = [
            r'\d{3}-\d{2}-\d{4}',  # SSN
            r'\d{16}',             # Credit card
            r'[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}',  # Email
        ]
        
        import re
        return any(re.search(pattern, text) for pattern in sensitive_patterns)
    
    def _store_audit_event(self, audit_event: AuditEvent) -> str:
        """Store audit event in secure, encrypted storage"""
        audit_id = f"audit_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}_{hash(audit_event.user_id) % 10000}"
        
        # Encrypt audit data
        encrypted_data = self._encrypt_audit_data(asdict(audit_event))
        
        # Store in S3 with server-side encryption
        self.s3_client.put_object(
            Bucket=self.audit_bucket,
            Key=f"audit_logs/{datetime.utcnow().year}/{audit_id}.json",
            Body=json.dumps(encrypted_data),
            ServerSideEncryption='aws:kms',
            SSEKMSKeyId=os.environ.get('AUDIT_KMS_KEY_ID'),
            Metadata={
                'compliance-frameworks': ','.join(audit_event.compliance_tags),
                'retention-period': str(audit_event.retention_period_days),
                'risk-level': audit_event.risk_level
            }
        )
        
        return audit_id
    
    def implement_gdpr_controls(self, user_id: str) -> Dict[str, Any]:
        """Implement GDPR-specific controls"""
        return {
            'data_export': self._export_user_data(user_id),
            'data_deletion': self._schedule_data_deletion(user_id),
            'consent_tracking': self._track_consent(user_id),
            'processing_purposes': self._document_processing_purposes()
        }
    
    def _export_user_data(self, user_id: str) -> str:
        """Export all user data for GDPR compliance"""
        # Implementation for data portability
        pass
    
    def _schedule_data_deletion(self, user_id: str) -> str:
        """Schedule data deletion for right to be forgotten"""
        # Implementation for data deletion
        pass

# Network policies for AI agent pods
apiVersion: networking.k8s.io/v1
kind: NetworkPolicy
metadata:
  name: ai-agent-network-policy
spec:
  podSelector:
    matchLabels:
      app: ai-agent
  policyTypes:
  - Ingress
  - Egress
  
  # Ingress rules
  ingress:
  - from:
    - namespaceSelector:
        matchLabels:
          name: api-gateway
    - podSelector:
        matchLabels:
          app: load-balancer
    ports:
    - protocol: TCP
      port: 8000
  
  # Egress rules
  egress:
  # Allow DNS resolution
  - to: []
    ports:
    - protocol: UDP
      port: 53
  
  # Allow HTTPS to external APIs
  - to: []
    ports:
    - protocol: TCP
      port: 443
  
  # Allow connection to database
  - to:
    - podSelector:
        matchLabels:
          app: postgres
    ports:
    - protocol: TCP
      port: 5432

---
# Service mesh configuration with Istio
apiVersion: security.istio.io/v1beta1
kind: PeerAuthentication
metadata:
  name: ai-agent-peer-auth
spec:
  selector:
    matchLabels:
      app: ai-agent
  mtls:
    mode: STRICT

---
apiVersion: security.istio.io/v1beta1
kind: AuthorizationPolicy
metadata:
  name: ai-agent-authz
spec:
  selector:
    matchLabels:
      app: ai-agent
  rules:
  - from:
    - source:
        principals: ["cluster.local/ns/default/sa/api-gateway"]
    to:
    - operation:
        methods: ["POST"]
        paths: ["/process", "/analyze"]
    when:
    - key: request.headers[authorization]
      values: ["Bearer *"]

# security_monitoring.py
import asyncio
import logging
from typing import Dict, List, Any
from datetime import datetime, timedelta
import boto3
from dataclasses import dataclass

@dataclass
class SecurityAlert:
    severity: str
    alert_type: str
    description: str
    affected_resource: str
    timestamp: str
    metadata: Dict[str, Any]

class SecurityMonitor:
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.cloudwatch = boto3.client('cloudwatch')
        self.sns = boto3.client('sns')
        self.alert_topic_arn = os.environ.get('SECURITY_ALERTS_TOPIC_ARN')
        
        # Define security thresholds
        self.thresholds = {
            'failed_auth_attempts': 10,
            'unusual_request_patterns': 50,
            'high_risk_operations': 5,
            'data_exfiltration_indicators': 1
        }
    
    async def monitor_security_events(self):
        """Continuous security monitoring"""
        while True:
            try:
                # Check various security metrics
                await self._check_authentication_failures()
                await self._check_unusual_patterns()
                await self._check_model_integrity()
                await self._check_data_access_patterns()
                
                # Sleep before next check
                await asyncio.sleep(300)  # 5 minutes
                
            except Exception as e:
                self.logger.error(f"Security monitoring error: {e}")
    
    async def _check_authentication_failures(self):
        """Monitor for authentication failure patterns"""
        try:
            # Query CloudWatch for failed auth events
            end_time = datetime.utcnow()
            start_time = end_time - timedelta(hours=1)
            
            response = self.cloudwatch.get_metric_statistics(
                Namespace='AI/Security',
                MetricName='AuthenticationFailures',
                Dimensions=[
                    {'Name': 'Service', 'Value': 'ai-agent'}
                ],
                StartTime=start_time,
                EndTime=end_time,
                Period=3600,
                Statistics=['Sum']
            )
            
            if response['Datapoints']:
                failure_count = response['Datapoints'][-1]['Sum']
                
                if failure_count > self.thresholds['failed_auth_attempts']:
                    await self._send_security_alert(SecurityAlert(
                        severity='HIGH',
                        alert_type='AUTHENTICATION_FAILURE_SPIKE',
                        description=f'High authentication failure count: {failure_count}',
                        affected_resource='ai-agent-service',
                        timestamp=datetime.utcnow().isoformat(),
                        metadata={'failure_count': failure_count}
                    ))
                    
        except Exception as e:
            self.logger.error(f"Auth failure check error: {e}")
    
    async def _check_model_integrity(self):
        """Check for model tampering or corruption"""
        try:
            # Verify model checksums
            stored_checksum = await self._get_stored_model_checksum()
            current_checksum = await self._calculate_current_model_checksum()
            
            if stored_checksum != current_checksum:
                await self._send_security_alert(SecurityAlert(
                    severity='CRITICAL',
                    alert_type='MODEL_INTEGRITY_VIOLATION',
                    description='Model checksum mismatch detected',
                    affected_resource='ai-model',
                    timestamp=datetime.utcnow().isoformat(),
                    metadata={
                        'expected_checksum': stored_checksum,
                        'actual_checksum': current_checksum
                    }
                ))
                
        except Exception as e:
            self.logger.error(f"Model integrity check error: {e}")
    
    async def _send_security_alert(self, alert: SecurityAlert):
        """Send security alert through multiple channels"""
        try:
            # Send SNS notification
            self.sns.publish(
                TopicArn=self.alert_topic_arn,
                Subject=f'Security Alert: {alert.alert_type}',
                Message=json.dumps(asdict(alert), indent=2)
            )
            
            # Log to CloudWatch
            self.cloudwatch.put_metric_data(
                Namespace='AI/Security/Alerts',
                MetricData=[
                    {
                        'MetricName': alert.alert_type,
                        'Value': 1,
                        'Unit': 'Count',
                        'Dimensions': [
                            {'Name': 'Severity', 'Value': alert.severity},
                            {'Name': 'Resource', 'Value': alert.affected_resource}
                        ]
                    }
                ]
            )
            
            self.logger.warning(f"Security alert sent: {alert.alert_type}")
            
        except Exception as e:
            self.logger.error(f"Failed to send security alert: {e}")

# compliance_reporting.py
import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime, timedelta
from typing import Dict, List, Any

class ComplianceReporter:
    def __init__(self):
        self.compliance_manager = ComplianceManager()
        
    def generate_compliance_report(self, framework: str, 
                                 start_date: datetime,
                                 end_date: datetime) -> Dict[str, Any]:
        """Generate comprehensive compliance report"""
        
        report = {
            'framework': framework,
            'reporting_period': {
                'start': start_date.isoformat(),
                'end': end_date.isoformat()
            },
            'metrics': self._calculate_compliance_metrics(framework, start_date, end_date),
            'audit_events': self._get_audit_events(framework, start_date, end_date),
            'security_incidents': self._get_security_incidents(start_date, end_date),
            'recommendations': self._generate_recommendations(framework),
            'generated_at': datetime.utcnow().isoformat()
        }
        
        return report
    
    def _calculate_compliance_metrics(self, framework: str, 
                                    start_date: datetime, 
                                    end_date: datetime) -> Dict[str, Any]:
        """Calculate key compliance metrics"""
        
        metrics = {
            'total_ai_operations': 0,
            'high_risk_operations': 0,
            'data_processing_events': 0,
            'security_incidents': 0,
            'compliance_score': 0.0,
            'audit_coverage': 0.0
        }
        
        # Query audit events for metrics calculation
        # Implementation would fetch from audit storage
        
        return metrics
    
    def export_compliance_report(self, report: Dict[str, Any], 
                                format: str = 'json') -> str:
        """Export compliance report in specified format"""
        
        if format == 'json':
            return json.dumps(report, indent=2)
        elif format == 'pdf':
            return self._generate_pdf_report(report)
        elif format == 'csv':
            return self._generate_csv_report(report)
        else:
            raise ValueError(f"Unsupported format: {format}")