File size: 3,925 Bytes
2e82565 8fe992b 2e82565 9b5b26a 2e82565 8c01ffb 2e82565 8c01ffb 2e82565 13d500a 2e82565 8c01ffb 2e82565 8c01ffb 2e82565 8c01ffb 2e82565 9b5b26a 2e82565 8fe992b 2e82565 9b5b26a 2e82565 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
#!/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 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 agents.coordinator_agent import search_vulnerabilities_for_software, get_vulnerability_details
from tools import utils
# Configure logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = utils.setup_logger("gradio_ui")
# 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 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 main():
"""Main entry point for the Gradio UI."""
args = create_parser().parse_args()
# Configure logging level
log_level = logging.DEBUG if args.verbose else logging.INFO
logging.basicConfig(level=log_level)
# Initialize the model
model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id=args.model,
custom_role_conversions=None,
)
# Initialize the agent con las herramientas ya existentes y las que hemos creado
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,
)
# Create Gradio UI
ui = GradioUI(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(),
},
)
if __name__ == "__main__":
main() |