removed weights from indexes for mongodb 5.0 compatibility and black …

…formatting
adulau · Aug 31, 2021 · 094f21f · 094f21f
1 parent b6f3ce9
commit 094f21f
Show file tree

Hide file tree

Showing 51 changed files with 589 additions and 491 deletions.
diff --git a/bin/cve_refs.py b/bin/cve_refs.py
@@ -11,6 +11,7 @@
 import os
 import sys
 import argparse
+
 runPath = os.path.dirname(os.path.realpath(__file__))
 sys.path.append(os.path.join(runPath, ".."))
 
@@ -23,30 +24,34 @@
 except:
     sys.exit(1)
 
-argparser = argparse.ArgumentParser(description='Lookup the NIST ref database')
-argparser.add_argument('-c', help='CVE id to lookup', default=False)
-argparser.add_argument('-u', action='store_true', help='Enable URL expansion', default=False)
-argparser.add_argument('-v', action='store_true', help='verbose output', default=False)
+argparser = argparse.ArgumentParser(description="Lookup the NIST ref database")
+argparser.add_argument("-c", help="CVE id to lookup", default=False)
+argparser.add_argument(
+    "-u", action="store_true", help="Enable URL expansion", default=False
+)
+argparser.add_argument("-v", action="store_true", help="verbose output", default=False)
 args = argparser.parse_args()
 
 if not args.c:
     sys.exit("CVE id missing")
 
-ref_urls = {"MS": "https://technet.microsoft.com/library/security/",
-            "SECUNIA": "http://secunia.com/advisories/",
-            "SREASON": "http://securityreason.com/security_alert",
-            "CERT": "http://www.cert.org/advisories",
-            "BID": "http://www.securityfocus.com/bid/",
-            "AIXAPART": "",
-            "ALLAIRE": "",
-            "APPLE": "",
-            "ASCEND": "",
-            "ATSTAKE": "",
-            "AUSCERT": "",
-            "BEA": "",
-            "BINDVIEW": "",
-            "SECTRACK": "http://www.securitytracker.com/id/",
-            "MANDRIVA": "http://www.mandriva.com/security/advisories?name="}
+ref_urls = {
+    "MS": "https://technet.microsoft.com/library/security/",
+    "SECUNIA": "http://secunia.com/advisories/",
+    "SREASON": "http://securityreason.com/security_alert",
+    "CERT": "http://www.cert.org/advisories",
+    "BID": "http://www.securityfocus.com/bid/",
+    "AIXAPART": "",
+    "ALLAIRE": "",
+    "APPLE": "",
+    "ASCEND": "",
+    "ATSTAKE": "",
+    "AUSCERT": "",
+    "BEA": "",
+    "BINDVIEW": "",
+    "SECTRACK": "http://www.securitytracker.com/id/",
+    "MANDRIVA": "http://www.mandriva.com/security/advisories?name=",
+}
 
 refs = r.smembers(args.c)
 
@@ -57,10 +62,10 @@
     if args.u:
         (provider, refid) = ref.split(":", 1)
         if provider in ref_urls.keys():
-            print ("{}{}".format(ref_urls[provider], refid))
-        elif provider == 'CONFIRM':
-            print ("{}".format(refid))
+            print("{}{}".format(ref_urls[provider], refid))
+        elif provider == "CONFIRM":
+            print("{}".format(refid))
         else:
-            print (ref)
+            print(ref)
     else:
-        print (ref)
+        print(ref)
diff --git a/bin/db_dump.py b/bin/db_dump.py
@@ -21,25 +21,38 @@
 from lib.DatabaseLayer import getCVEIDs
 
 
-argParser = argparse.ArgumentParser(description='Dump database in JSON format')
-argParser.add_argument('-r', default=False, action='store_true', help='Include ranking value')
-argParser.add_argument('-v', default=False, action='store_true', help='Include via4 map')
-argParser.add_argument('-c', default=False, action='store_true', help='Include CAPEC information')
-argParser.add_argument('-l', default=False, type=int, help='Limit output to n elements (default: unlimited)')
+argParser = argparse.ArgumentParser(description="Dump database in JSON format")
+argParser.add_argument(
+    "-r", default=False, action="store_true", help="Include ranking value"
+)
+argParser.add_argument(
+    "-v", default=False, action="store_true", help="Include via4 map"
+)
+argParser.add_argument(
+    "-c", default=False, action="store_true", help="Include CAPEC information"
+)
+argParser.add_argument(
+    "-l",
+    default=False,
+    type=int,
+    help="Limit output to n elements (default: unlimited)",
+)
 args = argParser.parse_args()
 
 rankinglookup = args.r
 via4lookup = args.v
 capeclookup = args.c
 
-cves = CveHandler(rankinglookup=rankinglookup, via4lookup=via4lookup, capeclookup=capeclookup)
+cves = CveHandler(
+    rankinglookup=rankinglookup, via4lookup=via4lookup, capeclookup=capeclookup
+)
 
 for cveid in getCVEIDs(limit=args.l):
     item = cves.getcve(cveid=cveid)
