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/__init__.py

@@ -4,17 +4,23 @@
 # 2.0.
 
 """Detection rules."""
-from . import devtools
-from . import docs
-from . import eswrap
-from . import kbwrap
-from . import main
-from . import mappings
-from . import misc
-from . import rule_formatter
-from . import rule_loader
-from . import schemas
-from . import utils
+import sys
+
+assert (3, 8) <= sys.version_info < (4, 0), "Only Python 3.8+ supported"
+
+from . import (  # noqa: E402
+    devtools,
+    docs,
+    eswrap,
+    kbwrap,
+    main,
+    mappings,
+    misc,
+    rule_formatter,
+    rule_loader,
+    schemas,
+    utils
+)
 
 __all__ = (
     'devtools',

detection_rules/__main__.py

@@ -6,8 +6,13 @@
 # coding=utf-8
 """Shell for detection-rules."""
 import os
+import sys
+
 import click
-from .main import root
+
+assert (3, 8) <= sys.version_info < (4, 0), "Only Python 3.8+ supported"
+
+from .main import root  # noqa: E402
 
 CURR_DIR = os.path.dirname(os.path.abspath(__file__))
 CLI_DIR = os.path.dirname(CURR_DIR)

detection_rules/cli_utils.py

@@ -4,24 +4,27 @@
 # 2.0.
 
 import copy
+import datetime
 import os
+from pathlib import Path
 
 import click
 
 import kql
 from . import ecs
 from .attack import matrix, tactics, build_threat_map_entry
-from .rule import Rule
+from .rule import TOMLRule, TOMLRuleContents
 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
 
+    creation_date = datetime.date.today().strftime("%Y/%m/%d")
     if verbose and path:
         click.echo(f'[+] Building rule for {path}')
 
@@ -32,8 +35,7 @@ def rule_prompt(path=None, rule_type=None, required_only=True, save=True, verbos
         kwargs.update(kwargs.pop('rule'))
 
     rule_type = rule_type or kwargs.get('type') or \
-        click.prompt('Rule type ({})'.format(', '.join(CurrentSchema.RULE_TYPES)),
-                     type=click.Choice(CurrentSchema.RULE_TYPES))
+        click.prompt('Rule type', type=click.Choice(CurrentSchema.RULE_TYPES))
 
     schema = CurrentSchema.get_schema(role=rule_type)
     props = schema['properties']
@@ -96,11 +98,10 @@ def rule_prompt(path=None, rule_type=None, required_only=True, save=True, verbos
 
     suggested_path = os.path.join(RULES_DIR, contents['name'])  # TODO: UPDATE BASED ON RULE STRUCTURE
     path = os.path.realpath(path or input('File path for rule [{}]: '.format(suggested_path)) or suggested_path)
-
-    rule = None
+    meta = {'creation_date': creation_date, 'updated_date': creation_date, 'maturity': 'development'}
 
     try:
-        rule = Rule(path, {'rule': contents})
+        rule = TOMLRule(path=Path(path), contents=TOMLRuleContents.from_dict({'rule': contents, 'metadata': meta}))
     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 +114,8 @@ 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=Path(path),
+                                contents=TOMLRuleContents.from_dict({'rule': contents, 'metadata': meta}))
             except kql.KqlParseError as e:
                 click.secho(e.args[0], fg='red', err=True)
                 click.pause()
@@ -127,7 +129,7 @@ def rule_prompt(path=None, rule_type=None, required_only=True, save=True, verbos
             break
 
     if save:
-        rule.save(verbose=True, as_rule=True)
+        rule.save_toml()
 
     if skipped:
         print('Did not set the following values because they are un-required when set to the default value')

detection_rules/devtools.py

@@ -4,6 +4,7 @@
 # 2.0.
 
 """CLI commands for internal detection_rules dev team."""
+import dataclasses
 import hashlib
 import io
 import json
@@ -23,10 +24,9 @@
 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, TOMLRuleContents, BaseQueryRuleData
 from .rule_loader import get_rule
-from .utils import get_path
-
+from .utils import get_path, dict_hash
 
 RULES_DIR = get_path('rules')
 
@@ -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']: 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)))
@@ -309,17 +309,27 @@ def deprecate_rule(ctx: click.Context, rule_file: str):
     version_info = load_versions()
     rule_file = Path(rule_file)
     contents = pytoml.loads(rule_file.read_text())
-    rule = Rule(path=rule_file, contents=contents)
+    rule = TOMLRule(path=rule_file, contents=contents)
 
     if rule.id not in version_info:
         click.echo('Rule has not been version locked and so does not need to be deprecated. '
                    'Delete the file or update the maturity to `development` instead')
         ctx.exit()
 
     today = time.strftime('%Y/%m/%d')
-    rule.metadata.update(updated_date=today, deprecation_date=today, maturity='deprecated')
+
+    new_meta = dataclasses.replace(rule.contents.metadata,
+                                   updated_date=today,
+                                   deprecation_date=today,
+                                   maturity='deprecated')
+    contents = dataclasses.replace(rule.contents, metadata=new_meta)
     deprecated_path = get_path('rules', '_deprecated', rule_file.name)
-    rule.save(new_path=deprecated_path, as_rule=True)
+
+    # create the new rule and save it
+    new_rule = TOMLRule(contents=contents, path=Path(deprecated_path))
+    new_rule.save_toml()
+
+    # remove the old rule
     rule_file.unlink()
     click.echo(f'Rule moved to {deprecated_path} - remember to git add this file')
 
@@ -375,27 +385,31 @@ def event_search(query, index, language, date_range, count, max_results, verbose
 def rule_event_search(ctx, rule_file, rule_id, date_range, count, max_results, verbose,
                       elasticsearch_client: Elasticsearch = None):
     """Search using a rule file against an Elasticsearch instance."""
-    rule = None
+    rule: TOMLRule
 
     if rule_id:
         rule = get_rule(rule_id, verbose=False)
     elif rule_file:
-        rule = Rule(rule_file, load_dump(rule_file))
+        rule = TOMLRule(path=rule_file, contents=TOMLRuleContents.from_dict(load_dump(rule_file)))
     else:
         client_error('Must specify a rule file or rule ID')
 
-    if rule.query and rule.contents.get('language'):
+    if isinstance(rule.contents.data, BaseQueryRuleData):
         if verbose:
             click.echo(f'Searching rule: {rule.name}')
 
-        rule_lang = rule.contents.get('language')
+        data = rule.contents.data
+        rule_lang = data.language
+
         if rule_lang == 'kuery':
-            language = None
+            language_flag = None
         elif rule_lang == 'eql':
-            language = True
+            language_flag = True
         else:
-            language = False
-        ctx.invoke(event_search, query=rule.query, index=rule.contents.get('index', ['*']), language=language,
+            language_flag = False
+
+        index = data.index or ['*']
+        ctx.invoke(event_search, query=data.query, index=index, language=language_flag,
                    date_range=date_range, count=count, max_results=max_results, verbose=verbose,
                    elasticsearch_client=elasticsearch_client)
     else:

detection_rules/docs.py

@@ -6,22 +6,24 @@
 """Create summary documents for a rule package."""
 from collections import defaultdict
 from pathlib import Path
+from typing import Optional, List
 
 import xlsxwriter
 
 from .attack import technique_lookup, matrix, attack_tm, tactics
 from .packaging import Package
+from .rule import ThreatMapping, TOMLRule
 
 
 class PackageDocument(xlsxwriter.Workbook):
     """Excel document for summarizing a rules package."""
 
-    def __init__(self, path, package):
+    def __init__(self, path, package: Package):
         """Create an excel workbook for the package."""
         self._default_format = {'font_name': 'Helvetica', 'font_size': 12}
         super(PackageDocument, self).__init__(path)
 
-        self.package: Package = package
+        self.package = package
         self.deprecated_rules = package.deprecated_rules
         self.production_rules = package.rules
 
@@ -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))
@@ -115,7 +117,7 @@ def add_summary(self):
             worksheet.write(row, 3, f'{num_techniques}/{total_techniques}', self.right_align)
             row += 1
 
-    def add_rule_details(self, rules=None, name='Rule Details'):
+    def add_rule_details(self, rules: Optional[List[TOMLRule]] = None, name='Rule Details'):
         """Add a worksheet for detailed metadata of rules."""
         if rules is None:
             rules = 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(rule.contents.to_api_format())
 
             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