File size: 5,247 Bytes
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 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
"""
Utility functions for the Vulnerability Intelligence Agent (VIA).
"""
import os
import logging
import json
import datetime
from typing import Dict, List, Any, Optional, Union
def setup_logger(name: str) -> logging.Logger:
"""
Set up a logger with the specified name.
Args:
name: Name of the logger
Returns:
Configured logger instance
"""
logger = logging.getLogger(name)
if not logger.handlers:
handler = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)
return logger
def merge_vulnerability_data(results: List[Dict[str, Any]]) -> Dict[str, Any]:
"""
Merge vulnerability data from multiple sources.
Args:
results: List of dictionaries with vulnerability data from different sources
Returns:
Merged dictionary with all vulnerabilities
"""
# Simple implementation - in a real system, this would be more sophisticated
if not results:
return {"software": "", "version": "", "vulnerabilities": []}
merged = {
"software": results[0].get("software", ""),
"version": results[0].get("version", ""),
"vulnerabilities": []
}
# Simple merge - just combine all vulnerabilities
for result in results:
if "vulnerabilities" in result:
merged["vulnerabilities"].extend(result["vulnerabilities"])
return merged
def save_report(data: Dict[str, Any], filename: str, report_dir: str = "reports") -> str:
"""
Save vulnerability data to a JSON file.
Args:
data: Vulnerability data to save
filename: Base filename (without extension)
report_dir: Directory to save the report in
Returns:
Path to the saved JSON file
"""
# Ensure the reports directory exists
os.makedirs(report_dir, exist_ok=True)
# Add timestamp to filename to avoid overwriting
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
json_filename = f"{filename}_{timestamp}.json"
json_path = os.path.join(report_dir, json_filename)
# Save the data to a JSON file
with open(json_path, 'w') as f:
json.dump(data, f, indent=2)
return json_path
def generate_markdown_report(data: Dict[str, Any], filename: str, report_dir: str = "reports") -> str:
"""
Generate a Markdown report from vulnerability data.
Args:
data: Vulnerability data
filename: Base filename (without extension)
report_dir: Directory to save the report in
Returns:
Path to the generated Markdown file
"""
# Ensure the reports directory exists
os.makedirs(report_dir, exist_ok=True)
# Add timestamp to filename to avoid overwriting
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
md_filename = f"{filename}_{timestamp}.md"
md_path = os.path.join(report_dir, md_filename)
with open(md_path, 'w') as f:
# Write title
f.write(f"# Vulnerability Report: {data['software']} {data['version']}\n\n")
f.write(f"*Generated on: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*\n\n")
# Write summary
f.write("## Summary\n\n")
vuln_count = len(data.get('vulnerabilities', []))
f.write(f"Found **{vuln_count}** vulnerabilities for {data['software']} {data['version']}.\n\n")
# Write vulnerabilities
if vuln_count > 0:
f.write("## Vulnerabilities\n\n")
for i, vuln in enumerate(data['vulnerabilities'], 1):
f.write(f"### {i}. {vuln.get('id', 'Unknown ID')}\n\n")
f.write(f"**Severity:** {vuln.get('severity', 'Unknown')}")
if 'cvss' in vuln:
f.write(f" (CVSS: {vuln['cvss']})")
f.write("\n\n")
f.write(f"**Description:** {vuln.get('description', 'No description available.')}\n\n")
if 'date' in vuln:
f.write(f"**Published:** {vuln['date']}\n\n")
if 'recommendation' in vuln:
f.write(f"**Recommendation:** {vuln['recommendation']}\n\n")
if 'source' in vuln:
f.write(f"**Source:** [{vuln['source']}]({vuln['source']})\n\n")
f.write("---\n\n")
else:
f.write("## No vulnerabilities found\n\n")
f.write("No known vulnerabilities were found for this software and version.\n\n")
# Write footer
f.write("## References\n\n")
f.write("- [CVE (Common Vulnerabilities and Exposures)](https://cve.mitre.org/)\n")
f.write("- [NVD (National Vulnerability Database)](https://nvd.nist.gov/)\n")
f.write("- [CISA Known Exploited Vulnerabilities Catalog](https://www.cisa.gov/known-exploited-vulnerabilities-catalog)\n")
f.write("- [CWE (Common Weakness Enumeration)](https://cwe.mitre.org/)\n")
return md_path |