Implement the appropriate LoopProgress progress bar.

Refactor the error handling logic in the code.

Signed-off-by: ziadhany <ziadhany2016@gmail.com>

Author: ziadhany

vulnerabilities/pipelines/enhance_with_exploitdb.py

@@ -0,0 +1,149 @@
+import csv
+import io
+import logging
+
+import requests
+from aboutcode.pipeline import LoopProgress
+from dateutil import parser as dateparser
+from django.db import DataError
+
+from vulnerabilities.models import Alias
+from vulnerabilities.models import Exploit
+from vulnerabilities.models import VulnerabilityReference
+from vulnerabilities.models import VulnerabilityRelatedReference
+from vulnerabilities.pipelines import VulnerableCodePipeline
+
+
+class ExploitDBImproverPipeline(VulnerableCodePipeline):
+    """
+    ExploitDB Improver Pipeline: Fetch ExploitDB data, iterate over it to find the vulnerability with
+    the specified alias, and create or update the ref and ref-type accordingly.
+    """
+
+    license_expression = "GPL-2.0"
+
+    @classmethod
+    def steps(cls):
+        return (
+            cls.fetch_exploits,
+            cls.add_exploit,
+        )
+
+    def fetch_exploits(self):
+        exploit_db_url = (
+            "https://gitlab.com/exploit-database/exploitdb/-/raw/main/files_exploits.csv"
+        )
+
+        try:
+            response = requests.get(exploit_db_url)
+            response.raise_for_status()
+        except requests.exceptions.HTTPError as http_err:
+            self.log(
+                f"Failed to fetch the Exploit-DB Exploits: {exploit_db_url} - {http_err}",
+                level=logging.ERROR,
+            )
+            raise
+
+        self.exploit_data = io.StringIO(response.text)
+
+    def add_exploit(self):
+        csvreader = csv.reader(self.exploit_data)
+        header = next(csvreader)
+
+        raw_data = list(csvreader)
+        fetched_exploit_count = len(raw_data)
+
+        vulnerability_exploit_count = 0
+        progress = LoopProgress(total_iterations=fetched_exploit_count, logger=self.log)
+
+        for row in progress.iter(raw_data):
+            vulnerability_exploit_count += add_vulnerability_exploit(row, header, self.log)
+
+        self.log(
+            f"Successfully added {vulnerability_exploit_count:,d} exploit-db vulnerability exploit"
+        )
+
+
+def add_vulnerability_exploit(row, header, logger):
+    vulnerability = None
+    aliases = row[11].split(";")
+
+    for raw_alias in aliases:
+        try:
+            if alias := Alias.objects.get(alias=raw_alias):
+                vulnerability = alias.vulnerability
+                break
+        except Alias.DoesNotExist:
+            continue
+
+    if not vulnerability:
+        logger(f"No vulnerability found for aliases {aliases}")
+        return 0
+
+    add_exploit_references(row[11], row[16], row[1], vulnerability, logger)
+
+    date_added = parse_date(row[header.index("date_added")])
+    source_date_published = parse_date(row[header.index("date_published")])
+    source_date_updated = parse_date(row[header.index("date_updated")])
+
+    try:
+        Exploit.objects.update_or_create(
+            vulnerability=vulnerability,
+            data_source="Exploit-DB",
+            defaults={
+                "date_added": date_added,
+                "description": row[header.index("description")],
+                "known_ransomware_campaign_use": row[header.index("verified")],
+                "source_date_published": source_date_published,
+                "exploit_type": row[header.index("type")],
+                "platform": row[header.index("platform")],
+                "source_date_updated": source_date_updated,
+                "source_url": row[header.index("source_url")],
+            },
+        )
+    except DataError as e:
+        logger(
+            f"Failed to Create the Vulnerability Exploit-DB: {e}",
+            level=logging.ERROR,
+        )
+    return 1
+
+
+def add_exploit_references(ref_id, direct_url, path, vul, logger):
+    url_map = {
+        "file_url": f"https://gitlab.com/exploit-database/exploitdb/-/blob/main/{path}",
+        "direct_url": direct_url,
+    }
+
+    for key, url in url_map.items():
+        if url:
+            try:
+                ref, created = VulnerabilityReference.objects.update_or_create(
+                    url=url,
+                    defaults={
+                        "reference_id": ref_id,
+                        "reference_type": VulnerabilityReference.EXPLOIT,
+                    },
+                )
+
+                if created:
+                    VulnerabilityRelatedReference.objects.get_or_create(
+                        vulnerability=vul,
+                        reference=ref,
+                    )
+
+            except DataError as e:
+                logger(
+                    f"Failed to Create the Vulnerability Reference For Exploit-DB: {e}",
+                    level=logging.ERROR,
+                )
+
+
+def parse_date(date_string):
+    if date_string:
+        try:
+            date_obj = dateparser.parse(date_string).date()
+            return date_obj.strftime("%Y-%m-%d")
+        except (ValueError, TypeError, Exception) as e:
+            logging.error(f"Error while parsing ExploitDB date '{date_string}': {e}")
+    return

vulnerabilities/pipelines/enhance_with_kev.py

