db stuff complete (TODO: more testing)

vendors/debian/DebianAdvisory.py

@@ -102,7 +102,6 @@ class DebianAdvisory:
 
         dsa_id2string = '%03d' % dsa_id
 
-        print(dsa_id2string)
         flag = True
         while flag:
             try:
@@ -178,12 +177,11 @@ class DebianAdvisory:
         with open(suf) as su:
             for line in su:
                 sp_line = line.strip().split("->")
-                print(sp_line)
                 if re.compile(sp_line[0]).match(name):
                     return sp_line[1]
         return name
 
-
+    @staticmethod
     def fixDSAquirks(dsa_id, dsa_state):
         """
         TODO:
@@ -250,7 +248,6 @@ class DebianAdvisory:
                         "LOCAL-05/18/11"]
         elif dsa_id == 2286:
             new_names = ["phpmyadmin"]
-            print(str(dsa_id) + 'whatsapp??')
         elif dsa_id == 1977:
             new_names = ["python3.5"]
         elif (

vendors/debian/DebianModel.py

@@ -16,18 +16,24 @@ class DebianModel:
     """
     module_path = os.path.dirname(__file__)
 
+    """
+    TODO: Tables to manage.
+    """
+
+    dsatable = dict()
+    src2dsa = dict()
+    dsa2cve = dict()
+    cvetable = dict()
+    src2month = dict()
+    src2sloccount = dict()
+    src2pop = dict()
+    src2deps = dict()
+    pkg_with_cvss = dict()
+    src2sum = dict()
+
     def __init__(self, action, configfile=os.path.join(module_path, 'config_default.txt')):
         ## DBs to track
-        self.dsatable = dict()
-        self.src2dsa = dict()
-        self.dsa2cve = dict()
-        self.cvetable = dict()
-        self.src2month = dict()
-        self.src2sloccount = dict()
-        self.src2pop = dict()
-        self.src2deps = dict()
-        self.pkg_with_cvss = dict()
-        self.src2sum = dict()
+
 
         ## config
         self.configfile = configfile
@@ -36,21 +42,38 @@ class DebianModel:
             raise IOError('Cannot open configuration file: ')
         (self.state, self.err) = self.load_state()
 
-        client = MongoClient()
+        self.client = MongoClient()
 
         if action == 'update':
-            (self.dsatable, self.src2dsa, self.dsa2cve, self.cvetable, self.src2month, self.src2sloccount, self.src2pop, self.src2deps) = self.load_dbs()
-            self.update_dbs(self.dsatable, client, self.src2dsa, self.dsa2cve, self.src2month, self.cvetable, self.pkg_with_cvss)
-            self.save_DBs(self.dsatable, self.src2dsa, self.dsa2cve, self.cvetable, self.src2month, self.src2sloccount, self.src2pop, self.src2deps, self.src2sum)
+            self.load_dbs()
+            self.update_dbs()
+            self.store_dbs()
             self.save_state(self.state)
             # lstm.predict(src2month, src2sloccount, src2pop, src2deps)
+            """
+            with open('dsatable.txt', 'w') as file:
+                file.write(str(sorted(self.dsatable.keys(), key=lambda x: str(x).lower())))
+            with open('src2dsa.txt', 'w') as file:
+                file.write(str(sorted(self.src2dsa.keys(), key=lambda x: str(x).lower())))
+            with open('dsa2cve.txt', 'w') as file:
+                file.write(str(sorted(self.dsa2cve.keys(), key=lambda x: str(x).lower())))
+            with open('cvetable.txt', 'w') as file:
+                file.write(str(sorted(self.cvetable.keys(), key=lambda x: str(x).lower())))
+            with open('src2month.txt', 'w') as file:
+                file.write(str(sorted(self.src2month.keys(), key=lambda x: str(x).lower())))
+            with open('src2sloccount.txt', 'w') as file:
+                file.write(str(sorted(self.src2sloccount.keys(), key=lambda x: str(x).lower())))
+            with open('src2pop.txt', 'w') as file:
+                file.write(str(sorted(self.src2pop.keys(), key=lambda x: str(x).lower())))
+            with open('src2deps.txt', 'w') as file:
+                file.write(str(sorted(self.src2deps.keys(), key=lambda x: str(x).lower())))
+                """
         elif action == 'status':
