Skip to content

Commit c0af222

Browse files
authored
Move Rule into a dataclass (#1029)
* WIP: Convert Rule to a dataclass * Fix make release * Lint fixes * Remove dead code * Fix lint and tests * Use Python 3.8 in GitHub actions * Update README to 3.8+ * Add Python 3.8 assertion * Fix is_dirty property * Remove incorrect pop from contents * Add mixin with from_dict() and to_dict() methods * Bypass validation for deprecated rules * Fix rule_prompt * Fix dict_hash usage * Fix rule_event_search * Switch to definitions.Date * Fix toml-lint command, ignoring 'unneeded defaults' * Moved severity Literal to definitions.Severity * Remove BaseMarshmallowDataclass * Fix lint and tests * Add maturity to metadata for rule prompt loop * Fix typo in devtools * Use rule loader to load single rule in toml-lint * Add Schema hint to __schema method * Add MITREAttackURL definition * Fix is_dirty to compare sha<-->sha * Normalize the autoformatted rule output for API and toml-lint * Make the package hash match * Make the rule object mutable but not rule contents * Restore the rules
1 parent cc6711c commit c0af222

27 files changed

+829
-696
lines changed

.github/workflows/pythonpackage.yml

+2-2
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,10 @@ jobs:
1414
steps:
1515
- uses: actions/checkout@v2
1616

17-
- name: Set up Python 3.7
17+
- name: Set up Python 3.8
1818
uses: actions/setup-python@v2
1919
with:
20-
python-version: 3.7
20+
python-version: 3.8
2121

2222
- name: Install dependencies
2323
run: |

Makefile

+1-1
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ all: release
1414

1515
$(VENV):
1616
pip install virtualenv
17-
virtualenv $(VENV) --python=python3.7
17+
virtualenv $(VENV) --python=python3.8
1818
$(PIP) install -r requirements.txt
1919
$(PIP) install setuptools -U
2020

README.md

+2-2
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
[![Supported Python versions](https://img.shields.io/badge/python-3.7+-yellow.svg)](https://www.python.org/downloads/)
1+
[![Supported Python versions](https://img.shields.io/badge/python-3.8+-yellow.svg)](https://www.python.org/downloads/)
22
[![Unit Tests](https://github.com/elastic/detection-rules/workflows/Unit%20Tests/badge.svg)](https://github.com/elastic/detection-rules/actions)
33
[![Chat](https://img.shields.io/badge/chat-%23security--detection--rules-blueviolet)](https://ela.st/slack)
44

@@ -35,7 +35,7 @@ Detection Rules contains more than just static rule files. This repository also
3535

3636
## Getting started
3737

38-
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:
38+
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:
3939
```console
4040
$ pip install -r requirements.txt
4141
Collecting jsl==0.2.4

detection_rules/__init__.py

+17-11
Original file line numberDiff line numberDiff line change
@@ -4,17 +4,23 @@
44
# 2.0.
55

66
"""Detection rules."""
7-
from . import devtools
8-
from . import docs
9-
from . import eswrap
10-
from . import kbwrap
11-
from . import main
12-
from . import mappings
13-
from . import misc
14-
from . import rule_formatter
15-
from . import rule_loader
16-
from . import schemas
17-
from . import utils
7+
import sys
8+
9+
assert (3, 8) <= sys.version_info < (4, 0), "Only Python 3.8+ supported"
10+
11+
from . import ( # noqa: E402
12+
devtools,
13+
docs,
14+
eswrap,
15+
kbwrap,
16+
main,
17+
mappings,
18+
misc,
19+
rule_formatter,
20+
rule_loader,
21+
schemas,
22+
utils
23+
)
1824

1925
__all__ = (
2026
'devtools',

detection_rules/__main__.py

+6-1
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,13 @@
66
# coding=utf-8
77
"""Shell for detection-rules."""
88
import os
9+
import sys
10+
911
import click
10-
from .main import root
12+
13+
assert (3, 8) <= sys.version_info < (4, 0), "Only Python 3.8+ supported"
14+
15+
from .main import root # noqa: E402
1116

1217
CURR_DIR = os.path.dirname(os.path.abspath(__file__))
1318
CLI_DIR = os.path.dirname(CURR_DIR)

detection_rules/cli_utils.py

+11-9
Original file line numberDiff line numberDiff line change
@@ -4,24 +4,27 @@
44
# 2.0.
55

66
import copy
7+
import datetime
78
import os
9+
from pathlib import Path
810

911
import click
1012

1113
import kql
1214
from . import ecs
1315
from .attack import matrix, tactics, build_threat_map_entry
14-
from .rule import Rule
16+
from .rule import TOMLRule, TOMLRuleContents
1517
from .schemas import CurrentSchema
1618
from .utils import clear_caches, get_path
1719

1820
RULES_DIR = get_path("rules")
1921

2022

21-
def rule_prompt(path=None, rule_type=None, required_only=True, save=True, verbose=False, **kwargs) -> Rule:
23+
def rule_prompt(path=None, rule_type=None, required_only=True, save=True, verbose=False, **kwargs) -> TOMLRule:
2224
"""Prompt loop to build a rule."""
2325
from .misc import schema_prompt
2426

27+
creation_date = datetime.date.today().strftime("%Y/%m/%d")
2528
if verbose and path:
2629
click.echo(f'[+] Building rule for {path}')
2730

@@ -32,8 +35,7 @@ def rule_prompt(path=None, rule_type=None, required_only=True, save=True, verbos
3235
kwargs.update(kwargs.pop('rule'))
3336

3437
rule_type = rule_type or kwargs.get('type') or \
35-
click.prompt('Rule type ({})'.format(', '.join(CurrentSchema.RULE_TYPES)),
36-
type=click.Choice(CurrentSchema.RULE_TYPES))
38+
click.prompt('Rule type', type=click.Choice(CurrentSchema.RULE_TYPES))
3739

3840
schema = CurrentSchema.get_schema(role=rule_type)
3941
props = schema['properties']
@@ -96,11 +98,10 @@ def rule_prompt(path=None, rule_type=None, required_only=True, save=True, verbos
9698

9799
suggested_path = os.path.join(RULES_DIR, contents['name']) # TODO: UPDATE BASED ON RULE STRUCTURE
98100
path = os.path.realpath(path or input('File path for rule [{}]: '.format(suggested_path)) or suggested_path)
99-
100-
rule = None
101+
meta = {'creation_date': creation_date, 'updated_date': creation_date, 'maturity': 'development'}
101102

102103
try:
103-
rule = Rule(path, {'rule': contents})
104+
rule = TOMLRule(path=Path(path), contents=TOMLRuleContents.from_dict({'rule': contents, 'metadata': meta}))
104105
except kql.KqlParseError as e:
105106
if e.error_msg == 'Unknown field':
106107
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
113114
while True:
114115
try:
115116
contents['query'] = click.edit(contents['query'], extension='.eql')
116-
rule = Rule(path, {'rule': contents})
117+
rule = TOMLRule(path=Path(path),
118+
contents=TOMLRuleContents.from_dict({'rule': contents, 'metadata': meta}))
117119
except kql.KqlParseError as e:
118120
click.secho(e.args[0], fg='red', err=True)
119121
click.pause()
@@ -127,7 +129,7 @@ def rule_prompt(path=None, rule_type=None, required_only=True, save=True, verbos
127129
break
128130

129131
if save:
130-
rule.save(verbose=True, as_rule=True)
132+
rule.save_toml()
131133

132134
if skipped:
133135
print('Did not set the following values because they are un-required when set to the default value')

detection_rules/devtools.py

+29-15
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
# 2.0.
55

66
"""CLI commands for internal detection_rules dev team."""
7+
import dataclasses
78
import hashlib
89
import io
910
import json
@@ -23,10 +24,9 @@
2324
from .main import root
2425
from .misc import PYTHON_LICENSE, add_client, GithubClient, Manifest, client_error, getdefault
2526
from .packaging import PACKAGE_FILE, Package, manage_versions, RELEASE_DIR
26-
from .rule import Rule
27+
from .rule import TOMLRule, TOMLRuleContents, BaseQueryRuleData
2728
from .rule_loader import get_rule
28-
from .utils import get_path
29-
29+
from .utils import get_path, dict_hash
3030

3131
RULES_DIR = get_path('rules')
3232

@@ -96,7 +96,7 @@ def kibana_diff(rule_id, repo, branch, threads):
9696
repo_hashes = {r.id: r.get_hash() for r in rules.values()}
9797

9898
kibana_rules = {r['rule_id']: r for r in get_kibana_rules(repo=repo, branch=branch, threads=threads).values()}
99-
kibana_hashes = {r['rule_id']: Rule.dict_hash(r) for r in kibana_rules.values()}
99+
kibana_hashes = {r['rule_id']: dict_hash(r) for r in kibana_rules.values()}
100100

101101
missing_from_repo = list(set(kibana_hashes).difference(set(repo_hashes)))
102102
missing_from_kibana = list(set(repo_hashes).difference(set(kibana_hashes)))
@@ -309,17 +309,27 @@ def deprecate_rule(ctx: click.Context, rule_file: str):
309309
version_info = load_versions()
310310
rule_file = Path(rule_file)
311311
contents = pytoml.loads(rule_file.read_text())
312-
rule = Rule(path=rule_file, contents=contents)
312+
rule = TOMLRule(path=rule_file, contents=contents)
313313

314314
if rule.id not in version_info:
315315
click.echo('Rule has not been version locked and so does not need to be deprecated. '
316316
'Delete the file or update the maturity to `development` instead')
317317
ctx.exit()
318318

319319
today = time.strftime('%Y/%m/%d')
320-
rule.metadata.update(updated_date=today, deprecation_date=today, maturity='deprecated')
320+
321+
new_meta = dataclasses.replace(rule.contents.metadata,
322+
updated_date=today,
323+
deprecation_date=today,
324+
maturity='deprecated')
325+
contents = dataclasses.replace(rule.contents, metadata=new_meta)
321326
deprecated_path = get_path('rules', '_deprecated', rule_file.name)
322-
rule.save(new_path=deprecated_path, as_rule=True)
327+
328+
# create the new rule and save it
329+
new_rule = TOMLRule(contents=contents, path=Path(deprecated_path))
330+
new_rule.save_toml()
331+
332+
# remove the old rule
323333
rule_file.unlink()
324334
click.echo(f'Rule moved to {deprecated_path} - remember to git add this file')
325335

@@ -375,27 +385,31 @@ def event_search(query, index, language, date_range, count, max_results, verbose
375385
def rule_event_search(ctx, rule_file, rule_id, date_range, count, max_results, verbose,
376386
elasticsearch_client: Elasticsearch = None):
377387
"""Search using a rule file against an Elasticsearch instance."""
378-
rule = None
388+
rule: TOMLRule
379389

380390
if rule_id:
381391
rule = get_rule(rule_id, verbose=False)
382392
elif rule_file:
383-
rule = Rule(rule_file, load_dump(rule_file))
393+
rule = TOMLRule(path=rule_file, contents=TOMLRuleContents.from_dict(load_dump(rule_file)))
384394
else:
385395
client_error('Must specify a rule file or rule ID')
386396

387-
if rule.query and rule.contents.get('language'):
397+
if isinstance(rule.contents.data, BaseQueryRuleData):
388398
if verbose:
389399
click.echo(f'Searching rule: {rule.name}')
390400

391-
rule_lang = rule.contents.get('language')
401+
data = rule.contents.data
402+
rule_lang = data.language
403+
392404
if rule_lang == 'kuery':
393-
language = None
405+
language_flag = None
394406
elif rule_lang == 'eql':
395-
language = True
407+
language_flag = True
396408
else:
397-
language = False
398-
ctx.invoke(event_search, query=rule.query, index=rule.contents.get('index', ['*']), language=language,
409+
language_flag = False
410+
411+
index = data.index or ['*']
412+
ctx.invoke(event_search, query=data.query, index=index, language=language_flag,
399413
date_range=date_range, count=count, max_results=max_results, verbose=verbose,
400414
elasticsearch_client=elasticsearch_client)
401415
else:

detection_rules/docs.py

+15-13
Original file line numberDiff line numberDiff line change
@@ -6,22 +6,24 @@
66
"""Create summary documents for a rule package."""
77
from collections import defaultdict
88
from pathlib import Path
9+
from typing import Optional, List
910

1011
import xlsxwriter
1112

1213
from .attack import technique_lookup, matrix, attack_tm, tactics
1314
from .packaging import Package
15+
from .rule import ThreatMapping, TOMLRule
1416

1517

1618
class PackageDocument(xlsxwriter.Workbook):
1719
"""Excel document for summarizing a rules package."""
1820

19-
def __init__(self, path, package):
21+
def __init__(self, path, package: Package):
2022
"""Create an excel workbook for the package."""
2123
self._default_format = {'font_name': 'Helvetica', 'font_size': 12}
2224
super(PackageDocument, self).__init__(path)
2325

24-
self.package: Package = package
26+
self.package = package
2527
self.deprecated_rules = package.deprecated_rules
2628
self.production_rules = package.rules
2729

@@ -47,16 +49,16 @@ def _get_attack_coverage(self):
4749
coverage = defaultdict(lambda: defaultdict(lambda: defaultdict(int)))
4850

4951
for rule in self.package.rules:
50-
threat = rule.contents.get('threat')
52+
threat = rule.contents.data.threat
5153
sub_dir = Path(rule.path).parent.name
5254

5355
if threat:
5456
for entry in threat:
55-
tactic = entry['tactic']
56-
techniques = entry.get('technique', [])
57+
tactic = entry.tactic
58+
techniques = entry.technique or []
5759
for technique in techniques:
58-
if technique['id'] in matrix[tactic['name']]:
59-
coverage[tactic['name']][technique['id']][sub_dir] += 1
60+
if technique.id in matrix[tactic.name]:
61+
coverage[tactic.name][technique.id][sub_dir] += 1
6062

6163
return coverage
6264

@@ -85,10 +87,10 @@ def add_summary(self):
8587

8688
tactic_counts = defaultdict(int)
8789
for rule in self.package.rules:
88-
threat = rule.contents.get('threat')
90+
threat = rule.contents.data.threat
8991
if threat:
9092
for entry in threat:
91-
tactic_counts[entry['tactic']['name']] += 1
93+
tactic_counts[entry.tactic.name] += 1
9294

9395
worksheet.write(row, 0, "Total Production Rules")
9496
worksheet.write(row, 1, len(self.production_rules))
@@ -115,7 +117,7 @@ def add_summary(self):
115117
worksheet.write(row, 3, f'{num_techniques}/{total_techniques}', self.right_align)
116118
row += 1
117119

118-
def add_rule_details(self, rules=None, name='Rule Details'):
120+
def add_rule_details(self, rules: Optional[List[TOMLRule]] = None, name='Rule Details'):
119121
"""Add a worksheet for detailed metadata of rules."""
120122
if rules is None:
121123
rules = self.production_rules
@@ -134,9 +136,9 @@ def add_rule_details(self, rules=None, name='Rule Details'):
134136
)
135137

136138
for row, rule in enumerate(rules, 1):
137-
flat_mitre = rule.get_flat_mitre()
138-
rule_contents = {'tactics': flat_mitre['tactic_names'], 'techniques': flat_mitre['technique_ids']}
139-
rule_contents.update(rule.contents.copy())
139+
flat_mitre = ThreatMapping.flatten(rule.contents.data.threat)
140+
rule_contents = {'tactics': flat_mitre.tactic_names, 'techniques': flat_mitre.technique_ids}
141+
rule_contents.update(rule.contents.to_api_format())
140142

141143
for column, field in enumerate(metadata_fields):
142144
value = rule_contents.get(field)

detection_rules/eswrap.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
from .main import root
2222
from .misc import add_params, client_error, elasticsearch_options
2323
from .utils import format_command_options, normalize_timing_and_sort, unix_time_to_formatted, get_path
24-
from .rule import Rule
24+
from .rule import TOMLRule
2525
from .rule_loader import get_rule, rta_mappings
2626

2727

@@ -195,7 +195,7 @@ def search(self, query, language, index: Union[str, list] = '*', start_time=None
195195

196196
return results
197197

198-
def search_from_rule(self, *rules: Rule, start_time=None, end_time='now', size=None):
198+
def search_from_rule(self, *rules: TOMLRule, start_time=None, end_time='now', size=None):
199199
"""Search an elasticsearch instance using a rule."""
200200
from .misc import nested_get
201201

0 commit comments

Comments
 (0)