feat: Split CrewAI and logic in modules.

src/app.py

@@ -1,5 +1,6 @@
 import streamlit as st
-from crewai import Agent, Task, Crew, Process, LLM
+from incident_crew import IncidentReporterCrew
+
 
 # --- Interfaz de Usuario con Streamlit ---
 st.set_page_config(
@@ -11,38 +12,27 @@ st.set_page_config(
 if 'gemini_api_key' not in st.session_state:
     st.session_state['gemini_api_key'] = ''
 
-llm = LLM(
-    model="gemini/gemini-2.0-flash",
-    temperature=0.3
-)
-
-threat_hunter_agent = Agent(
-    role="Security Threat Hunter",
-    goal="Rapidly identify, extract, and analyze key Indicators of Compromise (IOCs) from initial security alerts to provide structured and actionable intelligence.",
-    backstory="Highly specialized in threat intelligence, capable of parsing raw incident data, identifying patterns, and performing virtual lookups on security intelligence databases (like VirusTotal, IPInfo, etc.). Your output must be a concise, structured summary.",
-    verbose=True,
-    allow_delegation=False,
-    llm=llm,
-)
-
-reporter_agent = Agent(
-    role="Incident Response Document Specialist",
-    goal="Generate a formal, professional, and comprehensive Incident Response Report based on the initial alert and the technical intelligence provided by the Threat Hunter.",
-    backstory="An expert in cybersecurity documentation, focusing on clarity, structure, and adherence to industry-standard reporting formats. You transform raw data into polished, readable documents suitable for management and technical teams.",
-    verbose=True,
-    llm=llm,
-)
-
 
 st.title("📄 Generador IA de Informes Post-Mortem")
 
-
 # --- Sidebar para la configuración de la API Key ---
 with st.sidebar:
     st.header("⚙️ Instrucciones")
 
-    st.write("1. 📄 Completa el formulario con los detalles de la incidencia (o pega el texto RAW y deja que la IA extraiga los campos)")
-    st.write("2. 🚀 Presiona el botón y la IA generará un informe técnico estructurado post-mortem de tu incidente.")
+    # Campo de entrada para la API Key
+    st.session_state['gemini_api_key'] = st.text_input(
+        "1. 🔑 Ingresa tu API Key de Gemini",
+        type="password",
+        value=st.session_state['gemini_api_key']
+    )
+    
+    if st.session_state['gemini_api_key']:
+        st.success("API Key cargada. ¡Puedes generar el informe!")
+    else:
+        st.warning("🚨 Por favor, ingresa tu API Key de Gemini.")
+
+    st.write("2. 📄 Completa el formulario con los detalles de la incidencia")
+    st.write("3. 🚀 Presiona el botón y la IA generará un informe técnico estructurado post-mortem de tu incidente.")
 
 
 # --- Formulario de entrada de datos principal ---
@@ -81,67 +71,31 @@ with st.form("post_mortem_form"):
 
 # --- Lógica de procesamiento ---
 if submitted:
+   
+    if not st.session_state.get('gemini_api_key'):
+        st.error("❌ Por favor, ingresa tu API Key de Gemini en el panel lateral (sidebar) antes de continuar.")
 
-    if not all([tipo_alerta, sistema_afectado, fecha_hora, impacto_detalle, acciones_tomadas]):
+    elif not all([tipo_alerta, sistema_afectado, fecha_hora, impacto_detalle, acciones_tomadas]):
         st.error("❌ Por favor, completa todos los campos del formulario para generar el informe.")
 
     else:
+        inputs = {
+            "tipo_alerta": tipo_alerta,
+            "sistema_afectado": sistema_afectado,
+            "fecha_hora": fecha_hora,
+            "impacto_detalle": impacto_detalle,
+            "acciones_tomadas": acciones_tomadas,
+        }
+
         with st.spinner("⏳ Generando informe técnico..."):
-            task_ioc_extraction = Task(
-                description=f"""
-                Input: Initial raw incident alert details
-                {impacto_detalle}
-                Process:
-                1. Extract 10 observable Indicators of Compromise (IOCs) including IPs, domains, file hashes (SHA256, MD5), and URLs from the input.
-                2. For each extracted IOC, simulate querying external intelligence services (like VirusTotal, IPInfo, etc.).
-                3. Synthesize the findings into a clear, structured intelligence summary.
-                Output Requirements: The output MUST be a JSON-like or clearly delimited text block, detailing each IOC, its type, and a summary of the associated risk/reputation found (e.g., "Malicious/Known C2," "Clean," "High Reputation," "Related to Phishing Campaign X"). This summary is the ONLY content that should be passed to the next agent.
-                """,
-                expected_output="A structured summary of technical findings including IPs, hashes, and domains, and their corresponding threat intelligence status, ready for the final report.",
-                agent=threat_hunter_agent,
-            )
-
-            task_report_drafting = Task(
-                description=f"""
-                Input:
-                1. The original user input:
-                - **Tipo de Alerta**: {tipo_alerta}
-                - **Sistema Afectado**: {sistema_afectado}
-                - **Fecha y Hora**: {fecha_hora}
-                - **Detalles del Incidente**: {impacto_detalle}
-                - **Acciones Tomadas (Mitigación Inmediata)**: {acciones_tomadas}
-                2. The structured technical intelligence summary provided by the Threat Hunter.
-
-                Process: Draft a complete, professional Incident Response Report.
-
-                Output Requirements: The report MUST be formatted in Markdown and include the following mandatory sections:
-                1. Executive Summary (A brief, high-level overview).
-                2. Initial Incident Details (Source, Time, Initial Observation).
-                3. Technical Analysis & Indicators of Compromise (IOCs) (A detailed section incorporating ALL findings from the Threat Hunter agent information context).
-                4. Impact Assessment (Potential or confirmed impact on systems/data).
-                5. Mitigation or remediation steps taken (if any).
-
-                The entire final output must be this Markdown report.
-                """,
-                expected_output="A complete Incident Response Report in Spanish language. Formatted as markdown without '```'",
-                agent=reporter_agent,
-                markdown=True,
-                context=[task_ioc_extraction] # Esta tarea espera el resultado de la anterior
-            )
-
-            incident_crew = Crew(
-                        agents=[threat_hunter_agent, reporter_agent],
-                        tasks=[task_ioc_extraction, task_report_drafting],
-                        process=Process.sequential, # Ejecución secuencial: Threat Hunter -> Reporter
-                        verbose=True # Nivel de detalle de la ejecución
-                    )
-
-            final_report = incident_crew.kickoff()
+            final_report = IncidentReporterCrew(
+                api_key=st.session_state['gemini_api_key'],
+                ).crew().kickoff(inputs=inputs)
         
         # Mostrar el resultado
         st.divider()
-        st.success("✅ Informe Post-Mortem Generado")
-        st.code(incident_crew.usage_metrics)
+        st.success("✅ Informe post-mortem generado correctamente.")
+        #st.code(incident_crew.usage_metrics)
         st.markdown(final_report)
         st.download_button(
             label="Descargar Informe (.md)",

src/custom_tools.py

@@ -0,0 +1,60 @@
+import os
+import requests
+from crewai.tools import BaseTool
+from typing import Type, List
+from pydantic import BaseModel, Field
+
+# --- 1. IPInfo Geo Lookup Tool ---
+class IPInfoToolInput(BaseModel):
+    """Input schema for IPInfo Geo Lookup Tool."""
+    target: List[str] = Field(
+        ...,
+        description="A list of IP addresses (e.g., ['8.8.8.8', '1.1.1.1']) to query."
+    )
+
+class IPInfoGeoLookup(BaseTool):
+    # See https://ipinfo.io/developers/lite-api
+    name: str = "IPInfo Geo Lookup"
+    description: str = "Looks up geolocation, Internet Service Provider (ISP), and network details for an IP address. Useful for determining the geographic location of a digital asset."
+    args_schema: Type[BaseModel] = IPInfoToolInput
+
+    def _run(self, target: List[str]) -> str:
+        api_key = os.getenv("IPINFO_APIKEY")
+        if not api_key:
+            return "Error: La variable de entorno IPINFO_APIKEY no está configurada."
+
+        results = []
+        for ip_address in target:
+            try:
+                url = f"https://api.ipinfo.io/lite/{ip_address}?token={api_key}"
+                response = requests.get(url)
+                response.raise_for_status()  # Lanza una excepción para respuestas de error (4xx o 5xx)
+                data = response.json()
+                # Formateamos la salida para que sea legible y útil para el agente
+                result_str = (
+                    f"IP: {data.get('ip')}, Country: {data.get('country')}, ASN: {data.get('as_name')}"
+                )
+                results.append(result_str)
+            except requests.exceptions.RequestException as e:
+                results.append(f"Error al consultar la IP {ip_address}: {e}")
+        return "\n".join(results)
+
+# --- 2. VirusTotal Scanner Tool ---
+class VirusTotalToolInput(BaseModel):
+    """Input schema for VirusTotal Scanner Tool."""
+    resources: List[str] = Field(
+        ..., 
+        description="A list of resources to scan. Valid resources include file hash (MD5, SHA256), a URL, a domain, or an IP address."
+    )
+
+class VirusTotalScanner(BaseTool):
+    # See https://blog.virustotal.com/2024/08/VT-S1-EffectiveResearch.html
+    # See https://docs.virustotal.com/reference/ip-info
+    name: str = "VirusTotal Scanner"
+    description: str = "Analyzes files, URLs, domains, or IP addresses for malware and other security threats using multiple antivirus engines and reputation services. Provides a detailed security analysis report."
+    args_schema: Type[BaseModel] = VirusTotalToolInput
+
+    def _run(self, resources: List[str]) -> str:
+        # Placeholder logic for the actual VirusTotal API call
+        print(f"DEBUG: Running VirusTotal Scan for {resources}")
+        return f"Simulation: Resource {resources} analyzed. 0/90 engines flagged it as malicious. Reputation: Clean."

src/incident_crew.py

@@ -0,0 +1,84 @@
+from crewai import Agent, Task, Crew, Process, LLM
+from custom_tools import IPInfoGeoLookup
+
+class IncidentReporterCrew:
+    def __init__(self, api_key):
+        self.api_key = api_key
+
+        self.llm = LLM(
+            model="gemini/gemini-2.0-flash",
+            temperature=0.3,
+            api_key=self.api_key
+        )
+
+    def threat_hunter_agent(self) -> Agent:
+        return Agent(
+            role="Security Threat Hunter",
+            goal="Rapidly identify, extract, and analyze key Indicators of Compromise (IOCs) from initial security alerts to provide structured and actionable intelligence.",
+            backstory="Highly specialized in threat intelligence, capable of parsing raw incident data, identifying patterns, and performing virtual lookups on security intelligence databases (like VirusTotal, IPInfo, etc.). Your output must be a concise, structured summary.",
+            verbose=True,
+            tools=[IPInfoGeoLookup()],
+            allow_delegation=False,
+            llm=self.llm,
+        )
+
+    def reporter_agent(self) -> Agent:
+        return Agent(
+            role="Incident Response Document Specialist",
+            goal="Generate a formal, professional, and comprehensive Incident Response Report based on the initial alert and the technical intelligence provided by the Threat Hunter.",
+            backstory="An expert in cybersecurity documentation, focusing on clarity, structure, and adherence to industry-standard reporting formats. You transform raw data into polished, readable documents suitable for management and technical teams.",
+            verbose=True,
+            llm=self.llm,
+        )
+
+    def task_ioc_extraction(self) -> Task:
+        return Task(
+            description="""
+            Input: Initial raw incident alert details
+            {impacto_detalle}
+            Process:
+            1. Extract 10 observable Indicators of Compromise (IOCs) including IPs, domains, file hashes (SHA256, MD5), and URLs from the input.
+            2. For each extracted IOC, try querying external intelligence services (like VirusTotal, IPInfo, etc.).
+            3. Synthesize the findings into a clear, structured intelligence summary.
+            Output Requirements: The output MUST be a JSON-like or clearly delimited text block, detailing each IOC, its type, and a summary of the associated risk/reputation found (e.g., "Malicious/Known C2," "Clean," "High Reputation," "Related to Phishing Campaign X"). This summary is the ONLY content that should be passed to the next agent.
+            """,
+            expected_output="A structured summary of technical findings including IPs, hashes, and domains, and their corresponding threat intelligence status, ready for the final report.",
+            agent=self.threat_hunter_agent(),
+        )
+    
+    def task_report_drafting(self) -> Task:
+        return Task(
+            description="""
+            Input:
+            1. The original user input:
+            - **Tipo de Alerta**: {tipo_alerta}
+            - **Sistema Afectado**: {sistema_afectado}
+            - **Fecha y Hora**: {fecha_hora}
+            - **Detalles del Incidente**: {impacto_detalle}
+            - **Acciones Tomadas (Mitigación Inmediata)**: {acciones_tomadas}
+            2. The structured technical intelligence summary provided by the Threat Hunter in context.
+
+            Process: Draft a complete, professional Incident Response Report.
+
+            Output Requirements: The report MUST be formatted in Markdown and include the following mandatory sections:
+            1. Executive Summary (A brief, high-level overview).
+            2. Initial Incident Details (Source, Time, Initial Observation).
+            3. Technical Analysis & Indicators of Compromise (IOCs) (A detailed section incorporating ALL findings from the Threat Hunter agent information context).
+            4. Impact Assessment (Potential or confirmed impact on systems/data).
+            5. Mitigation or remediation steps taken (if any).
+
+            The entire final output must be this Markdown report.
+            """,
+            expected_output="A complete Incident Response Report in Spanish language. Formatted as markdown without '```'",
+            agent=self.reporter_agent(),
+            markdown=True,
+            #context=[self.task_ioc_extraction()] # Esta tarea espera el resultado de la anterior
+        )
+    
+    def crew(self) -> Crew:
+        return Crew(
+            agents=[self.threat_hunter_agent(), self.reporter_agent()],
+            tasks=[self.task_ioc_extraction(), self.task_report_drafting()],
+            process=Process.sequential, # Ejecución secuencial: Threat Hunter -> Reporter
+            verbose=True
+        )