import argparse
import requests

def get_cve_data(cpe):
    base_url = "https://services.nvd.nist.gov/rest/json/cves/2.0"
    query_params = {"cpeName": cpe}
    response = requests.get(base_url, params=query_params)

    if response.status_code == 200:
        cve_data = response.json()
        if "vulnerabilities" in cve_data:
            return cve_data["vulnerabilities"]
        else:
            print("No vulnerabilities found in the response.")
            return []
    else:
        print(f"Error in HTTP request: {response.status_code}")
        return []

def get_cve_id(cve):
    try:
        return cve["cve"]["id"]
    except (KeyError, TypeError, ValueError):
        return "N/A"

def get_cve_metric_version(cve):
    if "metrics" in cve["cve"]:
        if "cvssMetricV31" in cve["cve"]["metrics"]:
            return "3.1"
        elif "cvssMetricV30" in cve["cve"]["metrics"]:
            return "3.0"
        elif "cvssMetricV2" in cve["cve"]["metrics"]:
            return "2.0"
    return "N/A"

def get_cve_score(cve):
    try:
        metrics = cve["cve"]["metrics"]
        if "cvssMetricV31" in metrics:
            return float(metrics["cvssMetricV31"][0]["cvssData"]["baseScore"])
        elif "cvssMetricV30" in metrics:
            return float(metrics["cvssMetricV30"][0]["cvssData"]["baseScore"])
        elif "cvssMetricV2" in metrics:
            return float(metrics["cvssMetricV2"][0]["cvssData"]["baseScore"])
        else:
            return 0.0  # 如果没有找到有效的 CVSS 版本，返回默认值 0.0
    except (KeyError, TypeError, ValueError):
        return 0.0  # 如果提取分数失败，返回默认值 0.0

def get_cve_severity(cve):
    try:
        metrics = cve["cve"]["metrics"]
        if "cvssMetricV31" in metrics:
            return metrics["cvssMetricV31"][0]["cvssData"]["baseSeverity"]
        elif "cvssMetricV30" in metrics:
            return metrics["cvssMetricV30"][0]["cvssData"]["baseSeverity"]
        elif "cvssMetricV2" in metrics:
            return metrics["cvssMetricV2"][0]["baseSeverity"]
        else:
            return "N/A"  # 如果没有找到有效的 CVSS 版本，返回默认值 "N/A"
    except (KeyError, TypeError, ValueError):
        return "N/A"  # 如果提取严重性失败，返回默认值 "N/A"

def main():
    parser = argparse.ArgumentParser(description="Get and sort CVEs from a CPE")
    parser.add_argument("-c", "--cpe", required=True, help="CPE from which to retrieve CVEs")
    args = parser.parse_args()

    cve_data = get_cve_data(args.cpe)

    if cve_data:
        sorted_cve = sorted(cve_data, key=get_cve_score, reverse=True)
        i = 1
        for cve in sorted_cve:
            cve_id = get_cve_id(cve)
            score = get_cve_score(cve)
            severity = get_cve_severity(cve)
            print(f"[{i}] ID: {cve_id}, Score: {score}, Severity: {severity}")
            i += 1
    else:
        print("No CVE data found.")

if __name__ == "__main__":
    main()