# Enfoque tradicional: función personalizada para cada herramienta
def analyze_security_logs(log_path):
    # Implementación específica y no reutilizable
    pass

def check_ip_reputation(ip):
    # Otra implementación específica
    pass

# El LLM debe conocer estas funciones específicas
available_tools = [analyze_security_logs, check_ip_reputation]

# Enfoque MCP: servidor estandarizado
@app.tool()
def analyze_log_file(file_path: str, analysis_type: str) -> str:
    """Analiza un archivo de logs para identificar actividad sospechosa."""
    # Implementación con interfaz estandarizada
    
@app.tool() 
def check_ip_reputation(ip_address: str) -> str:
    """Verifica la reputación de una dirección IP."""
    # Implementación con interfaz estandarizada

# Aplicación 1: ChatGPT
def chatgpt_log_analyzer(logs):
    # Implementación específica para ChatGPT
    return analysis

def chatgpt_ip_checker(ip):
    # Implementación específica para ChatGPT
    return reputation

# Aplicación 2: Claude  
def claude_log_analyzer(logs):
    # Implementación específica para Claude
    return analysis

def claude_ip_checker(ip):
    # Implementación específica para Claude  
    return reputation

# Aplicación 3: Tu app personalizada
def custom_log_analyzer(logs):
    # Implementación específica para tu app
    return analysis

# 1 servidor MCP para todas las herramientas (n)
app = FastMCP("Universal Tool Server")

@app.tool()
def analyze_log_file(file_path: str, analysis_type: str) -> str:
    """Una implementación que funciona con CUALQUIER cliente MCP."""
    return json.dumps(analysis_results)

@app.tool() 
def check_ip_reputation(ip_address: str) -> str:
    """Una implementación que funciona con CUALQUIER cliente MCP."""
    return json.dumps(reputation_data)

async def refresh_available_tools(self):
    """Descubre dinámicamente todas las herramientas del servidor MCP"""
    
    if not self.session:
        raise Exception("No MCP session established")
    
    try:
        # El cliente pregunta: "¿Qué herramientas tienes disponibles?"
        response = await asyncio.wait_for(
            self.session.list_tools(),
            timeout=5.0
        )
        
        self.available_tools = []
        for tool in response.tools:
            # Cada herramienta se autodescribe con metadatos ricos
            tool_spec = {
                'toolSpec': {
                    'name': tool.name,
                    'description': tool.description,
                    'inputSchema': {
                        'json': tool.inputSchema  # Schema JSON completo
                    }
                }
            }
            self.available_tools.append(tool_spec)
        
        # Logging automático de capacidades descubiertas
        print(f"🔧 Discovered {len(self.available_tools)} specialized tools:")
        for tool in self.available_tools:
            print(f"  • {tool['toolSpec']['name']}: {tool['toolSpec']['description']}")
                
    except Exception as e:
        print(f"❌ Failed to discover tools: {e}")
        raise

#!/usr/bin/env python3
"""
Servidor MCP para Análisis Forense Digital
Expone herramientas especializadas mediante protocolo estandarizado
"""

import json
from datetime import datetime
from typing import Dict, Any

try:
    from mcp.server.fastmcp import FastMCP
    MCP_AVAILABLE = True
except ImportError:
    print("❌ Install MCP: pip install mcp")
    MCP_AVAILABLE = False

# Inicializar servidor con herramientas forenses
if MCP_AVAILABLE:
    app = FastMCP("Digital Forensics MCP Server")

@app.tool()
def analyze_log_file(file_path: str, analysis_type: str = "security") -> str:
    """
    Analiza un archivo de logs para identificar actividad sospechosa.
    
    Args:
        file_path: Ruta al archivo de log
        analysis_type: Tipo de análisis (security, network, authentication)
    
    Returns:
        JSON con análisis detallado del log
    """
    
    # NOTA: En un entorno real, esto leería archivos reales
    # Para la demo, usamos datos simulados que representan patrones típicos
    sample_security_events = [
        "2025-01-20 14:23:15 [WARNING] Authentication failure from 192.168.1.100",
        "2025-01-20 14:26:45 [CRITICAL] Suspicious PowerShell execution on WORKSTATION-01",
        "2025-01-20 14:27:10 [WARNING] Outbound connection to evil-domain.com",
        "2025-01-20 14:29:15 [CRITICAL] Process injection detected: PID 1234 → PID 5678"
    ]
    
    analysis_results = {
        "file_analyzed": file_path,
        "analysis_type": analysis_type,
        "timestamp": datetime.now().isoformat(),
        "findings": [],
        "risk_score": 0,
        "recommendations": []
    }
    
    # Detectar patrones de ataque usando lógica de análisis forense
    for log_entry in sample_security_events:
        if "Authentication failure" in log_entry:
            analysis_results["findings"].append({
                "type": "brute_force_attack",
                "severity": "HIGH",
                "description": "Multiple authentication failures detected",
                "indicators": ["credential_stuffing", "automated_attack"]
            })
            analysis_results["risk_score"] += 25
            
        elif "PowerShell execution" in log_entry:
            analysis_results["findings"].append({
                "type": "living_off_the_land",
                "severity": "CRITICAL",
                "description": "Suspicious PowerShell activity",
                "indicators": ["fileless_malware", "encoded_commands"]
            })
            analysis_results["risk_score"] += 40
    
    # Generar recomendaciones basadas en hallazgos
    if analysis_results["risk_score"] > 80:
        analysis_results["recommendations"].extend([
            "Immediate incident response required",
            "Isolate affected systems from network",
            "Deploy additional monitoring on critical assets"
        ])
    
    return json.dumps(analysis_results, indent=2)

