epss_dataset_generator.py

# -*- coding: utf-8 -*-
"""
Created on Sat Aug 13 05:13:54 2022

@author: Arman Hossain
"""

# ! pip install pyexploitdb

import pandas as pd

from df_utilities import add_columns_to_df
from pyExploitDb import PyExploitDb
from utilities import read_csv_dataset
from nvd_downloader import download_data
from merging_16_22_databases import merge_all
# from epss_top_vendors_prods import generate_top_vendor_prod
import gzip
from PyQt5.QtWidgets import (
    QApplication)

class P_ven_prod():
    def __init__(self,cpe_series):
        self.cpe_series = cpe_series
        self.vendors = []
        self.prods = []
        self.vendor_df = pd.DataFrame()
        self.prod_df = pd.DataFrame()
        self.vendor_info = {}
        
    def perse(self,top_vendors,top_prods):
        '''
        this will add columns vendors from top_vendors and prods from top_prds with respect to all cve ids.
        example:
            id microsoft chrome
            1     5        6
            2     0        0 [because cve-2 is do not affect microsoft or chorme]

        '''
        
        lst = self.cpe_series.values.tolist()
        # vendor_lst1 = ['microsoft','adobe','ibm','hp','apache','apple','linux','oracle','opensuse','cisco','huawei','canonical','redhat','debian']
        # vendor_lst = ['microsoft','adobe','ibm','hp','apache','apple','google']
        df = pd.DataFrame()
        # for vendor in vendor_lst:
        #     df[vendor] = [0]*len(cpe_series)
        
        index = 0
        for item in lst:
            if(not isinstance(item, str)):
                index+=1
                continue
            cpe_lst = item.split('#arman#')
            
            for a_cpe in cpe_lst:
                a_cpe_sp = a_cpe.split(':')
                vendor= a_cpe_sp[3]
                prod = vendor+"_"+a_cpe_sp[4]
                if vendor not in self.vendors:
                    self.vendors.append(vendor)
                    if vendor in top_vendors.index:
                        self.vendor_df[vendor] = [0]*len(self.cpe_series)
                if vendor in top_vendors.index:
                    self.vendor_df[vendor][index] = top_vendors['count'][vendor]
                
                # if prod not in self.prods:
                #     self.prods.append(prod)
                #     if prod in top_prods.index:
                #         self.prod_df[prod] = [0]*len(self.cpe_series)
                # if prod in top_prods.index:
                #     self.prod_df[prod][index] = top_prods['count'][prod]
                
                
            index+=1
    def gen_vendor_info(self):
        
        '''
        will provide all the vendors and count in dataset cpe2uri
        microsoft 34234
        juniper 234
        ...
        '''
        
        lst = self.cpe_series.values.tolist()
        index = 0
        for item in lst:
            if(not isinstance(item, str)):
                index+=1
                continue
            cpe_lst = item.split('#arman#')
            for a_cpe in cpe_lst:
                a_cpe_sp = a_cpe.split(':')
                vendor= a_cpe_sp[3]
                
                if vendor not in self.vendor_info.keys():
                    self.vendor_info[vendor] = 0
                self.vendor_info[vendor] += 1
                
            index+=1
            
        return pd.DataFrame.from_dict(self.vendor_info, orient='index') 
    
    
    
    # def gen_vendor_info2(self):
    #     lst = self.cpe_series.values.tolist()
    #     index = 0
    #     for item in lst:
    #         if(not isinstance(item, str)):
    #             index+=1
    #             continue
    #         cpe_lst = item.split('#arman#')
    #         temp_list = []
    #         for a_cpe in cpe_lst:
    #             a_cpe_sp = a_cpe.split(':')
    #             vendor= a_cpe_sp[3]
                
    #             if vendor not in self.vendor_info.keys():
    #                 self.vendor_info[vendor] = 0
    #             if(vendor):
    #             self.vendor_info[vendor] += 1
                
    #         index+=1
            
    #     return pd.DataFrame.from_dict(self.vendor_info, orient='index')
    



    

