NLP/Models+K-Means/main/CVSS_Calculator.py

import math

from cvss import CVSS2
from cvss import CVSS3
from cvss import CVSS4

import re

import pandas as pd
import numpy as np

data1 = pd.DataFrame(pd.read_json(r"E:\pythonProject_open\data\SIR_dataset_processed.json"))
vecStr = data1["vectorString"]
impactScores = data1["impactScore"]
exploitabilityScores = data1["exploitabilityScore"]

print("----")
for i, j, k in zip(vecStr, impactScores, exploitabilityScores) :
    cvssVer = re.findall(':(.*?)/', i)
    impactScore = float(j)
    exploitabilityScore = float(k)

    if float(cvssVer[0]) == 2:
        cvss = CVSS2(i)
    elif 2 <= float(cvssVer[0]) < 4:
        cvss = CVSS3(i)
    else:
        cvss = CVSS4(i)
    cvss_baseScore = cvss.base_score
    print(cvss_baseScore)

    if impactScore <= 0:
        cvss_baseScore = 0
    elif 0 < impactScore + exploitabilityScore < 10:
        cvss_baseScore = math.ceil((impactScore + exploitabilityScore) * 10) / 10
    else:
        cvss_baseScore = 10
    print(f"baseScore:{cvss_baseScore}, impactScore:{impactScore}, exploitabilityScore:{exploitabilityScore}")