class ForensicMCPClient:
    """Cliente MCP que conecta herramientas forenses con Bedrock"""
    
    def __init__(self, mcp_server_path: str, aws_region: str = "us-east-1"):
        self.mcp_server_path = mcp_server_path
        self.aws_region = aws_region
        self.available_tools = []  # Se puebla dinámicamente via auto-descubrimiento
        self.conversation_history = []
        self.mcp_connected = False
        
        # Inicializar cliente Bedrock
        self.bedrock_client = boto3.client('bedrock-runtime', region_name=aws_region)
        
        # Modelos disponibles
        self.available_models = {
            "claude-3-7-sonnet": "us.anthropic.claude-3-7-sonnet-20250219-v1:0",
            "claude-3-5-haiku": "us.anthropic.claude-3-5-haiku-20241022-v1:0", 
            "nova-pro": "us.amazon.nova-pro-v1:0",
            "nova-lite": "us.amazon.nova-lite-v1:0"
        }
        self.current_model = self.available_models["claude-3-7-sonnet"]

async def query_bedrock(self, user_prompt: str, system_prompt: str = None) -> Dict[str, Any]:
    """Consulta Bedrock usando herramientas MCP auto-descubiertas"""
    
    # Construir mensajes para Bedrock
    messages = []
    
    # Agregar historial de conversación
    for msg in self.conversation_history:
        messages.append(msg)
    
    # Agregar mensaje del usuario
    messages.append({
        "role": "user",
        "content": [{"text": user_prompt}]
    })
    
    # Sistema prompt especializado para análisis forense por defecto
    if not system_prompt:
        system_prompt = """Eres un experto en análisis forense digital y cyberseguridad. 

    Tienes acceso a herramientas especializadas que fueron auto-descubiertas:
    - Análisis de logs de seguridad
    - Verificación de reputación de IPs
    - Extracción de indicadores de compromiso (IOCs)  
    - Generación de timelines de incidentes
    - Creación de reportes forenses

    Usa estas herramientas de manera inteligente para investigar incidentes."""
    
    try:
        # Bedrock recibe las herramientas auto-descubiertas
        response = self.bedrock_client.converse(
            modelId=self.current_model,
            messages=messages,
            system=[{"text": system_prompt}],
            toolConfig={
                "tools": self.available_tools,  # Herramientas descubiertas dinámicamente
                "toolChoice": {"auto": {}}
            },
            inferenceConfig={
                "maxTokens": 4000,
                "temperature": 0.1,  # Precisión para análisis forense
                "topP": 0.9
            }
        )
        
        return response
        
    except ClientError as e:
        error_code = e.response['Error']['Code']
        if error_code == 'AccessDeniedException':
            raise Exception("Access denied to Bedrock. Check AWS credentials.")
        else:
            raise Exception(f"Bedrock error: {error_code}")

