Add a RuleCollection object instead of a "loader" module

detection_rules/attack.py

@@ -4,13 +4,13 @@
 # 2.0.
 
 """Mitre attack info."""
+import json
 import os
 import re
 import time
+from collections import OrderedDict
 
-import json
 import requests
-from collections import OrderedDict
 
 from .semver import Version
 from .utils import get_etc_path, get_etc_glob_path, read_gzip, gzip_compress

detection_rules/beats.py

@@ -4,16 +4,16 @@
 # 2.0.
 
 """ECS Schemas management."""
+import json
 import os
 import re
 from typing import List
 
-import kql
 import eql
-import json
 import requests
 import yaml
 
+import kql
 from .semver import Version
 from .utils import DateTimeEncoder, unzip, get_etc_path, gzip_compress, read_gzip, cached
 

detection_rules/devtools.py

@@ -17,15 +17,15 @@
 import click
 from elasticsearch import Elasticsearch
 from eql import load_dump
-from kibana.connector import Kibana
 
+from kibana.connector import Kibana
 from . import rule_loader
 from .eswrap import CollectEvents, add_range_to_dsl
 from .main import root
 from .misc import PYTHON_LICENSE, add_client, GithubClient, Manifest, client_error, getdefault
 from .packaging import PACKAGE_FILE, Package, manage_versions, RELEASE_DIR
-from .rule import TOMLRule, TOMLRuleContents, BaseQueryRuleData
-from .rule_loader import get_rule
+from .rule import TOMLRule, BaseQueryRuleData
+from .rule_loader import production_filter, RuleCollection
 from .utils import get_path, dict_hash
 
 RULES_DIR = get_path('rules')
@@ -68,7 +68,7 @@ def update_lock_versions(rule_ids):
     if not click.confirm('Are you sure you want to update hashes without a version bump?'):
         return
 
-    rules = [r for r in rule_loader.load_rules(verbose=False).values() if r.id in rule_ids]
+    rules = RuleCollection.default().filter(lambda r: r.id in rule_ids)
     changed, new = manage_versions(rules, exclude_version_update=True, add_new=False, save_changes=True)
 
     if not changed:
@@ -86,10 +86,12 @@ def kibana_diff(rule_id, repo, branch, threads):
     """Diff rules against their version represented in kibana if exists."""
     from .misc import get_kibana_rules
 
+    rules = RuleCollection.default()
+
     if rule_id:
-        rules = {r.id: r for r in rule_loader.load_rules(verbose=False).values() if r.id in rule_id}
+        rules = rules.filter(lambda r: r.id in rule_id)
     else:
-        rules = {r.id: r for r in rule_loader.get_production_rules()}
+        rules = rules.filter(production_filter)
 
     # add versions to the rules
     manage_versions(list(rules.values()), verbose=False)
