Updated schema with tags and better logic for the 'result' tag, lets force descriptionvictor_updates

2 files changed, 39 insertions, 26 deletions

diff --git a/src/soc_collector/schema.py b/src/soc_collector/schema.py
index 2c2dfb9..7688631 100644
--- a/src/soc_collector/schema.py
+++ b/src/soc_collector/schema.py
@@ -1,6 +1,10 @@
 """Our schema module"""
 from typing import Any, Dict, Optional, Union
-import jsonschema
+from jsonschema.exceptions import ValidationError
+
+# docker-compose require jsonschema < 4 so use Draft7 for now
+from jsonschema.validators import Draft7Validator
+
 from bson import ObjectId
 from bson.errors import InvalidId
 
@@ -8,12 +12,19 @@ from bson.errors import InvalidId
 # NOTE: Commented out properties are left intentionally, so it is easier to see
 #       what properties are optional.
 schema = {
-    "$schema": "http://json-schema.org/schema#",
+    "$schema": "https://json-schema.org/draft-07/schema#",
     "type": "object",
     "properties": {
-        "document_version":  {"type": "integer"},
+        "document_version":  {"type": "integer", "minimum": 2},
+        "tags": {
+            "type": "array",
+            "uniqueItems": True,
+            "items": {"type": "string", "enum": ["dev", "display", "hide", "notify"]},
+            "minItems": 1,
+            "maxItems": 3  # hide is incompatible with notify and/or display
+        },
         "ip":                {"type": "string"},
-        "port":              {"type": "integer"},
+        "port":              {"type": "integer", "minimum": 1, "maximum": 65535},
         "whois_description": {"type": "string"},
         "asn":               {"type": "string"},
         "asn_country_code":  {"type": "string"},
@@ -50,34 +61,29 @@ schema = {
                         "display_name":         {"type": "string"},
                         "vulnerable":           {"type": "boolean"},
                         "investigation_needed": {"type": "boolean"},
-                        "reliability":          {"type": "integer"},
+                        "reliability":          {"type": "integer", "minimum": 1, "maximum": 5},
+                        "severity":             {"type": "integer", "minimum": 1, "maximum": 5},
                         "description":          {"type": "string"},
                     },
                     "oneOf": [
                         {
-                            "required": [
-                                "display_name",
-                                "vulnerable",
-                                # "reliability", # TODO: reliability is required if vulnerable = true
-                                # "description",
-                            ]
+                            "properties": {"investigation_needed": {"const": True}},
+                            "required": ["display_name", "investigation_needed", "description"]
                         },
                         {
-                            "required":
-                            [
-                                "display_name",
-                                "investigation_needed",
-                                # "reliability",        # TODO: reliability is required if investigation_needed = true
-                                # "description",
-                            ]
-                        },
+                            "properties": {"vulnerable": {"type": "boolean"}},
+                            "if":   {"properties": {"vulnerable": {"const": True}}},
+                            "then": {"required": ["display_name", "vulnerable", "reliability", "severity", "description"]},
+                            "else": {"required": ["display_name", "vulnerable"]}
+                        }
                     ]
-                },
-            },
-        },
+                }
+            }
+        }
     },
     "required": [
         "document_version",
+        "tags",
         "ip",
         "port",
         "whois_description",
@@ -88,9 +94,9 @@ schema = {
         "domain",
         "timestamp",
         "display_name",
+        "result"
         # "description",
         # "custom_data",
-        "result",
     ],
 }
 
@@ -103,8 +109,9 @@ def valid_schema(json_data: Dict[str, Any]) -> bool:
     """
 
     try:
-        jsonschema.validate(json_data, schema, format_checker=jsonschema.FormatChecker())
-    except jsonschema.exceptions.ValidationError as exc:
+        # docker-compose require jsonschema < 4 so use Draft7 for now
+        Draft7Validator(schema).validate(json_data)
+    except ValidationError as exc:
         print(f"Validation failed with error: {exc.message}")
         return False
     return True
diff --git a/src/soc_collector/soc_collector_cli.py b/src/soc_collector/soc_collector_cli.py
index 4929655..85afce1 100644
--- a/src/soc_collector/soc_collector_cli.py
+++ b/src/soc_collector/soc_collector_cli.py
@@ -8,7 +8,7 @@ from sys import exit as app_exit
 import json
 import requests
 
-from .schema import object_id_from_data
+from .schema import object_id_from_data, valid_schema
 
 ROOT_CA_FILE = __file__.replace("soc_collector_cli.py", "data/collector_root_ca.crt")
 
@@ -136,6 +136,9 @@ def replace_action(data: str, api_key: str, base_url: str = "https://collector-d
         print("ERROR: Valid '_id' key not in data")
         app_exit(1)
 
+    if not valid_schema(json_data):
+        app_exit(1)
+
     req = requests.put(
         f"{base_url}/sc/v0", json=json_data, headers={"API-KEY": api_key}, timeout=5, verify=ROOT_CA_FILE
     )
@@ -166,6 +169,9 @@ def insert_action(data: str, api_key: str, base_url: str = "https://collector-de
         print("ERROR: '_id' key in data")
         app_exit(1)
 
+    if not valid_schema(json_data):
+        app_exit(1)
+
     req = requests.post(
         f"{base_url}/sc/v0", json=json_data, headers={"API-KEY": api_key}, timeout=5, verify=ROOT_CA_FILE
     )