|
|
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__() |
|
|
|
|
|
|
|
|
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) |
|
|
|
|
|
|
|
|
cvss_chart = self._generate_cvss_chart(data) |
|
|
timeline_chart = self._generate_timeline_chart(data) |
|
|
|
|
|
|
|
|
vulnerabilities_table = self._generate_vulnerabilities_table(data) |
|
|
|
|
|
|
|
|
summary = self._generate_summary(data, report_type) |
|
|
|
|
|
|
|
|
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); |
|
|
""" |
|
|
) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") |
|
|
filename = f"vulnerability_report_{report_type}_{timestamp}.html" |
|
|
|
|
|
|
|
|
reports_dir = "reports" |
|
|
if not os.path.exists(reports_dir): |
|
|
os.makedirs(reports_dir) |
|
|
|
|
|
|
|
|
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 [] |
|
|
|
|
|
|
|
|
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) |
|
|
|
|
|
|
|
|
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 |
