Add support for threshold rules (#65)

detection_rules/main.py

@@ -18,8 +18,9 @@ from eql import load_dump
 from .misc import PYTHON_LICENSE, nested_set
 from . import rule_loader
 from .packaging import PACKAGE_FILE, Package, manage_versions, RELEASE_DIR
-from .rule import RULE_TYPE_OPTIONS, Rule
+from .rule import Rule
 from .rule_formatter import toml_write
+from .schema import RULE_TYPES
 from .utils import get_path, clear_caches
 
 
@@ -35,7 +36,7 @@ def root():
 @click.argument('path', type=click.Path(dir_okay=False))
 @click.option('--config', '-c', type=click.Path(exists=True, dir_okay=False), help='Rule or config file')
 @click.option('--required-only', is_flag=True, help='Only prompt for required fields')
-@click.option('--rule-type', '-t', type=click.Choice(RULE_TYPE_OPTIONS), help='Type of rule to create')
+@click.option('--rule-type', '-t', type=click.Choice(RULE_TYPES), help='Type of rule to create')
 def create_rule(path, config, required_only, rule_type):
     """Create a detection rule."""
     config = load_dump(config) if config else {}

detection_rules/rule.py

@@ -14,12 +14,11 @@ import kql
 from . import ecs, beats
 from .attack import TACTICS, build_threat_map_entry, technique_lookup
 from .rule_formatter import nested_normalize, toml_write
-from .schema import metadata_schema, schema_validate, get_schema
+from .schema import RULE_TYPES, metadata_schema, schema_validate, get_schema
 from .utils import get_path, clear_caches, cached
 
 
 RULES_DIR = get_path("rules")
-RULE_TYPE_OPTIONS = ['machine_learning', 'query', 'saved_id']
 _META_SCHEMA_REQ_DEFAULTS = {}
 
 
@@ -209,8 +208,8 @@ class Rule(object):
 
         kwargs = copy.deepcopy(kwargs)
 
-        while rule_type not in RULE_TYPE_OPTIONS:
-            rule_type = click.prompt('Rule type ({})'.format(', '.join(RULE_TYPE_OPTIONS)))
+        while rule_type not in RULE_TYPES:
+            rule_type = click.prompt('Rule type ({})'.format(', '.join(RULE_TYPES)))
 
         schema = get_schema(rule_type)
         props = schema['properties']
@@ -245,6 +244,11 @@ class Rule(object):
                     contents[name] = threat_map
                 continue
 
+            if name == 'threshold':
+                contents[name] = {n: schema_prompt(f'threshold {n}', required=n in options['required'], **opts)
+                                  for n, opts in options['properties'].items()}
+                continue
+
             if kwargs.get(name):
                 contents[name] = schema_prompt(kwargs.pop(name))
                 continue

detection_rules/schema.py

@@ -33,9 +33,13 @@ NONFORMATTED_FIELDS = 'note',
 # enabled defaults to false instead of true
 # version is a required field that must exist
 
+# rule types
 MACHINE_LEARNING = 'machine_learning'
 SAVED_QUERY = 'saved_query'
 QUERY = 'query'
+THRESHOLD = 'threshold'
+
+RULE_TYPES = [MACHINE_LEARNING, SAVED_QUERY, QUERY, THRESHOLD]
 
 
 class FilterMetadata(jsl.Document):
@@ -100,6 +104,13 @@ class SeverityMapping(jsl.Document):
     severity = jsl.StringField(required=False)
 
 
+class ThresholdMapping(jsl.Document):
+    """Threshold mapping."""
+
+    field = jsl.StringField(required=False)
+    value = jsl.IntField(minimum=1, required=True)
+
+
 class ThreatTactic(jsl.Document):
     """Threat tactics."""
 
@@ -168,6 +179,11 @@ class SiemRuleApiSchema(jsl.Document):
         ml_scope.machine_learning_job_id = jsl.StringField(required=True)
         ml_scope.type = jsl.StringField(enum=[MACHINE_LEARNING], required=True, default=MACHINE_LEARNING)
 
+    with jsl.Scope(SAVED_QUERY) as saved_id_scope:
+        saved_id_scope.index = jsl.ArrayField(jsl.StringField(), required=False)
+        saved_id_scope.saved_id = jsl.StringField(required=True)
+        saved_id_scope.type = jsl.StringField(enum=[SAVED_QUERY], required=True, default=SAVED_QUERY)
+
     with jsl.Scope(QUERY) as query_scope:
         query_scope.index = jsl.ArrayField(jsl.StringField(), required=False)
         # this is not required per the API but we will enforce it here
@@ -175,10 +191,13 @@ class SiemRuleApiSchema(jsl.Document):
         query_scope.query = jsl.StringField(required=True)
         query_scope.type = jsl.StringField(enum=[QUERY], required=True, default=QUERY)
 
-    with jsl.Scope(SAVED_QUERY) as saved_id_scope:
-        saved_id_scope.index = jsl.ArrayField(jsl.StringField(), required=False)
-        saved_id_scope.saved_id = jsl.StringField(required=True)
-        saved_id_scope.type = jsl.StringField(enum=[SAVED_QUERY], required=True, default=SAVED_QUERY)
+    with jsl.Scope(THRESHOLD) as threshold_scope:
+        threshold_scope.index = jsl.ArrayField(jsl.StringField(), required=False)
+        # this is not required per the API but we will enforce it here
+        threshold_scope.language = jsl.StringField(enum=['kuery', 'lucene'], required=True, default='kuery')
+        threshold_scope.query = jsl.StringField(required=True)
+        threshold_scope.type = jsl.StringField(enum=[THRESHOLD], required=True, default=THRESHOLD)
+        threshold_scope.threshold = jsl.DocumentField(ThresholdMapping, required=True)
 
 
 class VersionedApiSchema(SiemRuleApiSchema):