| 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 |
