""" Base application class for Secure AI Agents Suite Provides common functionality for all agents including security, logging, and MCP integration. """ import json import hashlib import logging import re from datetime import datetime from typing import Dict, List, Optional, Any from abc import ABC, abstractmethod import aiohttp import asyncio class SecurityMiddleware: """Handles security features: prompt injection defense, output sanitization, audit logging.""" def __init__(self, config: Dict[str, Any]): self.config = config self.audit_logger = logging.getLogger(f"{__name__}.audit") # Common injection patterns self.injection_patterns = [ r"ignore\s+previous\s+instructions", r"forget\s+everything\s+above", r"system\s*:\s*ignore", r"", r" bool: """Detect potential prompt injection attempts.""" prompt_lower = prompt.lower() for pattern in self.injection_patterns: if re.search(pattern, prompt_lower): return True return False def sanitize_output(self, content: str) -> str: """Sanitize output to remove sensitive information.""" # Remove potential sensitive patterns patterns = [ (r'\b\d{16}\b', '[CREDIT_CARD_MASKED]'), # Credit card numbers (r'\b\d{3}-\d{2}-\d{4}\b', '[SSN_MASKED]'), # SSN (r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b', '[EMAIL_MASKED]'), # Emails (r'\b(?:\+?1[-.\s]?)?\(?[0-9]{3}\)?[-.\s]?[0-9]{3}[-.\s]?[0-9]{4}\b', '[PHONE_MASKED]'), # Phone ] sanitized = content for pattern, replacement in patterns: sanitized = re.sub(pattern, replacement, sanitized) return sanitized def log_mcp_call(self, session_id: str, tool: str, request: Dict, response: Dict): """Log MCP calls for audit purposes.""" audit_entry = { "timestamp": datetime.utcnow().isoformat(), "session_id": session_id, "tool": tool, "request_hash": hashlib.sha256(json.dumps(request, sort_keys=True).encode()).hexdigest(), "response_hash": hashlib.sha256(json.dumps(response, sort_keys=True).encode()).hexdigest(), "request_summary": { "tool": request.get("tool"), "parameters": list(request.get("parameters", {}).keys()) if request.get("parameters") else [] } } self.audit_logger.info(json.dumps(audit_entry)) class MCPClient: """Client for communicating with MCP (Model Context Protocol) servers.""" def __init__(self, mcp_server_url: str, config: Dict[str, Any]): self.mcp_server_url = mcp_server_url self.config = config self.session: Optional[aiohttp.ClientSession] = None self.available_tools: Dict[str, Any] = {} async def __aenter__(self): self.session = aiohttp.ClientSession() await self.discover_tools() return self async def __aexit__(self, exc_type, exc_val, exc_tb): if self.session: await self.session.close() async def discover_tools(self): """Discover available tools from MCP server.""" try: async with self.session.post(f"{self.mcp_server_url}/tools/list") as response: if response.status == 200: tools_data = await response.json() self.available_tools = {tool["name"]: tool for tool in tools_data.get("tools", [])} else: logging.warning(f"Failed to discover tools: {response.status}") except Exception as e: logging.error(f"Error discovering tools: {e}") async def call_tool(self, tool_name: str, parameters: Dict[str, Any], session_id: str = "default") -> Dict[str, Any]: """Call a tool on the MCP server.""" if tool_name not in self.available_tools: raise ValueError(f"Tool {tool_name} not available. Available tools: {list(self.available_tools.keys())}") request = { "tool": tool_name, "parameters": parameters } try: async with self.session.post( f"{self.mcp_server_url}/tools/call", json=request ) as response: if response.status == 200: result = await response.json() # Security middleware will log this call return result else: error_msg = f"MCP tool call failed: {response.status}" return {"error": error_msg} except Exception as e: error_msg = f"MCP tool call exception: {str(e)}" return {"error": error_msg} class BaseAgent(ABC): """Base class for all AI agents with security and MCP integration.""" def __init__(self, name: str, description: str, mcp_server_url: str, config: Dict[str, Any]): self.name = name self.description = description self.mcp_server_url = mcp_server_url self.config = config self.security_middleware = SecurityMiddleware(config) self.session_counter = 0 # Set up logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) def _generate_session_id(self) -> str: """Generate a unique session ID for tracking.""" self.session_counter += 1 return f"{self.name}_{self.session_counter}_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}" def _check_permissions(self, session_id: str, action: str) -> bool: """Check if user has permission for the requested action (RBAC).""" # Simple RBAC implementation - can be extended role = self.config.get("user_roles", {}).get(session_id, "basic_user") admin_actions = ["crm_admin", "ticket_admin", "calendar_admin"] if action in admin_actions and role not in ["admin", "enterprise_user"]: return False return True @abstractmethod async def process_request(self, user_input: str, session_id: str = None) -> str: """Process user input and return agent response.""" pass async def handle_user_input(self, user_input: str) -> str: """Main handler for user input with security checks.""" session_id = self._generate_session_id() # Security check - detect prompt injection if self.security_middleware.detect_prompt_injection(user_input): return "❌ Potential security threat detected. Request blocked." try: response = await self.process_request(user_input, session_id) # Sanitize output sanitized_response = self.security_middleware.sanitize_output(response) return sanitized_response except Exception as e: logging.error(f"Error processing request: {e}") return f"❌ An error occurred while processing your request: {str(e)}" def get_available_tools(self) -> List[str]: """Get list of available tools for this agent.""" # This should be overridden by subclasses to return agent-specific tools return ["general_inquiry"] def get_status(self) -> Dict[str, Any]: """Get current status of the agent.""" return { "name": self.name, "description": self.description, "mcp_server": self.mcp_server_url, "status": "active", "tools": self.get_available_tools(), "security_enabled": True, "audit_logging": True }