Amazon Security Lake MCP Server

from abc import ABC, abstractmethod from typing import Any, Dict, List import structlog from ..athena.client import AthenaClient from ..config.settings import Settings logger = structlog.get_logger(__name__) class BaseTool(ABC): def __init__(self, settings: Settings): self.settings = settings self.athena_client = AthenaClient(settings) @abstractmethod async def execute(self, **kwargs) -> Dict[str, Any]: pass def _format_error_response(self, error: str, details: str = "") -> Dict[str, Any]: logger.error("Tool execution error", error=error, details=details) return { "success": False, "error": error, "details": details, "results": [] } def _format_success_response( self, results: List[Dict[str, Any]], metadata: Dict[str, Any] = None ) -> Dict[str, Any]: return { "success": True, "error": None, "results": results, "metadata": metadata or {}, "count": len(results) } def _validate_time_range(self, start_time: str = None, end_time: str = None) -> bool: if not start_time or not end_time: return True try: from dateutil import parser as date_parser start_dt = date_parser.parse(start_time) end_dt = date_parser.parse(end_time) return start_dt <= end_dt except Exception: return False def _truncate_results_if_needed( self, results: List[Dict[str, Any]] ) -> List[Dict[str, Any]]: if len(results) > self.settings.max_query_results: logger.warning( "Truncating results due to limit", original_count=len(results), limit=self.settings.max_query_results ) return results[:self.settings.max_query_results] return results