Faster Weakness Lookup

Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com>

Author: TG1999

vulnerabilities/models.py

@@ -11,6 +11,7 @@
 import hashlib
 import json
 import logging
+import xml.etree.ElementTree as ET
 from contextlib import suppress
 from functools import cached_property
 from itertools import groupby
@@ -21,6 +22,8 @@
 from cvss.exceptions import CVSS3MalformedError
 from cvss.exceptions import CVSS4MalformedError
 from cwe2.database import Database
+from cwe2.mappings import xml_database_path
+from cwe2.weakness import Weakness as DBWeakness
 from django.contrib.auth import get_user_model
 from django.contrib.auth.models import UserManager
 from django.core import exceptions
@@ -53,9 +56,6 @@
 from vulnerabilities.utils import normalize_purl
 from vulnerabilities.utils import purl_to_dict
 from vulnerablecode import __version__ as VULNERABLECODE_VERSION
-from cwe2.weakness import Weakness as DBWeakness
-from cwe2.mappings import xml_database_path
-import xml.etree.ElementTree as ET
 
 logger = logging.getLogger(__name__)
 
@@ -469,6 +469,7 @@ def get_severity_vectors_and_values(self):
 
         return severity_vectors, severity_values
 
+
 def get_cwes(self):
     """Yield CWE Weakness objects"""
     for cwe_category in self.cwe_files:
@@ -481,10 +482,20 @@ def get_cwes(self):
     for tag_num in [1, 2]:  # Categories , Views
         tag = root[tag_num]
         for child in tag:
-            yield DBWeakness(*[child.attrib["ID"], child.attrib.get("Name"),None,child.attrib.get("Status"),child[0].text])
+            yield DBWeakness(
+                *[
+                    child.attrib["ID"],
+                    child.attrib.get("Name"),
+                    None,
+                    child.attrib.get("Status"),
+                    child[0].text,
+                ]
+            )
+
 
 Database.get_cwes = get_cwes
 
+
 class Weakness(models.Model):
     """
     A Common Weakness Enumeration model

vulnerabilities/views.py

@@ -163,7 +163,9 @@ def get_queryset(self):
             .prefetch_related(
                 Prefetch(
                     "references",
-                    queryset=models.VulnerabilityReference.objects.only("reference_id", "reference_type", "url"),
+                    queryset=models.VulnerabilityReference.objects.only(
+                        "reference_id", "reference_type", "url"
+                    ),
                 ),
                 Prefetch(
                     "aliases",
@@ -175,7 +177,9 @@ def get_queryset(self):
                 ),
                 Prefetch(
                     "severities",
-                    queryset=models.VulnerabilitySeverity.objects.only("scoring_system", "value", "url", "scoring_elements", "published_at"),
+                    queryset=models.VulnerabilitySeverity.objects.only(
+                        "scoring_system", "value", "url", "scoring_elements", "published_at"
+                    ),
                 ),
                 Prefetch(
                     "exploits",
@@ -201,21 +205,27 @@ def get_context_data(self, **kwargs):
         ]
 
         valid_severities = self.object.severities.exclude(scoring_system=EPSS.identifier).filter(
-            scoring_elements__isnull=False, 
-            scoring_system__in=SCORING_SYSTEMS.keys()
+            scoring_elements__isnull=False, scoring_system__in=SCORING_SYSTEMS.keys()
         )
 
         severity_vectors = []
 
         for severity in valid_severities:
             try:
-                vector_values = SCORING_SYSTEMS[severity.scoring_system].get(severity.scoring_elements)
+                vector_values = SCORING_SYSTEMS[severity.scoring_system].get(
+                    severity.scoring_elements
+                )
                 if vector_values:
                     severity_vectors.append({"vector": vector_values, "origin": severity.url})
-            except (CVSS2MalformedError, CVSS3MalformedError, CVSS4MalformedError, NotImplementedError):
+            except (
+                CVSS2MalformedError,
+                CVSS3MalformedError,
+                CVSS4MalformedError,
+                NotImplementedError,
+            ):
                 logging.error(f"CVSSMalformedError for {severity.scoring_elements}")
-        
-        epss_severity = vulnerability.severities.filter(scoring_system='epss').first()
+
+        epss_severity = vulnerability.severities.filter(scoring_system="epss").first()
         epss_data = None
         if epss_severity:
             epss_data = {
@@ -235,7 +245,7 @@ def get_context_data(self, **kwargs):
                 "weaknesses": weaknesses_present_in_db,
                 "status": vulnerability.get_status_label,
                 "history": vulnerability.history,
-                "epss_data": epss_data
+                "epss_data": epss_data,
             }
         )
         return context