Merge pull request #34 from ocsf/export-categories

Wiring category events into compilation CLI options and making event …

Author: query-jeremy

src/ocsf/compile/__main__.py

@@ -32,6 +32,11 @@
   --set-observable      Set the observable field on attributes to the corresponding Observable Type ID where applicable
   --no-set-observable   Do not set the observable field on attributes to the corresponding Observable Type ID
                         where applicable
+  --category-classes    Include classes in the category section of the schema output. This duplicates classes in each
+                        category without attributes in the same way as the OCSF Server's category export.
+  --no-category-classes
+                        Do not include classes in the category section of the schema output. This is the default
+                        behavior.
 ```
 
 Examples:
@@ -125,6 +130,22 @@ def main():
         action="store_false",
         help="Do not set the observable field on attributes to the corresponding Observable Type ID where applicable",
     )
+    parser.add_argument(
+        "--category-classes",
+        action="store_true",
+        default=True,
+        help=(
+            "Include classes in the category section of the schema output. This "
+            "duplicates classes in each category without attributes in the same "
+            "way as the OCSF Server's category export."
+        ),
+    )
+    parser.add_argument(
+        "--no-category-classes",
+        dest="category_classes",
+        action="store_false",
+        help=("Do not include classes in the category section of the schema output. " "This is the default behavior."),
+    )
 
     args = parser.parse_args()
 
@@ -142,6 +163,7 @@ def main():
     options.prefix_extensions = args.prefix_extensions
     options.set_object_types = args.set_object_types
     options.set_observable = args.set_observable
+    options.map_events_to_categories = args.category_classes
 
     repo = read_repo(args.path, preserve_raw_data=False)
 

src/ocsf/compile/compiler.py

@@ -111,7 +111,7 @@ def __init__(self, repo: Repository, options: CompilationOptions | None = None):
                 # attribute of records easier. Only performed if
                 # options.set_observable is True.
                 MarkObservablesPlanner(self._proto, _options),
-                # Map events to categories in the categories.json file.
+                # Map events to the `classes` dictionary of each category.
                 MapEventToCategoryPlanner(self._proto, _options),
                 # Copy records that are ONLY defined in extensions to the core
                 # schema so that they are included by ProtoSchema.schema().

src/ocsf/compile/planners/set_category.py

@@ -21,11 +21,16 @@ def apply(self, schema: ProtoSchema) -> MergeResult:
             return []
 
         path = PurePath(target.path).parts
+        if path[-1] == SpecialFiles.BASE_EVENT.value:
+            return []
+
         if RepoPaths.EVENTS.value not in path:
             raise ValueError(f"Cannot assign category to non-event: {target.path}")
             return []
 
         category = path[path.index(RepoPaths.EVENTS.value) + 1]
+        if category == "base_event.json":
+            return []
 
         categories = schema[SpecialFiles.CATEGORIES].data
         assert isinstance(categories, CategoriesDefn)

src/ocsf/repository/helpers.py

@@ -41,13 +41,21 @@ class SpecialFiles(StrEnum):
     VERSION = "version.json"
     EXTENSION = "extension.json"
     OBSERVABLE = "objects/observable.json"
+    BASE_EVENT = "base_event.json"
 
     @staticmethod
     def contains(path: str) -> bool:
         return path in [e.value for e in SpecialFiles]
 
 
-SPECIAL_FILES = ("dictionary.json", "categories.json", "version.json", "extension.json", "objects/observable.json")
+SPECIAL_FILES = (
+    "dictionary.json",
+    "categories.json",
+    "version.json",
+    "extension.json",
+    "objects/observable.json",
+    "base_event.json",
+)
 """Tuple containing strings of the values in the SpecialFiles enum."""