Required Types via Annotated and Dataclass

Author: eric-forte-elastic

detection_rules/cli_utils.py

@@ -18,6 +18,7 @@
 from . import ecs
 from .attack import build_threat_map_entry, matrix, tactics
 from .config import parse_rules_config
+from .mixins import enforce_required_fields
 from .rule import BYPASS_VERSION_LOCK, TOMLRule, TOMLRuleContents
 from .rule_loader import DEFAULT_PREBUILT_BBR_DIRS, DEFAULT_PREBUILT_RULES_DIRS, RuleCollection, dict_filter
 from .schemas import definitions
@@ -166,6 +167,8 @@ def rule_prompt(  # noqa: PLR0912, PLR0913, PLR0915
     )
 
     target_data_subclass = TOMLRuleContents.get_data_subclass(rule_type_val)
+
+    enforce_required_fields(target_data_subclass)
     schema = target_data_subclass.jsonschema()
     props = schema["properties"]
     required_fields = schema.get("required", []) + additional_required

detection_rules/mixins.py

@@ -8,13 +8,14 @@
 import dataclasses
 import json
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, get_type_hints
 
 import marshmallow
 import marshmallow_dataclass
 import marshmallow_dataclass.union_field
 import marshmallow_jsonschema  # type: ignore[reportMissingTypeStubs]
 import marshmallow_union  # type: ignore[reportMissingTypeStubs]
+import typing_inspect  # type: ignore[reportMissingTypeStubs]
 from marshmallow import Schema, ValidationError, validates_schema
 from marshmallow import fields as marshmallow_fields
 from semver import Version
@@ -38,6 +39,27 @@ def _strip_none_from_dict(obj: Any) -> Any:
     return obj
 
 
+def enforce_required_fields(cls: Any) -> None:
+    """Enforce required fields based on both dataclass and type Annotations."""
+    hints = get_type_hints(cls, include_extras=True)
+    marshmallow_schema = marshmallow_dataclass.class_schema(cls)()
+    for dc_field in dataclasses.fields(cls):
+        mm_field = marshmallow_schema.fields.get(dc_field.name)
+        if mm_field is None or mm_field.required:
+            continue
+
+        if dc_field.default is not dataclasses.MISSING:
+            continue
+        if getattr(dc_field, "default_factory", dataclasses.MISSING) is not dataclasses.MISSING:
+            continue
+
+        hint = hints.get(dc_field.name)
+        if hint is not None and typing_inspect.is_optional_type(hint):  # type: ignore[reportMissingTypeStubs]
+            continue
+
+        mm_field.required = True  # or set to some new list that is required_fields
+
+
 def patch_jsonschema(obj: Any) -> dict[str, Any]:
     """Patch marshmallow-jsonschema output to look more like JSL."""
 

detection_rules/schemas/definitions.py

@@ -243,8 +243,8 @@ def validator_wrapper(value: Any) -> Any:
     list[NonEmptyStr], fields.List(NON_EMPTY_STRING_FIELD, validate=validate.Length(min=1, max=3))
 ]
 PositiveInteger = Annotated[int, fields.Integer(validate=validate.Range(min=1))]
-RiskScore = Annotated[int, fields.Integer(validate=validate.Range(min=1, max=100), required=True)]
-RuleName = Annotated[str, fields.String(validate=elastic_rule_name_regexp(NAME_PATTERN), required=True)]
+RiskScore = Annotated[int, fields.Integer(validate=validate.Range(min=1, max=100))]
+RuleName = Annotated[str, fields.String(validate=elastic_rule_name_regexp(NAME_PATTERN))]
 SemVer = Annotated[str, fields.String(validate=validate.Regexp(VERSION_PATTERN))]
 SemVerMinorOnly = Annotated[str, fields.String(validate=validate.Regexp(MINOR_SEMVER))]
 Sha256 = Annotated[str, fields.String(validate=validate.Regexp(SHA256_PATTERN))]
@@ -254,7 +254,7 @@ def validator_wrapper(value: Any) -> Any:
 ThresholdValue = Annotated[int, fields.Integer(validate=validate.Range(min=1))]
 TimelineTemplateId = Annotated[str, fields.String(validate=elastic_timeline_template_id_validator())]
 TimelineTemplateTitle = Annotated[str, fields.String(validate=elastic_timeline_template_title_validator())]
-UUIDString = Annotated[str, fields.String(validate=validate.Regexp(UUID_PATTERN), required=True)]
+UUIDString = Annotated[str, fields.String(validate=validate.Regexp(UUID_PATTERN))]
 
 # experimental machine learning features and releases
 MachineLearningType = Literal[MACHINE_LEARNING_PACKAGES]