def count_ref(ref_source):
    
    lst = ref_source.values.tolist()
    df = pd.DataFrame()
    df['ref_count'] = [0]*len(ref_source)
    
    index = 0
    for item in lst:
        if(not isinstance(item, str)): # if None
            index+=1
            continue
        refcnt = len(item.split('#arman#'))
        
        df['ref_count'][index] = refcnt
        index+=1
    
    return df
    

# def exploit():
#     expltlst = ['poc_code','weaponized']
#     return 0


def tags_poc_from_exploit_db(cve_id_series):
    
    # taglst = ['remote','dos','webapps','local'] # 'code_execution','memory_corruption' not in exploitdb
    taglst = [] # 'code_execution','memory_corruption' not in exploitdb
    pEdb = PyExploitDb()
    pEdb.debug = False
    pEdb.openFile()
    
    df = pd.DataFrame()
    # for tag in taglst:
    #     df[tag] = [0]*len(cve_id_series)
        
    # df['poc_code'] = [0]*len(cve_id_series) 
    df['code_link'] = [None]*len(cve_id_series) 
    
    index = 0
    for id in cve_id_series:    
        results = pEdb.searchCve(id)
        # results = pEdb.searchCve('CVE-2018-14592')
        # results = pEdb.searchCve('CVE-2021-1167')
        if (results==[] or type(results) != dict):
            index+=1
            continue
        else:
            # df['poc_code'][index] = 1
            df['code_link'][index] = results['exploit']
            # if("tag_"+results['platform'] not in taglst):
            #     taglst.append("tag_"+results['platform'])
            #     df["tag_"+results['platform']] = [0]*len(cve_id_series)
            # df["tag_"+results['platform']][index] = 1
                
            
        index+=1
        
        
    return df
        


def get_epss(data):
    data["epss"] = [0]*len(data)
    
    epss = ''
    with gzip.open('./downloaded/epss_scores-current.csv.gz') as f:
        epss = pd.read_csv(f)
    epss.columns = ['epss', 'percentile']
    for id in epss.index:
        if id in data.index:
            score = epss['epss'][id]
            data['epss'][id] = score
    return data

def get_exploitdb(total_df,isExcel=True):
    exploit = pd.read_csv('./downloaded/files_exploits.csv')
    # exploit.codes[0].split(';')[0].find('CVE')
    idx = 0
    for idd in exploit.codes:
        QApplication.processEvents()
        if type(idd) == float:
            idx+=1
            continue
        for code in idd.split(';'):
            if code.find('CVE') ==0:
                if code in total_df.index:
                    QApplication.processEvents()
                    if isExcel: 
                        # total_df['code_link'][code] = 'https://gitlab.com/exploit-database/exploitdb/-/tree/main/'+exploit.file[idx]
                        if type(total_df['code_link'][code]) != str:
                            total_df['code_link'][code] = 'https://exploit-db.com/exploits/'+str(exploit.id[idx])
                        else:
                            total_df['code_link'][code] += '; '+'https://exploit-db.com/exploits/'+str(exploit.id[idx])
                    total_df['code_link_count'][code] += 1
        idx+=1
        
    return total_df

def generate(data,top_vend,top_prod):
    
    total_df = pd.DataFrame()
    total_df['ID'] = data['ID']
    

    vendor_product_perser = P_ven_prod(data['cpe23Uri'])
    vendor_product_perser.perse(top_vend,top_prod)
    total_df = add_columns_to_df(total_df, vendor_product_perser.vendor_df)
    
    ref_cnt_df = count_ref(data['refsource'])
    total_df = add_columns_to_df(total_df, ref_cnt_df)
    
    # tag_poc_df = tags_poc_from_exploit_db(data['ID'])
    # total_df = add_columns_to_df(total_df, tag_poc_df)
    
    total_df.index = total_df["ID"]
    total_df.drop('ID', inplace=True, axis=1)

    total_df['code_link_count'] = [0]*len(total_df)    
    total_df = get_exploitdb(total_df,False)
    
    total_df = get_epss(total_df)
    # total_df.to_csv("166_22_desc_filtered_data.csv",index=False)
    
    return total_df

