Add redundant code to utils

Signed-off-by: ambuj <[email protected]>

Author: ambuj-1211

vulnerabilities/importers/apache_httpd.py

@@ -13,7 +13,6 @@
 
 import requests
 from bs4 import BeautifulSoup
-from cwe2.database import Database
 from packageurl import PackageURL
 from univers.version_constraint import VersionConstraint
 from univers.version_range import ApacheVersionRange
@@ -25,7 +24,8 @@
 from vulnerabilities.importer import Reference
 from vulnerabilities.importer import VulnerabilitySeverity
 from vulnerabilities.severity_systems import APACHE_HTTPD
-from vulnerabilities.utils import get_cwe_id
+from vulnerabilities.utils import create_weaknesses_list
+from vulnerabilities.utils import cwe_regex
 from vulnerabilities.utils import get_item
 
 logger = logging.getLogger(__name__)
@@ -234,33 +234,21 @@ def get_weaknesses(cve_data):
         >>> get_weaknesses(mock_cve_data2)
         [190, 200]
     """
-
     alias = get_item(cve_data, "CVE_data_meta", "ID")
-    cwe_id = []
-    db = Database()
+    cwe_strings = []
     if alias:
         problemtype_data = get_item(cve_data, "problemtype", "problemtype_data") or []
         for problem in problemtype_data:
-            for desc in problem["description"]:
+            for desc in problem.get("description", []):
                 value = desc.get("value", "")
-                cwe_pattern = r"CWE-\d+"
-                cwe_id_string_list = re.findall(cwe_pattern, value)
-                for cwe_id_string in cwe_id_string_list:
-                    cwe_id.append(get_cwe_id(cwe_id_string))
-
+                cwe_id_string_list = re.findall(cwe_regex, value)
+                cwe_strings.extend(cwe_id_string_list)
     else:
         problemTypes = cve_data.get("containers", {}).get("cna", {}).get("problemTypes", [])
         descriptions = problemTypes[0].get("descriptions", []) if len(problemTypes) > 0 else []
         for description in descriptions:
             cwe_id_string = description.get("cweId", "")
-            cwe_id.append(get_cwe_id(cwe_id_string))
-
-    weaknesses = []
-    for cwe in cwe_id:
-        try:
-            db.get(cwe)
-            weaknesses.append(cwe)
-        except Exception:
-            logger.error("Invalid CWE id")
+            cwe_strings.append(cwe_id_string)
 
+    weaknesses = create_weaknesses_list(cwe_strings)
     return weaknesses

vulnerabilities/importers/debian.py

@@ -24,8 +24,8 @@
 from vulnerabilities.importer import AffectedPackage
 from vulnerabilities.importer import Importer
 from vulnerabilities.importer import Reference
+from vulnerabilities.utils import create_weaknesses_list
 from vulnerabilities.utils import dedupe
-from vulnerabilities.utils import get_cwe_id
 from vulnerabilities.utils import get_item
 
 logger = logging.getLogger(__name__)
@@ -178,14 +178,5 @@ def get_cwe_from_debian_advisory(record):
     description = record.get("description") or ""
     pattern = r"CWE-\d+"
     cwe_strings = re.findall(pattern, description)
-    weaknesses = []
-    db = Database()
-    for cwe_string in cwe_strings:
-        if cwe_string:
-            cwe_id = get_cwe_id(cwe_string)
-            try:
-                db.get(cwe_id)
-                weaknesses.append(cwe_id)
-            except Exception:
-                logger.error("Invalid CWE id")
+    weaknesses = create_weaknesses_list(cwe_strings)
     return weaknesses

vulnerabilities/importers/fireeye.py

@@ -12,14 +12,13 @@
 from typing import Iterable
 from typing import List
 
-from cwe2.database import Database
-
 from vulnerabilities.importer import AdvisoryData
 from vulnerabilities.importer import Importer
 from vulnerabilities.importer import Reference
 from vulnerabilities.utils import build_description
+from vulnerabilities.utils import create_weaknesses_list
+from vulnerabilities.utils import cwe_regex
 from vulnerabilities.utils import dedupe
-from vulnerabilities.utils import get_cwe_id
 
 logger = logging.getLogger(__name__)
 
@@ -160,19 +159,8 @@ def get_weaknesses(cwe_data):
     """
     cwe_list = []
     for line in cwe_data:
-        cwe_ids = re.findall(r"CWE-\d+", line)
+        cwe_ids = re.findall(cwe_regex, line)
         cwe_list.extend(cwe_ids)
 
-    weaknesses = []
-    db = Database()
-
-    for cwe_string in cwe_list:
-
-        if cwe_string:
-            cwe_id = get_cwe_id(cwe_string)
-            try:
-                db.get(cwe_id)
-                weaknesses.append(cwe_id)
-            except Exception:
-                logger.error("Invalid CWE id")
+    weaknesses = create_weaknesses_list(cwe_list)
     return weaknesses

vulnerabilities/tests/test_fireeye.py

@@ -226,7 +226,6 @@ def test_get_weaknesses(self):
                 "CWE-362: Concurrent Execution using Shared Resource with Improper Synchronization ('Race Condition')",
             ]
         ) == [379, 362]
-
         assert (
             get_weaknesses(
                 [

vulnerabilities/utils.py

@@ -29,6 +29,8 @@
 import saneyaml
 import toml
 import urllib3
+from cwe2.database import Database
+from cwe2.database import InvalidCWEError
 from packageurl import PackageURL
 from packageurl.contrib.django.utils import without_empty_values
 from univers.version_range import RANGE_CLASS_BY_SCHEMES
@@ -42,6 +44,7 @@
 cve_regex = re.compile(r"CVE-[0-9]{4}-[0-9]{4,19}", re.IGNORECASE)
 is_cve = cve_regex.match
 find_all_cve = cve_regex.findall
+cwe_regex = r"CWE-\d+"
 
 
 @dataclasses.dataclass(order=True, frozen=True)
@@ -399,6 +402,29 @@ def get_cwe_id(cwe_string: str) -> int:
     return int(cwe_id)
 
 
+def create_weaknesses_list(cwe_strings: str):
+    """
+    Convert the CWE string to CWE ids and store them to weaknesses list.
+    >>> create_weaknesses_list(["CWE-125","CWE-379"])
+    [125, 379]
+    """
+    weaknesses = []
+    db = Database()
+    for cwe_string in cwe_strings:
+        if not cwe_string:
+            continue
+        cwe_id = get_cwe_id(cwe_string)
+        if not cwe_id:
+            logger.error("Invalid CWE id: No CWE ID found")
+            continue
+        try:
+            db.get(cwe_id)
+            weaknesses.append(cwe_id)
+        except InvalidCWEError as e:
+            logger.error(f"Error: {e}")
+    return weaknesses
+
+
 def clean_nginx_git_tag(tag):
     """
     Return a cleaned ``version`` string from an nginx git tag.