@@ -388,9 +390,11 @@ def rule_event_search(ctx, rule_file, rule_id, date_range, count, max_results, v
     rule: TOMLRule
 
     if rule_id:
-        rule = get_rule(rule_id, verbose=False)
+        rule = RuleCollection().load_by_id(rule_id)
+        if rule is None:
+            client_error(f"Unable to find rule with id {rule_id}")
     elif rule_file:
-        rule = TOMLRule(path=rule_file, contents=TOMLRuleContents.from_dict(load_dump(rule_file)))
+        rule = RuleCollection().load_file(rule_file)
     else:
         client_error('Must specify a rule file or rule ID')
 
@@ -431,18 +435,17 @@ def rule_survey(ctx: click.Context, query, date_range, dump_file, hide_zero_coun
     """Survey rule counts."""
     from eql.table import Table
     from kibana.resources import Signal
-    from . import rule_loader
     from .main import search_rules
 
     survey_results = []
     start_time, end_time = date_range
 
     if query:
-        rule_paths = [r['file'] for r in ctx.invoke(search_rules, query=query, verbose=False)]
-        rules = rule_loader.load_rules(rule_loader.load_rule_files(paths=rule_paths, verbose=False), verbose=False)
-        rules = rules.values()
+        rules = RuleCollection()
+        paths = [Path(r['file']) for r in ctx.invoke(search_rules, query=query, verbose=False)]
+        rules.load_files(paths)
     else:
-        rules = rule_loader.load_rules(verbose=False).values()
+        rules = RuleCollection.default().filter(production_filter)
 
     click.echo(f'Running survey against {len(rules)} rules')
     click.echo(f'Saving detailed dump to: {dump_file}')

detection_rules/ecs.py

@@ -6,13 +6,13 @@
 """ECS Schemas management."""
 import copy
 import glob
+import json
 import os
 import shutil
-import json
 
-import requests
 import eql
 import eql.types
+import requests
 import yaml
 
 from .semver import Version

detection_rules/eswrap.py

@@ -7,8 +7,8 @@
 import json
 import os
 import time
-from contextlib import contextmanager
 from collections import defaultdict
+from contextlib import contextmanager
 from pathlib import Path
 from typing import Union
 
@@ -20,10 +20,9 @@
 import kql
 from .main import root
 from .misc import add_params, client_error, elasticsearch_options
-from .utils import format_command_options, normalize_timing_and_sort, unix_time_to_formatted, get_path
 from .rule import TOMLRule
-from .rule_loader import get_rule, rta_mappings
-
+from .rule_loader import rta_mappings, RuleCollection
+from .utils import format_command_options, normalize_timing_and_sort, unix_time_to_formatted, get_path
 
 COLLECTION_DIR = get_path('collections')
 MATCH_ALL = {'bool': {'filter': [{'match_all': {}}]}}
@@ -88,7 +87,8 @@ def evaluate_against_rule_and_update_mapping(self, rule_id, rta_name, verbose=Tr
         """Evaluate a rule against collected events and update mapping."""
         from .utils import combine_sources, evaluate
 
-        rule = get_rule(rule_id, verbose=False)
+        rule = RuleCollection.load_by_id(rule_id)
+        assert rule is not None, f"Unable to find rule with ID {rule_id}"
         merged_events = combine_sources(*self.events.values())
         filtered = evaluate(rule, merged_events)
 

detection_rules/kbwrap.py

@@ -4,13 +4,17 @@
 # 2.0.
 
 """Kibana cli commands."""
+import uuid
+from pathlib import Path
+
 import click
+
 import kql
 from kibana import Kibana, Signal, RuleResource
-
 from .main import root
 from .misc import add_params, client_error, kibana_options
-from .rule_loader import load_rule_files, load_rules
+from .rule_loader import RuleCollection
+from .schemas import downgrade
 from .utils import format_command_options
 
 
@@ -54,25 +58,26 @@ def kibana_group(ctx: click.Context, **kibana_kwargs):
 @click.pass_context
 def upload_rule(ctx, toml_files, replace_id):
     """Upload a list of rule .toml files to Kibana."""
-    from .packaging import manage_versions
 
     kibana = ctx.obj['kibana']
-    file_lookup = load_rule_files(paths=toml_files)
-    rules = list(load_rules(file_lookup=file_lookup).values())
-
-    # assign the versions from etc/versions.lock.json
-    # rules that have changed in hash get incremented, others stay as-is.
-    # rules that aren't in the lookup default to version 1
-    manage_versions(rules, verbose=False)
+    rules = RuleCollection()
+    rules.load_files(Path(p) for p in toml_files)
 
     api_payloads = []
 
     for rule in rules:
         try:
-            payload = rule.get_payload(include_version=True, replace_id=replace_id, embed_metadata=True,
-                                       target_version=kibana.version)
+            payload = rule.contents.to_api_format()
+            payload.setdefault("meta", {}).update(rule.contents.metadata.to_dict())
+
+            if replace_id:
+                payload["rule_id"] = str(uuid.uuid4())
+
+            payload = downgrade(payload, target_version=kibana.version)
+
         except ValueError as e:
             client_error(f'{e} in version:{kibana.version}, for rule: {rule.name}', e, ctx=ctx)
+
         rule = RuleResource(payload)
         api_payloads.append(rule)