def generate_for_ui(data):
    
    total_df = pd.DataFrame()
    total_df['ID'] = data['ID']
    
    total_df = add_columns_to_df(total_df, data['description'],'description')
    total_df = add_columns_to_df(total_df, data['url'],'url')
    
    
    
    ref_cnt_df = count_ref(data['refsource'])
    total_df = add_columns_to_df(total_df, ref_cnt_df)
    
    total_df = add_columns_to_df(total_df, data['tags'],'tags')
    total_df = add_columns_to_df(total_df, data['cpe23Uri'],'cpeUri')
    total_df = add_columns_to_df(total_df, data['vectorString'],'vectorString_v31')
    total_df = add_columns_to_df(total_df, data['vectorString_v2'],"vectorString_v2")
    total_df = add_columns_to_df(total_df, data['baseScore'],'CVSS_v31')
    total_df = add_columns_to_df(total_df, data['baseScore_v2'],'CVSS_v2')
    total_df = add_columns_to_df(total_df, data['publishedDate'],'publishedDate')
    # total_df = add_columns_to_df(total_df, data['epss'],'epss')
    
    # data['CVSS31'] = mdata['baseScore']
    # data['CVSS2'] = mdata['baseScore_v2']
    
    # tag_poc_df = tags_poc_from_exploit_db(data['ID'])
    # total_df = add_columns_to_df(total_df, tag_poc_df)
    
    total_df.index = total_df["ID"]
    total_df.drop('ID', inplace=True, axis=1)
    
    total_df['code_link'] = [None]*len(total_df)
    total_df['code_link_count'] = [0]*len(total_df)
    total_df = get_exploitdb(total_df)
    
    
    total_df = get_epss(total_df)
    
    return total_df

def dataforexcel():
    data = read_csv_dataset('./data/2002_22_nvd.csv')
    excel_data = generate_for_ui(data)    
    excel_data.to_csv('./data/data_for_excel.csv')
    
def epss_thesholding(data,a=0.00949,b=0.019):
    
    # a = 0.00949
    # b = 0.019
    
    l = len(data[data['epss']<= a])/len(data)
    r = len(data[data['epss']>= b])/len(data)
    m = 1-l-r
    print('a: ',a,' b: ',b)
    print('epss: ','low rating-',l,'medium rating-',m,'critical rating-',r)
    
    for id in data.index:
        if data['epss'][id] <= a:
            data['epss'][id] = 1
        elif data['epss'][id] < b:
            data['epss'][id] = 2
        else:
            data['epss'][id] = 3
    return data

def cvss_thresholding(data,col_name = 'cvss'):
    
    one = len(data[data[col_name] < 0.1])/len(data)
    two = len(data[data[col_name] < 4.0])/len(data) - one
    three = len(data[data[col_name] < 7.0])/len(data) - two - one
    four = len(data[data[col_name] < 9.0])/len(data) - three - two - one
    five = len(data[data[col_name] >= 9.0])/len(data)
    
    print('one ',one, ' two ',two,' three ',three,' four ',four,' five ',five)

    
    for id in data.index:
        if data[col_name][id] < 0.1:
            data[col_name][id] = 1
        elif data[col_name][id] < 4.0:
            data[col_name][id] = 2
        elif data[col_name][id] < 7.0:
            data[col_name][id] = 3
        elif data[col_name][id] < 9.0:
            data[col_name][id] = 4
        else:
            data[col_name][id] = 5
    return data

def filter_data(data):
    data = data[data.attackVector != -1]
    data = data[data.attackComplexity != -1]
    data = data[data.privilegesRequired != -1]
    data = data[data.userInteraction != -1]
    data = data[data.scope != -1]
    data = data[data.confidentialityImpact != -1]
    data = data[data.integrityImpact != -1]
    data = data[data.availabilityImpact != -1]
    return data

def preprocess(a=0.00949,b=0.019,cvss = False):
    
    data = pd.read_csv('./data/data2/epss_dataset.csv',index_col = 0)
    
    if cvss:
        del data['epss'] # comment to include epss
    else:
        del data['cvss'] # comment to include cvss
        data = data[data.epss!=0]
    
    # data['cvss'] = data['cvss'].round()
    # data['epss'].astype(float)
    # data['epss'] = data['epss'].round(decimals=2)
    
    
    data.dropna(inplace=True) # cvss contains null value
    data = filter_data(data)
    if cvss:
        data = cvss_thresholding(data)
        data.to_csv('./data/data2/cvss_dataset_final_epss.csv')
    else:
        data2 = epss_thesholding(data,a,b)
        data2.to_csv('./data/data2/epss_dataset_final_epss.csv')
        # data.to_csv('./data/data2/cvss_epss.csv')
  


def valid_index(lst):
    gotZero = False
    first = 0
    for i in range (0,len(lst)):
        if lst[i]==-1:
            if gotZero:
                return first,i
        elif not gotZero:
            gotZero = True
            first = i
    return first, len(lst)


def collectepsstraindata():
    mdata = read_csv_dataset('./data/2016_22_nvd.csv')
    # mdata.to_csv('./data/2016_22m_nvd.csv')
    # generate_top_vendor_prod()
    
    top_vend = pd.read_csv('./data/top_vendors.csv',index_col=0)
    top_prod = pd.read_csv('./data/top_products.csv',index_col=0)
    
    
    data = generate(mdata,top_vend,top_prod)
    arr = ['attackVector','attackComplexity','privilegesRequired','userInteraction','scope',
           'confidentialityImpact','integrityImpact','availabilityImpact']
    
    for i in arr:
        codes, uniques = pd.factorize(mdata[i])
        data[i] = codes
    
    mdata.index = mdata["ID"]
    mdata.drop('ID', inplace=True, axis=1)
    
    topics_data = pd.read_csv("./data/epss_topics_16_22.csv")
    data = add_columns_to_df(data,topics_data)
    
    dat = mdata['baseScore']
    data['cvss'] = dat #list(dat.round())
    
    
    
    dat = data['epss'].astype(float)
    data.drop('epss', inplace=True, axis=1)
    data['epss'] = dat #list(dat.round(decimals=2))
    data.to_csv('./data/epss_dataset.csv.csv')
    
    return data
    
#%% main
if __name__ == "__main__":
    # for EPSS model generation data
    
    affected_years = download_data()
    data = merge_all('./downloaded',2016,2050) # affected year should be used and merging with original database inorder to reduce time
    data.to_csv('./data/2016_22_nvd.csv',index=False)
    mdata = read_csv_dataset('./data/2016_22_nvd.csv')
    mdata.to_csv('./data/2016_22m_nvd.csv')
    
    data = collectepsstraindata()
    # end
    
    
            
# saving data

    import pickle
    with open('all_data.pickle', 'wb') as f:
        pickle.dump(data, f)
    
    data.to_csv('./data/epss_16_22_without_topic.csv')
    data2.to_csv('latest2.csv',index = False)
    
    
    data.to_csv("16_22_desc_filtered_data.csv",index=False)
    
    vendor_product_perser = P_ven_prod(data['cpe23Uri'])
    
    vendor_product_perser.gen_vendor_info()
    
    vendor_inf = vendor_product_perser.vendor_info
    
    
    df = pd.DataFrame.from_dict(vendor_inf, orient='index') 
    df.to_csv('vendor.csv')