IsKEV.py

# Find out if a CVE is in the KEV catalog
# Input:    PT-all-scores-all-approaches.csv
#           known_exploited_vulnerabilities.csv
# Output:   PT-all-scores-all-approaches-with-kev.csv
import pandas as pd

def append_is_in_kev(input_file, kev_file, output_file):
    # Load the input CSV
    df = pd.read_csv(input_file)

    # Ensure the required column is present
    if "CVEId" not in df.columns:
        print("Error: The input file does not contain the required 'CVEId' column.")
        return

    # Load the KEV catalog file
    kev_df = pd.read_csv(kev_file)

    # Ensure the KEV catalog has the required column
    if "cveID" not in kev_df.columns:
        print("Error: The KEV catalog file does not contain the required 'cveID' column.")
        return

    # Create a set of KEV CVE IDs for quick lookup
    kev_cves = set(kev_df["cveID"])

    # Check if each CVE ID in the input file is in the KEV catalog
    df["is-in-kev"] = df["CVEId"].apply(lambda cve: cve in kev_cves)

    # Save the updated file
    df.to_csv(output_file, index=False)
    print(f"File updated with 'is-in-kev' column and saved to {output_file}")


# Run the function
input_file = "PT-all-scores-all-approaches.csv"
kev_file = "known_exploited_vulnerabilities.csv"
output_file = "PT-all-scores-all-approaches-with-kev.csv"
append_is_in_kev(input_file, kev_file, output_file)