async def process_tool_use_response(self, response: Dict[str, Any]) -> Dict[str, Any]:
    """Procesa automáticamente el uso iterativo de herramientas por parte de Bedrock"""
    
    max_iterations = 10  # Prevenir bucles infinitos
    current_iteration = 0
    current_response = response
    
    # Bucle iterativo para manejar múltiples rondas de herramientas
    while (current_response.get("stopReason") == "tool_use" and 
           current_iteration < max_iterations):
        
        current_iteration += 1
        print(f"🔄 Processing tool use iteration {current_iteration}...")
        
        message = current_response["output"]["message"]
        tool_requests = message["content"]
        
        # Agregar mensaje del asistente al historial
        self.conversation_history.append(message)
        
        # Procesar cada solicitud de herramienta en esta iteración
        for tool_request in tool_requests:
            if "toolUse" in tool_request:
                tool_use = tool_request["toolUse"]
                tool_id = tool_use["toolUseId"]
                tool_name = tool_use["name"] 
                tool_input = tool_use["input"]
                
                print(f"🔧 Executing tool: {tool_name}")
                
                try:
                    # Ejecutar herramienta MCP
                    tool_result = await self.execute_mcp_tool(tool_name, tool_input)
                    
                    # Agregar resultado de herramienta al historial
                    self.conversation_history.append({
                        "role": "user",
                        "content": [{
                            "toolResult": {
                                "toolUseId": tool_id,
                                "content": [{"text": tool_result}]
                            }
                        }]
                    })
                    
                except Exception as e:
                    # Agregar error al historial para que Claude lo sepa
                    self.conversation_history.append({
                        "role": "user",
                        "content": [{
                            "toolResult": {
                                "toolUseId": tool_id,
                                "content": [{"text": f"Error executing tool: {str(e)}"}],
                                "status": "error"
                            }
                        }]
                    })
        
        # Obtener siguiente respuesta de Bedrock
        print(f"🤖 Getting Bedrock response after tool execution...")
        current_response = await self.query_bedrock_with_history()
    
    print(f"✅ Tool processing completed after {current_iteration} iterations")
    return current_response

async def analyze_security_incident(self, incident_description: str) -> str:
    """Flujo completo: descubrimiento → análisis → reporte"""
    
    print(f"🚨 INICIANDO ANÁLISIS DE INCIDENTE DE SEGURIDAD")
    print(f"📝 {incident_description}")
    print("=" * 60)
    
    # Verificar que MCP esté conectado y herramientas descubiertas
    if not self.mcp_connected:
        raise Exception("MCP not connected. Connect first.")
    
    # Limpiar historial para análisis limpio
    self.conversation_history = []
    
    analysis_prompt = f"""
    Analiza este incidente de seguridad usando todas las herramientas disponibles:
    
    INCIDENTE: {incident_description}
    
    Ejecuta un análisis forense COMPLETO en este orden específico:
    
    1. ANÁLISIS DE LOGS: Usa analyze_log_file() para examinar logs relevantes
    2. VERIFICACIÓN DE IPs: Usa check_ip_reputation() para todas las IPs mencionadas  
    3. EXTRACCIÓN DE IOCs: Usa extract_iocs() para identificar indicadores de compromiso
    4. TIMELINE: Usa generate_timeline() para crear cronología del ataque
    5. REPORTE: Usa generate_incident_report() para reporte ejecutivo final
    
    Proporciona recomendaciones específicas de contención y pasos de seguimiento.
    """
    
    # Bedrock automáticamente decide qué herramientas usar
    response = await self.query_bedrock(analysis_prompt)
    
    # Procesar uso de herramientas de manera iterativa
    if response.get("stopReason") == "tool_use":
        response = await self.process_tool_use_response(response)
    
    # Extraer respuesta final
    final_message = response["output"]["message"]
    self.conversation_history.append(final_message)
    
    # Combinar texto de respuesta
    response_text = ""
    for part in final_message.get("content", []):
        if "text" in part:
            response_text += part["text"]
    
    return response_text

class SecureForensicClient(ForensicMCPClient):
    """Cliente forense con controles de seguridad enterprise"""
    
    def sanitize_sensitive_data(self, forensic_data: Dict) -> Dict:
        """Sanitiza datos sensibles antes de enviar a LLM"""
        
        sanitized = forensic_data.copy()
        
        # Enmascarar IPs internas según política corporativa
        sanitized = self.mask_internal_ips(sanitized)
        
        # Redactar credenciales y secretos
        sanitized = self.redact_credentials(sanitized)
        
        # Hash hostnames internos para privacidad
        sanitized = self.hash_internal_hostnames(sanitized)
        
        return sanitized
    
    async def audit_tool_execution(self, tool_name: str, arguments: Dict, result: str):
        """Registra toda ejecución de herramientas para compliance"""
        
        audit_record = {
            "timestamp": datetime.now().isoformat(),
            "tool_name": tool_name,
            "arguments_hash": hashlib.sha256(str(arguments).encode()).hexdigest(),
            "result_length": len(result),
            "user_id": self.get_current_user(),
            "session_id": self.get_session_id(),
            "compliance_flags": self.check_compliance(tool_name, arguments)
        }
        
        await self.access_logger.log(audit_record)