161 lines
5.5 KiB
Python
161 lines
5.5 KiB
Python
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:", "<b>WARNING:</b>"), 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("<html><head><title>Vulnerability Report</title></head><body>\n")
|
|
file.write("<h1>Vulnerability Report</h1>\n")
|
|
for result in results:
|
|
file.write(f"<p>{result}</p>\n")
|
|
file.write("</body></html>")
|
|
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()
|
粤公网安备44010602013451号