@@ -0,0 +1,82 @@
+import logging
+
+import requests
+from aboutcode.pipeline import LoopProgress
+
+from vulnerabilities.models import Alias
+from vulnerabilities.models import Exploit
+from vulnerabilities.pipelines import VulnerableCodePipeline
+
+module_logger = logging.getLogger(__name__)
+
+
+class VulnerabilityKevPipeline(VulnerableCodePipeline):
+    """
+    Known Exploited Vulnerabilities Pipeline: Retrieve KEV data, iterate through it to identify vulnerabilities
+    by their associated aliases, and create or update the corresponding Exploit instances.
+    """
+
+    @classmethod
+    def steps(cls):
+        return (
+            cls.fetch_exploits,
+            cls.add_exploits,
+        )
+
+    def fetch_exploits(self):
+        kev_url = (
+            "https://www.cisa.gov/sites/default/files/feeds/known_exploited_vulnerabilities.json"
+        )
+        try:
+            response = requests.get(kev_url)
+            response.raise_for_status()
+        except requests.exceptions.HTTPError as http_err:
+            self.log(
+                f"Failed to fetch the KEV Exploits: {kev_url} - {http_err}",
+                level=logging.ERROR,
+            )
+            raise
+        self.kev_data = response.json()
+
+    def add_exploits(self):
+        fetched_exploit_count = self.kev_data.get("count")
+        self.log(f"Enhancing the vulnerability with {fetched_exploit_count:,d} exploit records")
+
+        vulnerability_exploit_count = 0
+        progress = LoopProgress(total_iterations=fetched_exploit_count, logger=self.log)
+
+        for record in progress.iter(self.kev_data.get("vulnerabilities", [])):
+            vulnerability_exploit_count += add_vulnerability_exploit(
+                kev_vul=record,
+                logger=self.log,
+            )
+
+        self.log(f"Successfully added {vulnerability_exploit_count:,d} kev exploit")
+
+
+def add_vulnerability_exploit(kev_vul, logger):
+    cve_id = kev_vul.get("cveID")
+
+    vulnerability = None
+    try:
+        if alias := Alias.objects.get(alias=cve_id):
+            vulnerability = alias.vulnerability
+    except Alias.DoesNotExist:
+        logger(f"No vulnerability found for aliases {cve_id}")
+        return 0
+
+    Exploit.objects.update_or_create(
+        vulnerability=vulnerability,
+        data_source="KEV",
+        defaults={
+            "description": kev_vul["shortDescription"],
+            "date_added": kev_vul["dateAdded"],
+            "required_action": kev_vul["requiredAction"],
+            "due_date": kev_vul["dueDate"],
+            "notes": kev_vul["notes"],
+            "known_ransomware_campaign_use": True
+            if kev_vul["knownRansomwareCampaignUse"] == "Known"
+            else False,
+        },
+    )
+    return 1

vulnerabilities/pipelines/enhance_with_metasploit.py

@@ -0,0 +1,103 @@
+import logging
+
+import requests
+import saneyaml
+from aboutcode.pipeline import LoopProgress
+from dateutil import parser as dateparser
+
+from vulnerabilities.models import Alias
+from vulnerabilities.models import Exploit
+from vulnerabilities.pipelines import VulnerableCodePipeline
+
+module_logger = logging.getLogger(__name__)
+
+
+class MetasploitImproverPipeline(VulnerableCodePipeline):
+    """
+    Metasploit Exploits Pipeline: Retrieve Metasploit data, iterate through it to identify vulnerabilities
+    by their associated aliases, and create or update the corresponding Exploit instances.
+    """
+
+    license_expression = "BSD-3-clause"
+
+    @classmethod
+    def steps(cls):
+        return (
+            cls.fetch_exploits,
+            cls.add_vulnerability_exploits,
+        )
+
+    def fetch_exploits(self):
+        url = "https://raw.githubusercontent.com/rapid7/metasploit-framework/master/db/modules_metadata_base.json"
+        self.log(f"Fetching {url}")
+        try:
+            response = requests.get(url)
+            response.raise_for_status()
+        except requests.exceptions.HTTPError as http_err:
+            self.log(
+                f"Failed to fetch the Metasploit Exploits: {url} - {http_err}", level=logging.ERROR
+            )
+            raise
+
+        self.metasploit_data = response.json()
+
+    def add_vulnerability_exploits(self):
+        fetched_exploit_count = len(self.metasploit_data)
+        self.log(f"Enhancing the vulnerability with {fetched_exploit_count:,d} exploit records")
+
+        vulnerability_exploit_count = 0
+        progress = LoopProgress(total_iterations=fetched_exploit_count, logger=self.log)
+        for _, record in progress.iter(self.metasploit_data.items()):
+            vulnerability_exploit_count += add_vulnerability_exploit(
+                record=record,
+                logger=self.log,
+            )
+        self.log(f"Successfully added {vulnerability_exploit_count:,d} vulnerability exploit")
+
+
+def add_vulnerability_exploit(record, logger):
+    vulnerability = None
+    references = record.get("references", [])
+    for ref in references:
+        if ref.startswith("OSVDB") or ref.startswith("URL-"):
+            # ignore OSV-DB and reference exploit for metasploit
+            continue
+
+        try:
+            if alias := Alias.objects.get(alias=ref):
+                vulnerability = alias.vulnerability
+                break
+        except Alias.DoesNotExist:
+            continue
+
+    if not vulnerability:
+        logger(f"No vulnerability found for aliases {references}")
+        return 0
+
+    description = record.get("description", "")
+    notes = record.get("notes", {})
+    platform = record.get("platform")
+
+    source_url = ""
+    if path := record.get("path"):
+        source_url = f"https://github.com/rapid7/metasploit-framework/tree/master{path}"
+    source_date_published = None
+
+    if disclosure_date := record.get("disclosure_date"):
+        try:
+            source_date_published = dateparser.parse(disclosure_date).date()
+        except ValueError:
+            logger(f"Error while parsing date {disclosure_date}", level=logging.ERROR)
+
+    Exploit.objects.update_or_create(
+        vulnerability=vulnerability,
+        data_source="Metasploit",
+        defaults={
+            "description": description,
+            "notes": saneyaml.dump(notes),
+            "source_date_published": source_date_published,
+            "platform": platform,
+            "source_url": source_url,
+        },
+    )
+    return 1