-            (dsatable, src2dsa, dsa2cve, cvetable, src2month, src2sloccount, src2pop, src2deps) = self.load_dbs()
+            self.load_dbs()
             # aptsec_status(sys.argv[2])
         elif action == 'show':
-            (dsatable, src2dsa, dsa2cve, cvetable, src2month, src2sloccount, src2pop, src2deps) = self.load_dbs()
-            # src2sum = plot_all(src2month)
-            # save_DBs(dsatable, src2dsa, dsa2cve, cvetable, src2month, src2sloccount, src2pop, src2deps, src2sum)
+            self.load_dbs()
+            self.store_dbs()
         else:
             self.print_help(self)
 
@@ -59,53 +82,60 @@ class DebianModel:
         Loads the required databases into the model. Can either be implemented as read from file, or read from DB.
         Currently reading it from files in the cache folder.
         """
-        cache_dir = os.path.join( self.module_path, self.config['DIR']['cache_dir'])
-
-        tables = ['dsatable', 'src2dsa', 'dsa2cve', 'cvetable', 'src2deps', 'src2month', 'src2sloccount', 'src2pop']
-        result = []
-
-        for i in range(0, len(tables)):
-            try:
-                with open(os.path.join(cache_dir, tables[i])) as t:
-                    result.append(json.load(t))
-            except (IOError, ValueError):
-                print('Read cache ' + tables[i] + ' failed!! Maybe first run of the system?')
-                result.append(dict())
-
-        return tuple(result)
-
-    def save_DBs(self, dsatable, src2dsa, dsa2cve, cvetable, src2month, src2sloccount, src2pop, src2deps, src2sum):
+        self.dsatable = self.load_single_db_from_cache('dsatable')
+        self.src2dsa = self.load_single_db_from_cache('src2dsa')
+        self.dsa2cve = self.load_single_db_from_cache('dsa2cve')
+        self.cvetable = self.load_single_db_from_cache('cvetable')
+        self.src2deps = self.load_single_db_from_cache('src2deps')
+        self.src2month = self.load_single_db_from_cache('src2month')
+        self.src2sloccount = self.load_single_db_from_cache('src2sloccount')
+        self.src2pop = self.load_single_db_from_cache('src2pop')
+
+    def load_single_db_from_cache(self, file_name):
         cache_dir = os.path.join(self.module_path, self.config['DIR']['cache_dir'])
-
-        tables = ['dsatable', 'src2dsa', 'dsa2cve', 'cvetable', 'src2deps', 'src2sloccount', 'src2pop']
-
-        for i in range(0, len(tables)):
-            try:
-                with open(os.path.join(cache_dir, tables[i]), 'w') as t:
-                    json.dump(dsatable, t, default=self.converter)
-            except IOError:
-                print('write cache dsatable failed!! Fatal error')
-
-        cache_src2month = os.path.join(cache_dir, 'src2month')
+        try:
+            with open(os.path.join(cache_dir, file_name)) as f:
+                return json.load(f)
+        except (IOError, ValueError):
+            print('Read cache ' + file_name + ' failed!! Maybe first run of the system?')
+
+    def store_dbs(self):
+        self.store_db_single('dsatable', self.dsatable)
+        self.store_db_single('src2dsa', self.src2dsa)
+        self.store_db_single('dsa2cve', self.dsa2cve)
+        self.store_db_single('cvetable', self.cvetable)
+        self.store_db_single('src2deps', self.src2deps)
+        self.store_db_single('src2sloccount', self.src2sloccount)
+        self.store_db_single('src2pop', self.src2pop)
+
+        # src2month needs special handling
+        cache_src2month = os.path.join(self.module_path, self.config['DIR']['cache_dir'], 'src2month')
         int_list = dict()
 
-        for element in src2month:
-            for i in range(len(src2month[element])):
+        for element in self.src2month:
+            for i in range(len(self.src2month[element])):
                 if element in int_list:
-                    int_list[element].append(int(src2month[element][i]))
+                    int_list[element].append(int(self.src2month[element][i]))
                 else:
                     int_list[element] = []
-                    int_list[element].append(int(src2month[element][i]))
+                    int_list[element].append(int(self.src2month[element][i]))
         try:
             with open(cache_src2month, 'w') as fp:
                 json.dump(int_list, fp, default=self.converter)
         except IOError:
             print('write cache src2month failed!! Fatal error')
 
+    def store_db_single(self, file_name, db):
+        cache_dir = os.path.join(self.module_path, self.config['DIR']['cache_dir'])
+        try:
+            with open(os.path.join(cache_dir, file_name), 'w') as f:
+                json.dump(db, f, default=self.converter)
+        except (IOError, ValueError):
+            print('Read cache ' + file_name + ' failed!! Maybe first run of the system?')
+
     def save_state(self, state):
         """Save state, different from DBs in that we always need it"""
         state_file = os.path.join(self.module_path, self.config['DIR']['cache_dir'], 'state')
-
         try:
             with open(state_file, 'w') as sf:
                 json.dump(state, sf)
@@ -119,56 +149,51 @@ class DebianModel:
         if isinstance(o, np.float):
             return o.astype(int)
 
-    def update_dbs(self, dsatable, client, src2dsa, dsa2cve, src2month, cvetable, pkg_with_cvss):
+    def update_dbs(self):
         now = datetime.datetime.now()
         new_adv = DebianAdvisory.checkDSAs(self.state, self.config)
 
         for id in new_adv:
-            if id in dsatable:
+            if id in self.dsatable:
                 logging.info(self.state['vendor'] + ' advisory ' + id + ' already known.\n')
             else:
                 ## store advisory and parse it
-                dsatable[id] = new_adv[id]
-                self.updateCVETables(id, dsatable, self.state, src2dsa, dsa2cve, cvetable, client)
+                self.dsatable[id] = new_adv[id]
+                self.updateCVETables(id)
 
         # recompute all pkg statistics
-        for srcpkg in src2dsa:
-            self.processCVEs(srcpkg, now, src2dsa, dsa2cve, src2month, cvetable, pkg_with_cvss, self.config)
-        
-        return 0
+        for srcpkg in self.src2dsa:
+            self.processCVEs(srcpkg, now)
 
-    def updateCVETables(self, myid, dsatable, state, src2dsa, dsa2cve, cvetable, client):
+    def updateCVETables(self, myid):
 
-        logging.info('Updating vulnerability database with advisory ' + state['vendor'] + str(myid) + ' \n')
+        logging.info('Updating vulnerability database with advisory ' + self.state['vendor'] + str(myid) + ' \n')
 
-        adv = dsatable[myid]
+        adv = self.dsatable[myid]
         dsastats = DebianAdvisory.parseDSAhtml(adv)
 
         dsastats = DebianAdvisory.fixDSAquirks(myid, dsastats)
 
         for srcpkg in dsastats[0]:
-            if srcpkg in src2dsa:
-                src2dsa[srcpkg].append(myid)
+            if srcpkg in self.src2dsa:
+                self.src2dsa[srcpkg].append(myid)
             else:
-                src2dsa[srcpkg] = []
-                src2dsa[srcpkg].append(myid)
+                self.src2dsa[srcpkg] = []
+                self.src2dsa[srcpkg].append(myid)
 
-            dsa2cve[str(myid)] = dsastats[2]
+            self.dsa2cve[str(myid)] = dsastats[2]
 
 
         for cve_id in dsastats[2]:
             # No fetch CVE We use mongodb and cve-search
-            cve = CVEParse.fetchCVE(cve_id, client)
+            cve = CVEParse.fetchCVE(cve_id, self.client)
             cvestats = CVEParse.parseCVE(cve_id, cve)
             finaldate = cvestats[0]
 
             if cvestats[0] > dsastats[1] or cvestats[0] == 0:
                 finaldate = dsastats[1]
 
-            cvedata = (finaldate, dsastats[1] - finaldate, cvestats[1], cvestats[2], cvestats[3])
-            cvetable[cve_id] = cvedata
-
-        return cvetable
+            self.cvetable[cve_id] = (finaldate, dsastats[1] - finaldate, cvestats[1], cvestats[2], cvestats[3])
 
     @staticmethod
     def print_help():
@@ -204,34 +229,32 @@ class DebianModel:
 
         return state, err
 
-
-    def processCVEs(self, pkg, now, src2dsa, dsa2cve, src2month, cvetable, pkg_with_cvss, config):
+    def processCVEs(self, srcpkg, now):
         """
         compute and store MTBF, MTBR and Scores of each src pkg
         output: %src2mtbf
         (srcpkg=> ())
         """
+
         stats = [now, 0, 0, 0, 0, 0, 0]
-        mylambda = config['TRUST']['lambda']
         cvestats = dict()
-        logging.info('Processing package: ' + pkg + '.\n')
+        logging.info('Processing package: ' + srcpkg + '.\n')
 
         ## keep track of the number of low-medium-high severity vulnerabilities
         ## TODO see how cvss affects vulnerability prediction - if some packages show patterns
-        temp_cvss = 0.0
         with_cvss = dict()
 
         ## To eliminate duplicate cves
         haveseen = dict()
 
         ## cvestats = (date: number)
-        for dsa_id in src2dsa[pkg]:
-            for cve_id in dsa2cve[str(dsa_id)]:
+        for dsa_id in self.src2dsa[srcpkg]:
+            for cve_id in self.dsa2cve[str(dsa_id)]:
                 if cve_id in haveseen:
                     continue
                 else:
                     haveseen[cve_id] = 1
-                    tt = cvetable[cve_id][0]
+                    tt = self.cvetable[cve_id][0]
                     if tt in cvestats:
                         cvestats[tt] += 1
                     else:
@@ -241,11 +264,11 @@ class DebianModel:
         ## Date at the moment taken from CVE? - not sure.
 
         ## with_cvss = (date: number low, number med, number high)
-        for dsa_id in src2dsa[pkg]:
-            for cve_id in dsa2cve[str(dsa_id)]:
-                tt = cvetable[cve_id][0]
+        for dsa_id in self.src2dsa[srcpkg]:
+            for cve_id in self.dsa2cve[str(dsa_id)]:
+                tt = self.cvetable[cve_id][0]
                 try:
-                    temp_cvss = float(cvetable[cve_id][2])
+                    temp_cvss = float(self.cvetable[cve_id][2])
                 except TypeError:
                     print(cve_id)
                     continue
@@ -270,27 +293,20 @@ class DebianModel:
         if stats[1] < 1:
             return
 
-        prev_date = 0
-        weight = 0
-
         dates = sorted(cvestats, key=cvestats.get)
         try:
             stats[0] = dates[0]
         except IndexError:
-            print(pkg + str(dates))
+            print(srcpkg + str(dates))
             stats[0] = 0
 
         count = sum(cvestats.values())
 
-        print(pkg + ' ' + str(count))
-
-        #    pkg_with_cvss[pkg] = with_cvss
-        self.format_data(pkg, with_cvss, pkg_with_cvss, True)
+        self.format_data(srcpkg, with_cvss, self.pkg_with_cvss, True)
 
-        self.format_data(pkg, cvestats, src2month, False)
+        self.format_data(srcpkg, cvestats, self.src2month, False)
 
     def format_data(self, pkg, cvestats, src2month, cvss):
-
         x = []
         y = []
         monthyear = []