1 files changed, 66 insertions, 43 deletions
diff --git a/src/schema.py b/src/schema.py
index 37da5aa..7e106d1 100644
--- a/src/schema.py
+++ b/src/schema.py
@@ -1,75 +1,98 @@
 import json
+import sys
 
 import jsonschema
 
+# fmt:off
+# NOTE: Commented out properties are left intentionally, so it is easier to see
+#       what properties are optional.
 schema = {
     "$schema": "http://json-schema.org/schema#",
     "type": "object",
     "properties": {
-        "document_version": {
-            "type": "integer"
-        },
-        "ip": {
-            "type": "string"
-        },
-        "port": {
-            "type": "integer"
-        },
-        "whois_description": {
-            "type": "string"
-        },
-        "asn": {
-            "type": "string"
-        },
-        "asn_country_code": {
-            "type": "string"
-        },
-        "ptr": {
-            "type": "string"
-        },
-        "abuse_mail": {
-            "type": "string"
-        },
-        "domain": {
-            "type": "string"
-        },
-        "timestamp_in_utc": {
-            "type": "string"
+        "document_version":  {"type": "integer"},
+        "ip":                {"type": "string"},
+        "port":              {"type": "integer"},
+        "whois_description": {"type": "string"},
+        "asn":               {"type": "string"},
+        "asn_country_code":  {"type": "string"},
+        "ptr":               {"type": "string"},
+        "abuse_mail":        {"type": "string"},
+        "domain":            {"type": "string"},
+        "timestamp_in_utc":  {"type": "string"},
+        "system_name":       {"type": "string"},
+        "description":       {"type": "string"},
+        "custom_data": {
+            "type": "object",
+            "patternProperties": {
+                ".*": {
+                    "type": "object",
+                    "properties": {
+                        "display_name": {"type": "string"},
+                        "data":         {"type": ["string", "boolean", "integer"]},
+                        "description":  {"type": "string"},
+                    },
+                    "required": [
+                        "display_name",
+                        "data",
+                        # "description"
+                    ]
+                },
+            },
         },
-        "user_presentation": {
+        "result": {
             "type": "object",
-            "properties": {
-                "description": {
-                    "type": "string"
+            "patternProperties": {
+                ".*": {
+                    "type": "object",
+                    "properties": {
+                        "display_name": {"type": "string"},
+                        "vulnerable":   {"type": "boolean"},
+                        "reliability":  {"type": "integer"},
+                        "description":  {"type": "string"},
+                    },
+                    "required": [
+                        "display_name",
+                        "vulnerable",
+                        # "reliability", # TODO: reliability is required if vulnerable = true
+                        # "description",
+                    ]
                 },
-                "data": {
-                    "type": "object"
-                }
-            }
-        }
+            },
+        },
     },
     "required": [
         "document_version",
         "ip",
         "port",
+        "whois_description",
+        "asn",
+        "asn_country_code",
+        "ptr",
+        "abuse_mail",
+        "domain",
         "timestamp_in_utc",
-        "user_presentation"
-    ]
+        "system_name",
+        # "description",
+        # "custom_data",
+        "result",
+    ],
 }
+# fmt:on
 
 
 def validate_collector_data(json_blob):
     try:
         jsonschema.validate(json_blob, schema)
     except jsonschema.exceptions.ValidationError as e:
-        print(f'Validation failed with error: {e}')
+        print(f"Validation failed with error: {e}")
         return False
 
     return True
 
 
-if __name__ == '__main__':
-    with open('example_data.json') as fd:
+if __name__ == "__main__":
+    with open(sys.argv[1]) as fd:
         json_data = json.loads(fd.read())
 
     validate_collector_data(json_data)