Move Rule into a dataclass

.github/workflows/pythonpackage.yml

@@ -14,10 +14,10 @@ jobs:
     steps:
     - uses: actions/checkout@v2
 
-    - name: Set up Python 3.7
+    - name: Set up Python 3.8
       uses: actions/setup-python@v2
       with:
-        python-version: 3.7
+        python-version: 3.8
 
     - name: Install dependencies
       run: |

Makefile

@@ -14,7 +14,7 @@ all: release
 
 $(VENV):
 	pip install virtualenv
-	virtualenv $(VENV) --python=python3.7
+	virtualenv $(VENV) --python=python3.8
 	$(PIP) install -r requirements.txt
 	$(PIP) install setuptools -U
 

README.md

@@ -1,4 +1,4 @@
-[![Supported Python versions](https://img.shields.io/badge/python-3.7+-yellow.svg)](https://www.python.org/downloads/)
+[![Supported Python versions](https://img.shields.io/badge/python-3.8+-yellow.svg)](https://www.python.org/downloads/)
 [![Unit Tests](https://github.com/elastic/detection-rules/workflows/Unit%20Tests/badge.svg)](https://github.com/elastic/detection-rules/actions)
 [![Chat](https://img.shields.io/badge/chat-%23security--detection--rules-blueviolet)](https://ela.st/slack)
 
@@ -35,7 +35,7 @@ Detection Rules contains more than just static rule files. This repository also
 
 ## Getting started
 
-Although rules can be added by manually creating `.toml` files, we don't recommend it. This repository also consists of a python module that aids rule creation and unit testing. Assuming you have Python 3.7+, run the below command to install the dependencies:
+Although rules can be added by manually creating `.toml` files, we don't recommend it. This repository also consists of a python module that aids rule creation and unit testing. Assuming you have Python 3.8+, run the below command to install the dependencies:
 ```console
 $ pip install -r requirements.txt
 Collecting jsl==0.2.4

detection_rules/cli_utils.py

@@ -11,14 +11,14 @@
 import kql
 from . import ecs
 from .attack import matrix, tactics, build_threat_map_entry
-from .rule import Rule
+from .rule import TOMLRule
 from .schemas import CurrentSchema
 from .utils import clear_caches, get_path
 
 RULES_DIR = get_path("rules")
 
 
-def rule_prompt(path=None, rule_type=None, required_only=True, save=True, verbose=False, **kwargs) -> Rule:
+def rule_prompt(path=None, rule_type=None, required_only=True, save=True, verbose=False, **kwargs) -> TOMLRule:
     """Prompt loop to build a rule."""
     from .misc import schema_prompt
 
@@ -100,7 +100,7 @@ def rule_prompt(path=None, rule_type=None, required_only=True, save=True, verbos
     rule = None
 
     try:
-        rule = Rule(path, {'rule': contents})
+        rule = TOMLRule(path, {'rule': contents})
     except kql.KqlParseError as e:
         if e.error_msg == 'Unknown field':
             warning = ('If using a non-ECS field, you must update "ecs{}.non-ecs-schema.json" under `beats` or '
@@ -113,7 +113,7 @@ def rule_prompt(path=None, rule_type=None, required_only=True, save=True, verbos
         while True:
             try:
                 contents['query'] = click.edit(contents['query'], extension='.eql')
-                rule = Rule(path, {'rule': contents})
+                rule = TOMLRule(path, {'rule': contents})
             except kql.KqlParseError as e:
                 click.secho(e.args[0], fg='red', err=True)
                 click.pause()

detection_rules/devtools.py

@@ -23,7 +23,7 @@
 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 Rule
+from .rule import TOMLRule
 from .rule_loader import get_rule
 from .utils import get_path
 
@@ -96,7 +96,7 @@ def kibana_diff(rule_id, repo, branch, threads):
     repo_hashes = {r.id: r.get_hash() for r in rules.values()}
 
     kibana_rules = {r['rule_id']: r for r in get_kibana_rules(repo=repo, branch=branch, threads=threads).values()}
-    kibana_hashes = {r['rule_id']: Rule.dict_hash(r) for r in kibana_rules.values()}
+    kibana_hashes = {r['rule_id']: TOMLRule.dict_hash(r) for r in kibana_rules.values()}
 
     missing_from_repo = list(set(kibana_hashes).difference(set(repo_hashes)))
     missing_from_kibana = list(set(repo_hashes).difference(set(kibana_hashes)))
@@ -354,7 +354,7 @@ def rule_event_search(ctx, rule_file, rule_id, date_range, count, max_results, v
     if rule_id:
         rule = get_rule(rule_id, verbose=False)
     elif rule_file:
-        rule = Rule(rule_file, load_dump(rule_file))
+        rule = TOMLRule(rule_file, load_dump(rule_file))
     else:
         client_error('Must specify a rule file or rule ID')
 

detection_rules/docs.py

@@ -9,8 +9,10 @@
 
 import xlsxwriter
 
+from . import utils
 from .attack import technique_lookup, matrix, attack_tm, tactics
 from .packaging import Package
+from .rule import ThreatMapping
 
 
 class PackageDocument(xlsxwriter.Workbook):
@@ -47,16 +49,16 @@ def _get_attack_coverage(self):
         coverage = defaultdict(lambda: defaultdict(lambda: defaultdict(int)))
 
         for rule in self.package.rules:
-            threat = rule.contents.get('threat')
+            threat = rule.contents.data.threat
             sub_dir = Path(rule.path).parent.name
 
             if threat:
                 for entry in threat:
-                    tactic = entry['tactic']
-                    techniques = entry.get('technique', [])
+                    tactic = entry.tactic
+                    techniques = entry.technique or []
                     for technique in techniques:
-                        if technique['id'] in matrix[tactic['name']]:
-                            coverage[tactic['name']][technique['id']][sub_dir] += 1
+                        if technique.id in matrix[tactic.name]:
+                            coverage[tactic.name][technique.id][sub_dir] += 1
 
         return coverage
 
@@ -85,10 +87,10 @@ def add_summary(self):
 
         tactic_counts = defaultdict(int)
         for rule in self.package.rules:
-            threat = rule.contents.get('threat')
+            threat = rule.contents.data.threat
             if threat:
                 for entry in threat:
-                    tactic_counts[entry['tactic']['name']] += 1
+                    tactic_counts[entry.tactic.name] += 1
 
         worksheet.write(row, 0, "Total Production Rules")
         worksheet.write(row, 1, len(self.production_rules))
@@ -134,9 +136,9 @@ def add_rule_details(self, rules=None, name='Rule Details'):
         )
 
         for row, rule in enumerate(rules, 1):
-            flat_mitre = rule.get_flat_mitre()
-            rule_contents = {'tactics': flat_mitre['tactic_names'], 'techniques': flat_mitre['technique_ids']}
-            rule_contents.update(rule.contents.copy())
+            flat_mitre = ThreatMapping.flatten(rule.contents.data.threat)
+            rule_contents = {'tactics': flat_mitre.tactic_names, 'techniques': flat_mitre.technique_ids}
+            rule_contents.update(utils.dataclass_to_dict(rule.contents.data))
 
             for column, field in enumerate(metadata_fields):
                 value = rule_contents.get(field)

detection_rules/eswrap.py

@@ -21,7 +21,7 @@
 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 Rule
+from .rule import TOMLRule
 from .rule_loader import get_rule, rta_mappings
 
 
@@ -195,7 +195,7 @@ def search(self, query, language, index: Union[str, list] = '*', start_time=None
 
         return results
 
-    def search_from_rule(self, *rules: Rule, start_time=None, end_time='now', size=None):
+    def search_from_rule(self, *rules: TOMLRule, start_time=None, end_time='now', size=None):
         """Search an elasticsearch instance using a rule."""
         from .misc import nested_get
 

detection_rules/main.py

@@ -19,7 +19,7 @@
 from . import rule_loader
 from .cli_utils import rule_prompt
 from .misc import client_error, nested_set, parse_config
-from .rule import Rule
+from .rule import TOMLRule
 from .rule_formatter import toml_write
 from .schemas import CurrentSchema, available_versions
 from .utils import get_path, clear_caches, load_rule_contents
@@ -119,7 +119,7 @@ def toml_lint(rule_file):
     """Cleanup files with some simple toml formatting."""
     if rule_file:
         contents = pytoml.load(rule_file)
-        rule = Rule(path=rule_file.name, contents=contents)
+        rule = TOMLRule(path=rule_file.name, contents=contents)
 
         # removed unneeded defaults
         for field in rule_loader.find_unneeded_defaults_from_rule(rule):
@@ -179,7 +179,7 @@ def view_rule(ctx, rule_id, rule_file, api_format, verbose=True):
         contents = {k: v for k, v in load_rule_contents(rule_file, single_only=True)[0].items() if v}
 
         try:
-            rule = Rule(rule_file, contents)
+            rule = TOMLRule(rule_file, contents)
         except jsonschema.ValidationError as e:
             client_error(f'Rule: {rule_id or os.path.basename(rule_file)} failed validation', e, ctx=ctx)
     else:
@@ -318,7 +318,7 @@ def search_rules(query, columns, language, count, verbose=True, rules: Dict[str,
             subtechnique_ids.extend([st['id'] for t in techniques for st in t.get('subtechnique', [])])
 
         flat.update(techniques=technique_ids, tactics=tactic_names, subtechniques=subtechnique_ids,
-                    unique_fields=Rule.get_unique_query_fields(rule_doc['rule']))
+                    unique_fields=TOMLRule.get_unique_query_fields(rule_doc['rule']))
         flattened_rules.append(flat)
 
     flattened_rules.sort(key=lambda dct: dct["name"])