daqc's picture
Upload 36 files
52b089b verified
import json
import os
from datetime import datetime
from typing import Dict, List, Optional
from smolagents import Tool
import plotly.graph_objects as go
import plotly.express as px
from jinja2 import Template
class ReportGeneratorTool(Tool):
"""Tool for generating interactive HTML vulnerability reports with charts."""
name = "generate_vulnerability_report"
description = "Generates an interactive HTML report with charts and vulnerability analysis. The report is generated from CVEDB search results."
inputs = {
"vulnerability_data": {
"type": "string",
"description": "Vulnerability data in JSON format",
},
"report_type": {
"type": "string",
"description": "Report type: 'cve' for a specific CVE or 'product' for a product",
}
}
output_type = "string"
def __init__(self):
super().__init__()
# Base HTML template
self.html_template = """
<!DOCTYPE html>
<html>
<head>
<title>Vulnerability Report</title>
<script src="https://cdn.plot.ly/plotly-latest.min.js"></script>
<style>
body { font-family: Arial, sans-serif; margin: 20px; }
.container { max-width: 1200px; margin: 0 auto; }
.header { text-align: center; margin-bottom: 30px; }
.section { margin-bottom: 40px; }
.chart { margin: 20px 0; }
.summary { background-color: #f5f5f5; padding: 20px; border-radius: 5px; }
.critical { color: #dc3545; }
.high { color: #fd7e14; }
.medium { color: #ffc107; }
.low { color: #28a745; }
</style>
</head>
<body>
<div class="container">
<div class="header">
<h1>Vulnerability Report</h1>
<p>Generated on {{ generation_date }}</p>
</div>
<div class="section">
<h2>Summary</h2>
<div class="summary">
{{ summary }}
</div>
</div>
<div class="section">
<h2>Severity Distribution (CVSS)</h2>
<div id="cvss_chart" class="chart"></div>
</div>
<div class="section">
<h2>Temporal Trend</h2>
<div id="timeline_chart" class="chart"></div>
</div>
<div class="section">
<h2>Vulnerability Details</h2>
{{ vulnerabilities_table }}
</div>
</div>
<script>
{{ plotly_js }}
</script>
</body>
</html>
"""
def forward(self, vulnerability_data: str, report_type: str) -> str:
"""Generates an HTML report with interactive charts from vulnerability data."""
try:
data = json.loads(vulnerability_data)
# Generate charts with Plotly
cvss_chart = self._generate_cvss_chart(data)
timeline_chart = self._generate_timeline_chart(data)
# Generate vulnerability table
vulnerabilities_table = self._generate_vulnerabilities_table(data)
# Generate summary
summary = self._generate_summary(data, report_type)
# Render template
template = Template(self.html_template)
html = template.render(
generation_date=datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
summary=summary,
vulnerabilities_table=vulnerabilities_table,
plotly_js=f"""
var cvssData = {cvss_chart};
var timelineData = {timeline_chart};
Plotly.newPlot('cvss_chart', cvssData.data, cvssData.layout);
Plotly.newPlot('timeline_chart', timelineData.data, timelineData.layout);
"""
)
# Save the report to the reports folder
# NOTE: Only saves to folder when running locally
# If deployed on a Hugging Face Space, it doesn't save files
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"vulnerability_report_{report_type}_{timestamp}.html"
# Create the reports folder if it doesn't exist
reports_dir = "reports"
if not os.path.exists(reports_dir):
os.makedirs(reports_dir)
# Save the file
filepath = os.path.join(reports_dir, filename)
with open(filepath, 'w', encoding='utf-8') as f:
f.write(html)
return f"Report generated and saved as: {filepath}\n\n{html}"
except Exception as e:
return f"Error generating report: {str(e)}"
def _generate_cvss_chart(self, data: Dict) -> Dict:
"""Generates a CVSS score distribution chart."""
if isinstance(data, list):
cvss_scores = [v.get('cvss', 0) for v in data if 'cvss' in v]
else:
cvss_scores = [data.get('cvss', 0)] if 'cvss' in data else []
fig = go.Figure()
fig.add_trace(go.Histogram(
x=cvss_scores,
nbinsx=10,
name='CVSS Scores'
))
fig.update_layout(
title='CVSS Score Distribution',
xaxis_title='CVSS Score',
yaxis_title='Number of Vulnerabilities',
showlegend=False
)
return fig.to_json()
def _generate_timeline_chart(self, data: Dict) -> Dict:
"""Generates a vulnerability timeline chart."""
if isinstance(data, list):
dates = [v.get('published_time', '') for v in data if 'published_time' in v]
else:
dates = [data.get('published_time', '')] if 'published_time' in data else []
# Convert dates to datetime format and count by month
from collections import Counter
from datetime import datetime
date_counts = Counter()
for date_str in dates:
try:
date = datetime.strptime(date_str, "%Y-%m-%dT%H:%M:%S")
month_key = date.strftime("%Y-%m")
date_counts[month_key] += 1
except:
continue
months = sorted(date_counts.keys())
counts = [date_counts[m] for m in months]
fig = go.Figure()
fig.add_trace(go.Scatter(
x=months,
y=counts,
mode='lines+markers',
name='Vulnerabilities'
))
fig.update_layout(
title='Vulnerability Timeline Trend',
xaxis_title='Month',
yaxis_title='Number of Vulnerabilities',
showlegend=False
)
return fig.to_json()
def _generate_vulnerabilities_table(self, data: Dict) -> str:
"""Generates an HTML table with vulnerability details."""
if isinstance(data, list):
vulnerabilities = data
else:
vulnerabilities = [data]
if not vulnerabilities:
return "<p>No vulnerability data available.</p>"
table_html = """
<table border="1" style="width: 100%; border-collapse: collapse;">
<thead>
<tr style="background-color: #f2f2f2;">
<th style="padding: 8px; text-align: left;">CVE ID</th>
<th style="padding: 8px; text-align: left;">CVSS Score</th>
<th style="padding: 8px; text-align: left;">EPSS Score</th>
<th style="padding: 8px; text-align: left;">Known Exploitable</th>
<th style="padding: 8px; text-align: left;">Publication Date</th>
<th style="padding: 8px; text-align: left;">Summary</th>
</tr>
</thead>
<tbody>
"""
for vuln in vulnerabilities:
cvss = vuln.get('cvss', 'Not available')
epss = vuln.get('epss', 'Not available')
kev = vuln.get('kev', False)
# Determine risk class based on CVSS score
risk_class = ""
if cvss != 'Not available' and isinstance(cvss, (int, float)):
if cvss >= 7.0:
risk_class = "critical"
elif cvss >= 4.0:
risk_class = "high"
else:
risk_class = "low"
table_html += f"""
<tr class="{risk_class}">
<td style="padding: 8px;">{vuln.get('id', 'Not available')}</td>
<td style="padding: 8px;">{cvss}</td>
<td style="padding: 8px;">{epss}</td>
<td style="padding: 8px;">{'Yes' if kev else 'No'}</td>
<td style="padding: 8px;">{vuln.get('published_time', 'Not available')}</td>
<td style="padding: 8px;">{vuln.get('summary', 'Not available')[:100]}...</td>
</tr>
"""
table_html += """
</tbody>
</table>
"""
return table_html
def _generate_summary(self, data: Dict, report_type: str) -> str:
"""Generates a summary of the vulnerability data."""
if isinstance(data, list):
total_vulns = len(data)
exploited = sum(1 for v in data if v.get('kev', False))
avg_cvss = sum(v.get('cvss', 0) for v in data if 'cvss' in v) / max(1, sum(1 for v in data if 'cvss' in v))
else:
total_vulns = 1
exploited = 1 if data.get('kev', False) else 0
avg_cvss = data.get('cvss', 0)
summary_html = f"""
<p>Found <strong>{total_vulns}</strong> vulnerabilities.</p>
<p>Exploited vulnerabilities: <strong>{exploited}</strong></p>
<p>Average CVSS Score: <strong>{avg_cvss:.2f}</strong></p>
<p>Publication date: <strong>{data.get('published_time', 'N/A')}</strong></p>
"""
return summary_html