-    if 'cvss' in item:
-        if type(item['cvss']) == str:
-            item['cvss'] = float(item['cvss'])
-    date_fields = ['cvss-time', 'Modified', 'Published']
+    if "cvss" in item:
+        if type(item["cvss"]) == str:
+            item["cvss"] = float(item["cvss"])
+    date_fields = ["cvss-time", "Modified", "Published"]
     for field in date_fields:
         if field in item:
             item[field] = str(item[field])

diff --git a/bin/dump_last.py b/bin/dump_last.py
@@ -82,7 +82,10 @@
     print("</head><body>")
 for x in cvelist.get(limit=last_items):
     if not (args.f == "html"):
-        item = {"title": str(x["id"]) + " " + x["summary"][:90] + "...", "description": x["summary"]}
+        item = {
+            "title": str(x["id"]) + " " + x["summary"][:90] + "...",
+            "description": x["summary"],
+        }
         if args.r and x.get("ranking"):
             item["description"] = item["description"] + " Ranking:" + str(x["ranking"])
         item["pubDate"] = time.localtime()

diff --git a/bin/search.py b/bin/search.py
@@ -79,7 +79,7 @@ def replace_special_chars(cpe):
     default=False,
     action="store_true",
     help='With this option, a strict vendor product search is executed. The values in "-p" should be formatted as '
