#!/usr/bin/env python
import json
import colorama
import sys
from termcolor import colored
from tabulate import tabulate

'''
Quick script to format the result of a NPM Audit JSON report:
    npm audit --json > npm-audit.json
    python generate-report-npm.py npm-audit.json

Dependencies:
    pip install tabulate colorama termcolor
'''


def get_color(severity):
    if severity in ["HIGH", "CRITICAL"]:
        severity_color = "red"
    elif severity in ["MEDIUM", "MODERATE"]:
        severity_color = "yellow"
    elif severity == "LOW":
        severity_color = "cyan"
    else:
        severity_color = "white"
    return severity_color


report = sys.argv[1]
colorama.init()
with open(report) as f:
    content = f.read()
v = json.loads(content)
table_rows = []
cve_ghsa_list = []
for vid in v["vulnerabilities"]:
    vias = v["vulnerabilities"][vid]["via"]
    for via in vias:
        if isinstance(via, str):
            mod_name = via
            severity = v["vulnerabilities"][vid]["severity"].upper()
            cve_ghsa = "NA"
        else:        
            mod_name = via["name"]
            severity = via["severity"].upper()
            cve_ghsa = via["url"].split("/")[-1]
            cve_ghsa_list.append(cve_ghsa)
        vuln_range = v["vulnerabilities"][vid]["range"]
        severity_color = get_color(severity)
        severity = colored(f"{severity}", severity_color, attrs=["bold"])
        table_rows.append([severity, mod_name, vuln_range, cve_ghsa])
tempo = []
table_rows_depuplicated = []
for row in table_rows:
    k = ""
    for val in row:
        k += val.upper().strip()
    if k not in tempo:
        tempo.append(k)
        table_rows_depuplicated.append(row)
table_rows_depuplicated.sort()
vulns_metadata = v["metadata"]
table_rows.clear()
levels = ["CRITICAL", "HIGH", "MODERATE", "LOW", "INFO"]
vuln_count = 0
for level in levels:
    vuln_count += vulns_metadata["vulnerabilities"][level.lower()]
    severity = colored(f"{level}", get_color(level), attrs=["bold"])
    table_rows.append([severity, vulns_metadata["vulnerabilities"][level.lower()]])
    if level == levels[-1]:
        table_rows.append(["TOTAL", vuln_count])
print(colored("[+] Vulnerabilities:", "yellow"))
table_headers = ["Severity", "Module name", "Affected versions range", "CVE/GHSA"]
print(tabulate(table_rows_depuplicated, headers=table_headers))
print("")
print(colored("[+] Vulnerabilities severity distribution:", "yellow"))
table_headers = ["Severity", "Vulnerability count"]
print(tabulate(table_rows, headers=table_headers))
print("")
print(colored(f"[+] List of CVE/GHSA ({len(cve_ghsa_list)}):", "yellow"))
print("\n".join(cve_ghsa_list))
