Implement json schema validation

Author: felddy

src/cyhy_kevsync/sync.py

@@ -1,52 +1,61 @@
+# Standard Python Libraries
 import json
 import logging
-import time
-import random
+from typing import Optional
 import urllib.request
 
-from rich.progress import track
-
+# Third-Party Libraries
 from cyhy_db.models import KEVDoc
+from jsonschema import SchemaError, ValidationError, validate
+from rich.progress import track
 
 # TODO rename this file to something better
 
+ALLOWED_URL_SCHEMES = ["http", "https"]
+
 logger = logging.getLogger(__name__)
 
-# def sync(url: str = DEFAULT_KEV_URL) -> None:
-#     """Synchronize the KEV data from the given URL."""
 
-#     for _ in track(
-#         range(100),
-#         description="KEV Syncing",
-#     ):
-#         time.sleep(random.uniform(0.01, 1))
+async def fetch_kev_data(
+    kev_json_url: str, kev_schema_url: Optional[str] = None
+) -> dict:
+    """Fetch the KEV data from the given URL."""
 
+    # Create a Request object so we can test the safety of the URL
+    key_json_request = urllib.request.Request(kev_json_url)
+    if key_json_request.type not in ALLOWED_URL_SCHEMES:
+        raise ValueError("Invalid URL scheme in json URL: %s" % key_json_request.type)
 
-async def fetch_kev_data(url: str) -> dict:
-    """Fetch the KEV data from the given URL."""
+    # Below we disable the bandit blacklist for the urllib.request.urlopen() function
+    # since we are checking the URL scheme before using.
 
-    # We disable the bandit blacklist for the urllib.request.urlopen() function
-    # because the URL is either the default (safe) URL or one provided in the
-    # Lambda configuration so we can assume it is safe.
-    with urllib.request.urlopen(url) as response:  # nosec B310
+    with urllib.request.urlopen(kev_json_url) as response:  # nosec B310
         if response.status != 200:
             raise Exception("Failed to retrieve KEV JSON.")
 
         kev_json = json.loads(response.read().decode("utf-8"))
 
-    # TODO: Check the data against the schema
-    # https://www.cisa.gov/sites/default/files/feeds/known_exploited_vulnerabilities_schema.json
-
-    # Sanity check the JSON data
-    if "vulnerabilities" not in kev_json:
-        raise ValueError(
-            "JSON does not look like valid KEV data. Missing vulnerabilities."
-        )
+    # If a schema URL was provided, we will validate the JSON data against it
+    if kev_schema_url:
+        # Create a Request object so we can test the safety of the URL
+        key_schema_request = urllib.request.Request(kev_schema_url)
+        if key_schema_request.type not in ALLOWED_URL_SCHEMES:
+            raise ValueError(
+                "Invalid URL scheme in schema URL: %s" % key_json_request.type
+            )
+        with urllib.request.urlopen(kev_schema_url) as response:  # nosec B310
+            if response.status != 200:
+                raise Exception("Failed to retrieve KEV JSON schema.")
+            kev_schema = json.loads(response.read().decode("utf-8"))
+            try:
+                validate(instance=kev_json, schema=kev_schema)
+                logger.info("KEV JSON is valid against the schema.")
+            except ValidationError as e:
+                logger.error("JSON validation error: %s", e.message)
+            except SchemaError as e:
+                logger.error("Schema error: %s", e.message)
 
     reported_vuln_count = kev_json.get("count")
-    if reported_vuln_count is None:
-        raise ValueError("JSON does not look like valid KEV data.  Missing count.")
-
     actual_vuln_count = len(kev_json["vulnerabilities"])
     if reported_vuln_count != actual_vuln_count:
         logger.warning(
@@ -67,8 +76,8 @@ async def create_kev_doc(kev_json: dict) -> str:
     """Add the provided KEV to the database and return its id."""
     cve_id = kev_json.get("cveID")
     if not cve_id:
-        raise ValueError("JSON does not look like valid KEV data.")
-    known_ransomware = kev_json.get("knownRansomwareCampaignUse").lower() == "known"
+        raise ValueError("cveID not found in KEV JSON.")
+    known_ransomware = kev_json["knownRansomwareCampaignUse"].lower() == "known"
     kev_doc = KEVDoc(id=cve_id, known_ransomware=known_ransomware)
     await kev_doc.save()
     logger.debug("Created KEV document with id: %s", cve_id)
@@ -83,7 +92,10 @@ async def remove_outdated_kevs() -> None:
 
 async def process_kev_json(kev_json: dict) -> None:
     """Process the KEV JSON data."""
-    for kev in kev_json["vulnerabilities"]:
+    for kev in track(
+        kev_json["vulnerabilities"],
+        description="Creating KEV docs",
+    ):
         try:
             await create_kev_doc(kev)
         except Exception as e: