269 lines
8.4 KiB
Python
269 lines
8.4 KiB
Python
import argparse
|
||
import requests
|
||
from bs4 import BeautifulSoup
|
||
from packaging.version import Version, InvalidVersion
|
||
import sys
|
||
from reportlab.lib.pagesizes import letter
|
||
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
|
||
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
|
||
from colorama import Fore, Style, init
|
||
from tqdm import tqdm
|
||
import html
|
||
import os
|
||
|
||
|
||
init(autoreset=True) # 初始化colorama，并在每次打印后自动重置颜色
|
||
|
||
|
||
def fetch_html(url: str) -> str:
|
||
try:
|
||
response = requests.get(url)
|
||
response.raise_for_status()
|
||
return response.text
|
||
except requests.RequestException as e:
|
||
print(f"Error fetching {url}: {e}")
|
||
return ""
|
||
|
||
|
||
def parse_html(html: str) -> list:
|
||
soup = BeautifulSoup(html, "html.parser")
|
||
table = soup.find("table", id="sortable-table")
|
||
if not table:
|
||
return []
|
||
|
||
rows = table.find_all("tr", class_="vue--table__row")
|
||
results = []
|
||
for row in rows:
|
||
info = {}
|
||
link = row.find("a")
|
||
chip = row.find("span", class_="vue--chip__value")
|
||
if link and chip:
|
||
info["link"] = link.get_text(strip=True)
|
||
info["chip"] = chip.get_text(strip=True)
|
||
results.append(info)
|
||
return results
|
||
|
||
|
||
def load_requirements(file_path: str) -> list:
|
||
requirements = []
|
||
try:
|
||
with open(file_path, "r") as file:
|
||
for line in file:
|
||
line = line.strip()
|
||
if line and not line.startswith("#"):
|
||
requirements.append(line)
|
||
except FileNotFoundError:
|
||
print(f"Error: File {file_path} not found.")
|
||
sys.exit(1)
|
||
return requirements
|
||
|
||
|
||
def version_in_range(version, range_str: str) -> bool:
|
||
if version is not None:
|
||
try:
|
||
v = Version(version)
|
||
except InvalidVersion:
|
||
return False
|
||
else:
|
||
if range_str[-2] == ",":
|
||
return True
|
||
|
||
ranges = range_str.split(",")
|
||
for range_part in ranges:
|
||
range_part = range_part.strip("[]()")
|
||
if range_part:
|
||
try:
|
||
if range_part.endswith(")"):
|
||
upper = Version(range_part[:-1])
|
||
if v >= upper:
|
||
return False
|
||
elif range_part.startswith("["):
|
||
lower = Version(range_part[1:])
|
||
if v < lower:
|
||
return False
|
||
except InvalidVersion:
|
||
return False
|
||
return True
|
||
|
||
|
||
def check_vulnerabilities(requirements: list, base_url: str) -> str:
|
||
results = []
|
||
for req in tqdm(requirements, desc="Checking vulnerabilities", unit="dependency"):
|
||
version = ""
|
||
if "==" in req:
|
||
package_name, version = req.split("==")
|
||
else:
|
||
package_name, version = req, None
|
||
url = f"{base_url}{package_name}"
|
||
# print(f"\nFetching data for {package_name} from {url}")
|
||
html_content = fetch_html(url)
|
||
if html_content:
|
||
extracted_data = parse_html(html_content)
|
||
if extracted_data:
|
||
relevant_vulns = []
|
||
for vuln in extracted_data:
|
||
if version_in_range(version, vuln["chip"]):
|
||
relevant_vulns.append(vuln)
|
||
if relevant_vulns:
|
||
result = f"Vulnerabilities found for {package_name}:\n"
|
||
for vuln in relevant_vulns:
|
||
result += f" - {vuln['link']}\n"
|
||
results.append(result)
|
||
return "\n".join(results)
|
||
|
||
|
||
def save_to_file(output_path: str, data: str):
|
||
if output_path.endswith(".html"):
|
||
save_as_html(output_path, data)
|
||
elif output_path.endswith(".pdf"):
|
||
save_as_pdf(output_path, data)
|
||
elif output_path.endswith(".md"):
|
||
save_as_markdown(output_path, data)
|
||
else:
|
||
save_as_txt(output_path, data)
|
||
|
||
|
||
def save_as_html(output_path: str, data: str):
|
||
escaped_data = html.escape(data)
|
||
html_content = f"""
|
||
<html>
|
||
<head>
|
||
<meta charset="UTF-8">
|
||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||
<link rel="icon" href="https://s2.loli.net/2024/05/30/WDc6MekjbuCU9Qo.png">
|
||
<title>Vulnerability Report</title>
|
||
<style>
|
||
body {{
|
||
font-family: Arial, sans-serif;
|
||
background-image: url('https://s2.loli.net/2024/05/30/85Mv7leB2IRWNp6.jpg');
|
||
background-size: cover;
|
||
color: #333;
|
||
margin: 0;
|
||
padding: 0;
|
||
display: flex;
|
||
justify-content: center;
|
||
align-items: center;
|
||
height: 100vh;
|
||
}}
|
||
.container {{
|
||
background: rgba(255, 255, 255, 0.8);
|
||
border-radius: 10px;
|
||
padding: 20px;
|
||
box-shadow: 0 0 10px rgba(0, 0, 0, 0.1);
|
||
max-width: 800px;
|
||
width: 100%;
|
||
margin: 20px;
|
||
overflow-y: auto;
|
||
max-height: 90vh;
|
||
}}
|
||
.title {{
|
||
font-size: 24px;
|
||
font-weight: bold;
|
||
text-align: center;
|
||
margin-bottom: 20px;
|
||
}}
|
||
pre {{
|
||
white-space: pre-wrap;
|
||
word-wrap: break-word;
|
||
font-size: 14px;
|
||
line-height: 1.5;
|
||
color: #333;
|
||
background: #f4f4f4;
|
||
padding: 10px;
|
||
border-radius: 5px;
|
||
border: 1px solid #ddd;
|
||
overflow: auto;
|
||
font-weight: bold;
|
||
}}
|
||
</style>
|
||
</head>
|
||
<body>
|
||
<div class="container">
|
||
<div class="title">Vulnerability Report</div>
|
||
<pre>{escaped_data}</pre>
|
||
</div>
|
||
</body>
|
||
</html>
|
||
"""
|
||
with open(output_path, "w", encoding="utf-8") as file:
|
||
file.write(html_content)
|
||
|
||
|
||
def save_as_pdf(output_path: str, data: str):
|
||
doc = SimpleDocTemplate(output_path, pagesize=letter)
|
||
story = []
|
||
styles = getSampleStyleSheet()
|
||
|
||
# Add the title centered
|
||
title_style = ParagraphStyle(
|
||
"Title",
|
||
parent=styles["Title"],
|
||
alignment=1, # Center alignment
|
||
fontSize=24,
|
||
leading=28,
|
||
spaceAfter=20,
|
||
fontName="Helvetica-Bold",
|
||
)
|
||
title = Paragraph("Vulnerability Report", title_style)
|
||
story.append(title)
|
||
|
||
# Normal body text style
|
||
normal_style = ParagraphStyle(
|
||
"BodyText", parent=styles["BodyText"], fontSize=12, leading=15, spaceAfter=12
|
||
)
|
||
|
||
# Add the vulnerability details
|
||
for line in data.split("\n"):
|
||
if line.strip(): # Skip empty lines
|
||
story.append(Paragraph(line, normal_style))
|
||
|
||
doc.build(story)
|
||
|
||
|
||
def save_as_markdown(output_path: str, data: str):
|
||
with open(output_path, "w") as file:
|
||
file.write("## Vulnerability Report: \n\n")
|
||
file.write(data)
|
||
|
||
|
||
def save_as_txt(output_path: str, data: str):
|
||
with open(output_path, "w") as file:
|
||
file.write("Vulnerability Report: \n\n")
|
||
file.write(data)
|
||
|
||
|
||
def print_separator(title, char="-", length=50, padding=2):
|
||
print(f"{title:^{length + 4*padding}}") # 居中打印标题，两侧各有padding个空格
|
||
print(char * (length + 2 * padding)) # 打印分割线，两侧各有padding个字符的空格
|
||
|
||
|
||
def modify_file_name(file_path: str) -> str:
|
||
"""
|
||
Modify the file name by adding '-re' before the file extension.
|
||
|
||
Args:
|
||
file_path (str): The original file path.
|
||
|
||
Returns:
|
||
str: The modified file path.
|
||
"""
|
||
directory, file_name = os.path.split(file_path)
|
||
name, ext = os.path.splitext(file_name)
|
||
new_file_name = f"{name}-re{ext}"
|
||
new_file_path = os.path.join(directory, new_file_name)
|
||
return new_file_path
|
||
|
||
|
||
def requirement_detection(requirement_path, output_path=None):
|
||
base_url = "https://security.snyk.io/package/pip/"
|
||
requirements = load_requirements(requirement_path)
|
||
results = check_vulnerabilities(requirements, base_url)
|
||
if output_path is not None:
|
||
new_path = modify_file_name(output_path)
|
||
save_to_file(new_path, results)
|
||
print(f"Vulnerability scan complete. Results saved to {output_path}")
|
||
print(f"Requirements scan complete. Results saved to {new_path}")
|
||
else:
|
||
print_separator("\nVulnerability Report", "=", 40, 5)
|
||
print(results)
|