import argparse
import os
import re
import sys
from packaging import version
from packaging.specifiers import SpecifierSet
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
def load_vulnerable_packages(filename):
"""从文件加载有漏洞的包信息"""
with open(filename, "r", encoding="utf-8") as file:
content = file.read()
vulnerabilities = {}
blocks = content.split("--------------------------------------------------")
for block in blocks:
name_match = re.search(r"Package Name: (.+)", block)
range_match = re.search(r"Version Ranges: (.+)", block)
if name_match and range_match:
package_name = name_match.group(1).strip()
version_range = range_match.group(1).strip()
version_range = ",".join(
[part.strip() for part in version_range.split(",")]
)
vulnerabilities[package_name] = SpecifierSet(version_range)
return vulnerabilities
def load_requirements(filename):
"""从文件加载项目的依赖信息"""
with open(filename, "r", encoding="utf-8") as file:
lines = file.readlines()
requirements = {}
for line in lines:
if "==" in line:
package_name, package_version = line.strip().split("==")
requirements[package_name] = package_version
return requirements
def output_pdf(results, file_name):
doc = SimpleDocTemplate(file_name, pagesize=letter)
story = []
styles = getSampleStyleSheet()
# Custom styles
title_style = styles["Title"]
title_style.alignment = 1 # Center alignment
warning_style = ParagraphStyle(
"WarningStyle", parent=styles["BodyText"], fontName="Helvetica-Bold"
)
normal_style = styles["BodyText"]
# Add the title
title = Paragraph("Vulnerability Report", title_style)
story.append(title)
story.append(Spacer(1, 20)) # Space after title
# Iterate through results to add entries
for result in results:
if "WARNING:" in result:
# Add warning text in bold
entry = Paragraph(
result.replace("WARNING:", "WARNING:"), warning_style
)
else:
# Add normal text
entry = Paragraph(result, normal_style)
story.append(entry)
story.append(Spacer(1, 12)) # Space between entries
doc.build(story)
def output_results(filename, results, format_type):
"""根据指定的格式输出结果"""
output_dir = os.path.dirname(filename)
if not os.path.exists(output_dir):
os.makedirs(output_dir)
with open(filename, "w", encoding="utf-8") as file:
if format_type == "html":
file.write("Vulnerability Report\n")
file.write("Vulnerability Report
\n")
for result in results:
file.write(f"{result}
\n")
file.write("")
elif format_type == "md":
file.write("# Vulnerability Report\n")
for result in results:
file.write(f"* {result}\n")
elif format_type == "pdf":
output_pdf(results, filename)
else: # 默认为txt
for result in results:
file.write(f"{result}\n")
def check_vulnerabilities(requirements, vulnerabilities, output_file):
"""检查依赖项是否存在已知漏洞，并输出结果"""
results_warning = [] # 存储有漏洞的依赖
results_ok = [] # 存储没有漏洞的依赖
for req_name, req_version in requirements.items():
if req_name in vulnerabilities:
spec = vulnerabilities[req_name]
if version.parse(req_version) in spec:
results_warning.append(
f"WARNING: {req_name}=={req_version} is vulnerable!"
)
else:
results_ok.append(f"OK: {req_name}=={req_version} is not affected.")
else:
results_ok.append(
f"OK: {req_name} not found in the vulnerability database."
)
# 合并结果，先输出所有警告，然后输出所有正常情况
results = results_warning + results_ok
if output_file:
filename, ext = os.path.splitext(output_file)
output_format = ext[1:] if ext[1:] else "txt"
if output_format not in ["txt", "md", "html", "pdf"]:
print("Warning: Invalid file format specified. Defaulting to TXT format.")
output_format = "txt" # 确保使用默认格式
output_file = filename + ".txt"
output_results(output_file, results, output_format)
else:
print("\n".join(results))
def main():
parser = argparse.ArgumentParser(
description="Check project dependencies for vulnerabilities."
)
parser.add_argument(
"vulnerabilities_file", help="Path to the file containing vulnerability data"
)
parser.add_argument(
"requirements_file", help="Path to the requirements file of the project"
)
parser.add_argument(
"-o",
"--output",
help="Output file path with extension, e.g., './output/report.txt'",
)
args = parser.parse_args()
vulnerabilities = load_vulnerable_packages(args.vulnerabilities_file)
requirements = load_requirements(args.requirements_file)
check_vulnerabilities(requirements, vulnerabilities, args.output)
if __name__ == "__main__":
main()