-         'vendor:product, e.g. microsoft:windows_7',
+    "vendor:product, e.g. microsoft:windows_7",
 )
 argParser.add_argument(
     "--lax",
@@ -137,7 +137,9 @@ def replace_special_chars(cpe):
 summary_text = args.s
 nlimit = args.i
 
-cves = CveHandler(rankinglookup=rankinglookup, namelookup=namelookup, capeclookup=capeclookup)
+cves = CveHandler(
+    rankinglookup=rankinglookup, namelookup=namelookup, capeclookup=capeclookup
+)
 
 
 def print_job(item):
@@ -161,7 +163,12 @@ def search_product(prod):
     if strict_vendor_product:
         search = prod.split(":")
         search = (search[0], search[1])
-        ret = cvesForCPE(search, lax=relaxSearch, vulnProdSearch=vulnerableProductSearch, strict_vendor_product=True)
+        ret = cvesForCPE(
+            search,
+            lax=relaxSearch,
+            vulnProdSearch=vulnerableProductSearch,
+            strict_vendor_product=True,
+        )
     else:
         ret = cvesForCPE(prod, lax=relaxSearch, vulnProdSearch=vulnerableProductSearch)
     for item in ret["results"]:

diff --git a/lib/ApiRequests.py b/lib/ApiRequests.py
@@ -80,7 +80,9 @@ def process(self, database_connection):
             results = database_connection.query_docs(**self.body)
 
             self.logger.debug(
-                "Retrieving from: {} -- records: {}".format(self.body["retrieve"], len(results))
+                "Retrieving from: {} -- records: {}".format(
+                    self.body["retrieve"], len(results)
+                )
             )
 
             self.request_results = results
@@ -100,13 +102,7 @@ def validate_body(self):
 
         required_keys = ["retrieve", "dict_filter"]
 
-        optional_keys = [
-            "sort",
-            "limit",
-            "skip",
-            "query_filter",
-            "sort_dir",
-        ]
+        optional_keys = ["sort", "limit", "skip", "query_filter", "sort_dir"]
 
         if self.body is None or len(self.body) == 0:
             return (

diff --git a/lib/Authentication.py b/lib/Authentication.py
@@ -33,7 +33,6 @@
 
 
 class AuthenticationMethod(ABC):
-
     @abstractmethod
     def validateUser(self, user, pwd):
         raise NotImplementedError
@@ -76,7 +75,9 @@ def _load_methods(self):
                 self.methods.append((method, authType.lower(), authMethod))
                 self.logger.info("Loaded Auth Method {}".format(x[0]))
             except Exception as e:
-                self.logger.error("Failed to load Auth Method {}:  -> {}".format(x[0], e))
+                self.logger.error(
+                    "Failed to load Auth Method {}:  -> {}".format(x[0], e)
+                )
 
     def isCVESearchUser(self, user):
         return self.dbh.connection.userExists(user)
@@ -100,7 +101,11 @@ def validateUser(self, user, password):
                     if authType == "sufficient" and result is WRONG_CREDS:
                         continue
                 except Exception as e:
-                    self.logger.error("Exception trying to authenticate user: {} -> {}".format(name, e))
+                    self.logger.error(
+                        "Exception trying to authenticate user: {} -> {}".format(
+                            name, e
+                        )
+                    )
         # If we reach here, all methods (if any) failed to authenticate the user
         #  so we check the user against the local database.
         return self.dbh.connection.verifyUser(user, password)

diff --git a/lib/CVEs.py b/lib/CVEs.py
@@ -17,7 +17,15 @@
 runPath = os.path.dirname(os.path.realpath(__file__))
 sys.path.append(os.path.join(runPath, ".."))
 
-from lib.DatabaseLayer import getCAPECFor, getCPE, getAlternativeCPE, getVIA4, getCVE, findRanking, getCVEs
+from lib.DatabaseLayer import (
+    getCAPECFor,
+    getCPE,
+    getAlternativeCPE,
+    getVIA4,
+    getCVE,
+    findRanking,
+    getCVEs,
+)
 
 
 class CveHandler(object):

diff --git a/lib/DatabaseHandler.py b/lib/DatabaseHandler.py
@@ -7,7 +7,9 @@
 class DatabaseHandler(object):
     def __init__(self):
         self.config = Configuration()
-        database_plugin = self.config.readSetting("Database", "PluginName", self.config.default["DatabasePluginName"])
+        database_plugin = self.config.readSetting(
+            "Database", "PluginName", self.config.default["DatabasePluginName"]
+        )
 
         self.dbpluginloader = DatabasePluginLoader()
 
@@ -16,7 +18,9 @@ def __init__(self):
         if isinstance(fetched_plugin, DatabasePluginBase):
             self.connection = fetched_plugin
         else:
-            raise TypeError("The provided plugin is not derived from the DatabasePluginBase class!")
+            raise TypeError(
+                "The provided plugin is not derived from the DatabasePluginBase class!"
+            )
 
     def handle_api_json_query(self, request):
         if not isinstance(request, JSONApiRequest):

diff --git a/lib/DatabaseLayer.py b/lib/DatabaseLayer.py
@@ -161,7 +161,9 @@ def target_version_is_included(target_version, cpe_version):
 
 
 # API Functions
-def cvesForCPE(cpe, lax=False, vulnProdSearch=False, limit=0, strict_vendor_product=False):
+def cvesForCPE(
+    cpe, lax=False, vulnProdSearch=False, limit=0, strict_vendor_product=False
+):
     if not cpe:
         return []
 
@@ -263,19 +265,25 @@ def cvesForCPE(cpe, lax=False, vulnProdSearch=False, limit=0, strict_vendor_prod
         if limit != 0:
             if mongo_version > [4, 4]:
                 cves = (
-                    colCVE.find({"vendors": vendor, "products": {"$regex": cpe_regex_string}})
+                    colCVE.find(
+                        {"vendors": vendor, "products": {"$regex": cpe_regex_string}}
+                    )
                     .limit(limit)
                     .sort("cvss", direction=pymongo.DESCENDING)
                     .allow_disk_use(True)
                 )
             else:
                 cves = (
-                    colCVE.find({"vendors": vendor, "products": {"$regex": cpe_regex_string}})
+                    colCVE.find(
+                        {"vendors": vendor, "products": {"$regex": cpe_regex_string}}
+                    )
                     .limit(limit)
                     .sort("cvss", direction=pymongo.DESCENDING)
                 )
         else:
-            cves = colCVE.find({"vendors": vendor, "products": {"$regex": cpe_regex_string}})
+            cves = colCVE.find(
+                {"vendors": vendor, "products": {"$regex": cpe_regex_string}}
+            )
 
         final_cves = cves
 
@@ -298,19 +306,25 @@ def cvesForCPE(cpe, lax=False, vulnProdSearch=False, limit=0, strict_vendor_prod
         if limit != 0:
             if mongo_version > [4, 4]:
                 cves = (
-                    colCVE.find({"{}".format(cpe_searchField): {"$regex": cpe_regex_string}})
+                    colCVE.find(
+                        {"{}".format(cpe_searchField): {"$regex": cpe_regex_string}}
+                    )
                     .limit(limit)
                     .sort("cvss", direction=pymongo.DESCENDING)
                     .allow_disk_use(True)
                 )
             else:
                 cves = (
-                    colCVE.find({"{}".format(cpe_searchField): {"$regex": cpe_regex_string}})
+                    colCVE.find(
+                        {"{}".format(cpe_searchField): {"$regex": cpe_regex_string}}
+                    )
                     .limit(limit)
                     .sort("cvss", direction=pymongo.DESCENDING)
                 )
         else:
-            cves = colCVE.find({"{}".format(cpe_searchField): {"$regex": cpe_regex_string}})
+            cves = colCVE.find(
+                {"{}".format(cpe_searchField): {"$regex": cpe_regex_string}}
+            )
 
         final_cves = cves
 
@@ -455,8 +469,14 @@ def getSearchResults(search):
 
     search = search.lower()
 
-    vendor_query = {"n": "Vendor", "d": getCVEs(query=[{"vendors": search.replace(" ", "")}])["results"]}
-    product_query = {"n": "Product", "d": getCVEs(query=[{"products": search.replace(" ", "_")}])["results"]}
+    vendor_query = {
+        "n": "Vendor",
+        "d": getCVEs(query=[{"vendors": search.replace(" ", "")}])["results"],
+    }
+    product_query = {
+        "n": "Product",
+        "d": getCVEs(query=[{"products": search.replace(" ", "_")}])["results"],
+    }
 
     for collection in [links, textsearch, vendor_query, product_query]:
         for item in collection["d"]:

diff --git a/lib/DatabasePluginBase.py b/lib/DatabasePluginBase.py
@@ -146,4 +146,4 @@ def isSingleMaster(self, **kwargs):
 
     @abstractmethod
     def getUser(self, **kwargs):
-        raise NotImplementedError
+        raise NotImplementedError