Add new configuration and tools

.env.template

@@ -0,0 +1,11 @@
+# API Keys and Endpoints
+NVD_API_KEY=""
+
+# Model Configuration
+MODEL_ID="Qwen/Qwen2.5-Coder-32B-Instruct"
+MAX_TOKENS=2096
+TEMPERATURE=0.3
+
+# Application Configuration
+ENV="development"
+DEBUG=true 

agent.json

@@ -1,9 +1,43 @@
 {
+    "agent_config": {
+        "name": "VulnerabilityIntelligenceAgent",
+        "description": "Agente especializado en análisis de vulnerabilidades y amenazas de ciberseguridad",
+        "model": {
+            "type": "HfApiModel",
+            "model_id": "Qwen/Qwen2.5-Coder-32B-Instruct",
+            "max_tokens": 2096,
+            "temperature": 0.3
+        },
+        "max_steps": 10,
+        "verbosity_level": 2
+    },
     "tools": [
-        "web_search",
-        "visit_webpage",
-        "final_answer"
+        {
+            "name": "final_answer",
+            "type": "FinalAnswerTool",
+            "description": "Proporciona una respuesta final al problema"
+        },
+        {
+            "name": "web_search",
+            "type": "DuckDuckGoSearchTool",
+            "description": "Realiza búsquedas web usando DuckDuckGo"
+        },
+        {
+            "name": "visit_webpage",
+            "type": "VisitWebpageTool",
+            "description": "Visita y extrae contenido de páginas web"
+        },
+        {
+            "name": "vuln_search",
+            "type": "VulnerabilitySearchTool",
+            "description": "Busca información sobre vulnerabilidades en múltiples fuentes"
+        }
     ],
+    "prompt_templates": {
+        "path": "prompts.yaml",
+        "default_system_prompt": "system_prompt",
+        "default_user_prompt": "user_prompt"
+    },
     "model": {
         "class": "HfApiModel",
         "data": {

app.py

@@ -1,113 +1,115 @@
-#!/usr/bin/env python3
-"""
-Gradio UI for the Vulnerability Intelligence Agent (VIA).
-This provides a chat interface to interact with the VIA using natural language.
-"""
 import os
-import sys
-import argparse
-import logging
-from typing import Dict, List, Any, Optional
-
+import json
+import yaml
+from dotenv import load_dotenv
 import gradio as gr
-from smolagents import CodeAgent, HfApiModel, GradioUI
-from smolagents.tools import load_tool, tool
-
-# Asegurarse de que el directorio actual esté en sys.path para que los imports funcionen
-sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+from smolagents import CodeAgent
+from smolagents.models import HfApiModel
+from tools.final_answer import FinalAnswerTool
+from tools.web_search import DuckDuckGoSearchTool
+from tools.visit_webpage import VisitWebpageTool
+from tools.vuln_search import VulnerabilitySearchTool
 
-from agents.coordinator_agent import search_vulnerabilities_for_software, get_vulnerability_details
-from tools import utils
+# Load environment variables
+load_dotenv()
 
-# Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-)
-logger = utils.setup_logger("gradio_ui")
+def load_agent_config():
+    """Load agent configuration from agent.json"""
+    with open('agent.json', 'r') as f:
+        return json.load(f)
 
-# Cargar las herramientas básicas usando las que ya existen en smolagents
-final_answer = load_tool("smolagents/final_answer", trust_remote_code=True)
-
-def get_agent_description():
-    """
-    Get the description for the agent.
-    """
-    return """
-    # 🔐 Vulnerability Intelligence Agent (VIA)
-    
-    I am an intelligent agent designed to help you find vulnerabilities in software and systems.
-    
-    ## What I can do:
-    - Search for known vulnerabilities in software by name and version
-    - Provide detailed information about specific vulnerabilities (CVE, CWE, etc.)
-    - Generate reports about vulnerabilities
-    
-    ## How to use me:
-    - Ask about vulnerabilities in specific software, e.g., "Find vulnerabilities in OpenSSL 1.1.1k"
-    - Ask about a specific vulnerability, e.g., "Tell me about CVE-2021-44228"
-    - Use natural language to describe what you're looking for
-    
-    ## Examples:
-    - "What vulnerabilities exist in Apache 2.4.54?"
-    - "Are there any critical vulnerabilities in log4j 2.14.1?"
-    - "Give me details about CVE-2021-44228"
-    - "What security issues should I be aware of in OpenSSL 1.1.1k?"
-    """
+def load_prompts():
+    """Load prompt templates from prompts.yaml"""
+    with open('prompts.yaml', 'r') as f:
+        return yaml.safe_load(f)
 
-def create_parser():
-    """Create command line argument parser."""
-    parser = argparse.ArgumentParser(description="Vulnerability Intelligence Agent (VIA) UI")
-    parser.add_argument("--port", type=int, default=7860, help="Port to run the Gradio app on")
-    parser.add_argument("--host", type=str, default="127.0.0.1", help="Host to run the Gradio app on")
-    parser.add_argument("--model", type=str, default="Qwen/Qwen2.5-Coder-32B-Instruct", 
-                        help="HuggingFace model ID to use")
-    parser.add_argument("--share", action="store_true", help="Create a public link")
-    parser.add_argument("--verbose", action="store_true", help="Enable verbose logging")
-    
-    return parser
+def initialize_tools():
+    """Initialize agent tools"""
+    tools = {
+        'final_answer': FinalAnswerTool(),
+        'web_search': DuckDuckGoSearchTool(),
+        'visit_webpage': VisitWebpageTool(),
+        'vuln_search': VulnerabilitySearchTool()
+    }
+    return tools
 
-def main():
-    """Main entry point for the Gradio UI."""
-    args = create_parser().parse_args()
+def create_agent():
+    """Create and configure the vulnerability agent"""
+    config = load_agent_config()
+    prompts = load_prompts()
     
-    # Configure logging level
-    log_level = logging.DEBUG if args.verbose else logging.INFO
-    logging.basicConfig(level=log_level)
-    
-    # Initialize the model
+    # Configure model
+    model_config = config['agent_config']['model']
     model = HfApiModel(
-        max_tokens=2096,
-        temperature=0.5,
-        model_id=args.model,
-        custom_role_conversions=None,
+        model_id=model_config['model_id'],
+        max_tokens=model_config['max_tokens'],
+        temperature=model_config['temperature']
     )
     
-    # Initialize the agent con las herramientas ya existentes y las que hemos creado
+    # Initialize tools
+    tools = initialize_tools()
+    
+    # Create agent
     agent = CodeAgent(
         model=model,
-        tools=[search_vulnerabilities_for_software, get_vulnerability_details, final_answer],
-        max_steps=10,
-        verbosity_level=2 if args.verbose else 1,
+        tools=tools,
+        max_steps=config['agent_config']['max_steps'],
+        verbosity_level=config['agent_config']['verbosity_level']
     )
     
-    # Create Gradio UI
-    ui = GradioUI(agent)
+    return agent, prompts
+
+def process_query(query, analysis_type="general"):
+    """Process a user query"""
+    agent, prompts = create_agent()
     
-    # Launch the UI
-    ui.launch(
-        share=args.share,
-        server_name=args.host,
-        server_port=args.port,
-        show_api=False,
-        favicon_path=None,
-        allowed_paths=[],
-        app_kwargs={
-            "title": "🔐 Vulnerability Intelligence Agent (VIA)",
-            "description": get_agent_description(),
-            "theme": gr.themes.Base(),
-        },
-    )
+    # Select appropriate template
+    if analysis_type == "vulnerability":
+        template = prompts['vulnerability_analysis']
+        formatted_prompt = template.format(cve_id=query)
+    elif analysis_type == "threat":
+        template = prompts['threat_report']
+        formatted_prompt = template.format(target=query)
+    else:
+        template = prompts['user_prompt']
+        formatted_prompt = template.format(query=query)
+    
+    # Execute agent
+    system_prompt = prompts['system_prompt']
+    result = agent.run(formatted_prompt, system_prompt=system_prompt)
+    
+    return result
+
+# Gradio Interface
+def create_interface():
+    """Create the Gradio user interface"""
+    with gr.Blocks(title="Vulnerability Intelligence Agent") as interface:
+        gr.Markdown("# Vulnerability Intelligence Agent (VIA)")
+        
+        with gr.Row():
+            with gr.Column():
+                query_input = gr.Textbox(
+                    label="Query",
+                    placeholder="Enter your security query..."
+                )
+                analysis_type = gr.Radio(
+                    choices=["general", "vulnerability", "threat"],
+                    label="Analysis Type",
+                    value="general"
+                )
+                submit_btn = gr.Button("Analyze")
+            
+            with gr.Column():
+                output = gr.Markdown(label="Result")
+        
+        submit_btn.click(
+            fn=process_query,
+            inputs=[query_input, analysis_type],
+            outputs=output
+        )
+    
+    return interface
 
 if __name__ == "__main__":
-    main()
+    interface = create_interface()
+    interface.launch()

prompts.yaml

@@ -186,7 +186,6 @@
 "planning":
   "initial_facts": |-
     Below I will present you a task.
-
     You will now build a comprehensive preparatory survey of which facts we have at our disposal and which ones we still need.
     To do so, you will have to read the task and identify things that must be discovered in order to successfully complete it.
     Don't make any assumptions. For each item, provide a thorough reasoning. Here is how you will structure this survey:
@@ -209,7 +208,6 @@
     Do not add anything else.
   "initial_plan": |-
     You are a world expert at making efficient plans to solve any task using a set of carefully crafted tools.
-
     Now for the given task, develop a step-by-step high-level plan taking into account the above inputs and list of facts.
     This plan should involve individual tasks based on the available tools, that if executed correctly will yield the correct answer.
     Do not skip steps, do not add any superfluous steps. Only write the high-level plan, DO NOT DETAIL INDIVIDUAL TOOL CALLS.
@@ -261,11 +259,9 @@
     ### 2. Facts that we have learned
     ### 3. Facts still to look up
     ### 4. Facts still to derive
-
     Now write your new list of facts below.
   "update_plan_pre_messages": |-
     You are a world expert at making efficient plans to solve any task using a set of carefully crafted tools.
-
     You have been given a task:
     ```
     {{task}}
@@ -279,7 +275,6 @@
     ```
     {{task}}
     ```
-
     You can leverage these tools:
     {%- for tool in tools.values() %}
     - {{ tool.name }}: {{ tool.description }}
@@ -319,7 +314,6 @@
     {{task}}
     ---
     You're helping your manager solve a wider task: so make sure to not provide a one-line answer, but give as much information as possible to give them a clear understanding of the answer.
-
     Your final_answer WILL HAVE to contain these parts:
     ### 1. Task outcome (short version):
     ### 2. Task outcome (extremely detailed version):
@@ -329,4 +323,4 @@
     And even if your task resolution is not successful, please return as much context as possible, so that your manager can act upon this feedback.
   "report": |-
     Here is the final answer from your managed agent '{{name}}':
-    {{final_answer}}
+    {{final_answer}}

tools/final_answer.py

@@ -11,4 +11,4 @@ class FinalAnswerTool(Tool):
         return answer
 
     def __init__(self, *args, **kwargs):
-        self.is_initialized = False
+        self.is_initialized = False 

tools/visit_webpage.py

@@ -3,6 +3,7 @@ from smolagents.tools import Tool
 import requests
 import markdownify
 import smolagents
+import re
 
 class VisitWebpageTool(Tool):
     name = "visit_webpage"
@@ -42,4 +43,4 @@ class VisitWebpageTool(Tool):
             return f"An unexpected error occurred: {str(e)}"
 
     def __init__(self, *args, **kwargs):
-        self.is_initialized = False
+        self.is_initialized = False 

tools/vuln_search.py

@@ -0,0 +1,62 @@
+import os
+from typing import Dict, Any, Optional, List
+import nvdlib
+from smolagents.tools import Tool
+
+class VulnerabilitySearchTool(Tool):
+    name = "vuln_search"
+    description = "Search for vulnerabilities in NVD (National Vulnerability Database)"
+    inputs = {
+        'query': {'type': 'str', 'description': 'Search term or CVE ID'},
+        'max_results': {'type': 'int', 'description': 'Maximum number of results', 'default': 5}
+    }
+    output_type = Dict[str, Any]
+
+    def __init__(self):
+        """Initialize NVD API connection"""
+        self.nvd_api_key = os.getenv('NVD_API_KEY')
+        
+        # Configure NVD
+        if self.nvd_api_key:
+            nvdlib.set_api_key(self.nvd_api_key)
+
+    def search_nvd(self, query: str, max_results: int) -> List[Dict[str, Any]]:
+        """Search vulnerabilities in NVD"""
+        try:
+            # If query looks like a CVE-ID, search directly
+            if query.startswith('CVE-'):
+                results = nvdlib.get_cve(query)
+                return [{
+                    'id': results.id,
+                    'description': results.descriptions[0].value,
+                    'severity': results.metrics.cvssMetricV31[0].cvssData.baseScore if results.metrics else None,
+                    'published': results.published,
+                    'references': [ref.url for ref in results.references]
+                }]
+            
+            # Otherwise, perform general search
+            results = nvdlib.searchCVE(
+                keyword=query,
+                limit=max_results
+            )
+            
+            return [{
+                'id': r.id,
+                'description': r.descriptions[0].value,
+                'severity': r.metrics.cvssMetricV31[0].cvssData.baseScore if r.metrics else None,
+                'published': r.published,
+                'references': [ref.url for ref in r.references]
+            } for r in results]
+        
+        except Exception as e:
+            return [{'error': f"Error in NVD search: {str(e)}"}]
+
+    def forward(self, query: str, max_results: int = 5) -> Dict[str, Any]:
+        """Process search in NVD"""
+        results = self.search_nvd(query, max_results)
+        
+        return {
+            'query': query,
+            'source': 'nvd',
+            'results': results
+        } 

tools/web_search.py

@@ -24,4 +24,4 @@ class DuckDuckGoSearchTool(Tool):
         if len(results) == 0:
             raise Exception("No results found! Try a less restrictive/shorter query.")
         postprocessed_results = [f"[{result['title']}]({result['href']})\n{result['body']}" for result in results]
-        return "## Search Results\n\n" + "\n\n".join(postprocessed_results)
+        return "## Search Results\n\n" + "\n\n".join(postprocessed_results)