# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one # or more contributor license agreements. Licensed under the Elastic License # 2.0; you may not use this file except in compliance with the Elastic License # 2.0. """Helper functions for managing rules in the repository.""" import copy import dataclasses import io import json import textwrap import typing from collections import OrderedDict import toml from .schemas import definitions from .utils import cached SQ = "'" DQ = '"' TRIPLE_SQ = SQ * 3 TRIPLE_DQ = DQ * 3 @cached def get_preserved_fmt_fields(): from .rule import BaseRuleData preserved_keys = set() for field in dataclasses.fields(BaseRuleData): # type: dataclasses.Field if field.type in (definitions.Markdown, typing.Optional[definitions.Markdown]): preserved_keys.add(field.metadata.get("data_key", field.name)) return preserved_keys def cleanup_whitespace(val): if isinstance(val, str): return " ".join(line.strip() for line in val.strip().splitlines()) return val def nested_normalize(d, skip_cleanup=False): if isinstance(d, str): return d if skip_cleanup else cleanup_whitespace(d) elif isinstance(d, list): return [nested_normalize(val) for val in d] elif isinstance(d, dict): for k, v in d.items(): if k == 'query': # TODO: the linter still needs some work, but once up to par, uncomment to implement - kql.lint(v) # do not normalize queries d.update({k: v}) elif k in get_preserved_fmt_fields(): # let these maintain newlines and whitespace for markdown support d.update({k: nested_normalize(v, skip_cleanup=True)}) else: d.update({k: nested_normalize(v)}) return d else: return d def wrap_text(v, block_indent=0, join=False): """Block and indent a blob of text.""" v = ' '.join(v.split()) lines = textwrap.wrap(v, initial_indent=' ' * block_indent, subsequent_indent=' ' * block_indent, width=120, break_long_words=False, break_on_hyphens=False) lines = [line + '\n' for line in lines] # If there is a single line that contains a quote, add a new blank line to trigger multiline formatting if len(lines) == 1 and '"' in lines[0]: lines = lines + [''] return lines if not join else ''.join(lines) class NonformattedField(str): """Non-formatting class.""" def preserve_formatting_for_fields(data: OrderedDict, fields_to_preserve: list) -> OrderedDict: """Preserve formatting for specified nested fields in an action.""" def apply_preservation(target: OrderedDict, keys: list) -> None: """Apply NonformattedField preservation based on keys path.""" for key in keys[:-1]: # Iterate to the key, diving into nested dictionaries if key in target and isinstance(target[key], dict): target = target[key] else: # Cannot preserve formatting for missing or non-dict intermediate return final_key = keys[-1] if final_key in target: # Apply NonformattedField to the target field if it exists target[final_key] = NonformattedField(target[final_key]) for field_path in fields_to_preserve: keys = field_path.split('.') apply_preservation(data, keys) return data class RuleTomlEncoder(toml.TomlEncoder): """Generate a pretty form of toml.""" def __init__(self, _dict=dict, preserve=False): """Create the encoder but override some default functions.""" super(RuleTomlEncoder, self).__init__(_dict, preserve) self._old_dump_str = toml.TomlEncoder().dump_funcs[str] self._old_dump_list = toml.TomlEncoder().dump_funcs[list] self.dump_funcs[str] = self.dump_str self.dump_funcs[type(u"")] = self.dump_str self.dump_funcs[list] = self.dump_list self.dump_funcs[NonformattedField] = self.dump_str def dump_str(self, v): """Change the TOML representation to multi-line or single quote when logical.""" initial_newline = ['\n'] if isinstance(v, NonformattedField): # first line break is not forced like other multiline string dumps lines = v.splitlines(True) initial_newline = [] else: lines = wrap_text(v) multiline = len(lines) > 1 raw = (multiline or (DQ in v and SQ not in v)) and TRIPLE_DQ not in v if multiline: if raw: return "".join([TRIPLE_DQ] + initial_newline + lines + [TRIPLE_DQ]) else: return "\n".join([TRIPLE_SQ] + [self._old_dump_str(line)[1:-1] for line in lines] + [TRIPLE_SQ]) elif raw: return u"'{:s}'".format(lines[0]) return self._old_dump_str(v) def _dump_flat_list(self, v): """A slightly tweaked version of original dump_list, removing trailing commas.""" if not v: return "[]" retval = "[" + str(self.dump_value(v[0])) + "," for u in v[1:]: retval += " " + str(self.dump_value(u)) + "," retval = retval.rstrip(',') + "]" return retval def dump_list(self, v): """Dump a list more cleanly.""" if all([isinstance(d, str) for d in v]) and sum(len(d) + 3 for d in v) > 100: dump = [] for item in v: if len(item) > (120 - 4 - 3 - 3) and ' ' in item: dump.append(' """\n{} """'.format(wrap_text(item, block_indent=4, join=True))) else: dump.append(' ' * 4 + self.dump_value(item)) return '[\n{},\n]'.format(',\n'.join(dump)) if v and all(isinstance(i, dict) for i in v): # Compact inline format for lists of dictionaries with proper indentation retval = "\n" + ' ' * 2 + "[\n" retval += ",\n".join([' ' * 4 + self.dump_inline_table(u).strip() for u in v]) retval += "\n" + ' ' * 2 + "]\n" return retval return self._dump_flat_list(v) def toml_write(rule_contents, outfile=None): """Write rule in TOML.""" def write(text, nl=True): if outfile: outfile.write(text) if nl: outfile.write(u"\n") else: print(text, end='' if not nl else '\n') encoder = RuleTomlEncoder() contents = copy.deepcopy(rule_contents) needs_close = False def order_rule(obj): if isinstance(obj, dict): obj = OrderedDict(sorted(obj.items())) for k, v in obj.items(): if isinstance(v, dict) or isinstance(v, list): obj[k] = order_rule(v) if isinstance(obj, list): for i, v in enumerate(obj): if isinstance(v, dict) or isinstance(v, list): obj[i] = order_rule(v) obj = sorted(obj, key=lambda x: json.dumps(x)) return obj def _do_write(_data, _contents): query = None threat_query = None if _data == 'rule': # - We want to avoid the encoder for the query and instead use kql-lint. # - Linting is done in rule.normalize() which is also called in rule.validate(). # - Until lint has tabbing, this is going to result in all queries being flattened with no wrapping, # but will at least purge extraneous white space query = contents['rule'].pop('query', '').strip() # - As tags are expanding, we may want to reconsider the need to have them in alphabetical order # tags = contents['rule'].get("tags", []) # # if tags and isinstance(tags, list): # contents['rule']["tags"] = list(sorted(set(tags))) threat_query = contents['rule'].pop('threat_query', '').strip() top = OrderedDict() bottom = OrderedDict() for k in sorted(list(_contents)): v = _contents.pop(k) if k == 'actions': # explicitly preserve formatting for message field in actions preserved_fields = ["params.message"] v = [preserve_formatting_for_fields(action, preserved_fields) for action in v] if v is not None else [] if k == 'filters': # explicitly preserve formatting for value field in filters preserved_fields = ["meta.value"] v = [preserve_formatting_for_fields(meta, preserved_fields) for meta in v] if v is not None else [] if k == 'note' and isinstance(v, str): # Transform instances of \ to \\ as calling write will convert \\ to \. # This will ensure that the output file has the correct number of backslashes. v = v.replace("\\", "\\\\") if k == 'setup' and isinstance(v, str): # Transform instances of \ to \\ as calling write will convert \\ to \. # This will ensure that the output file has the correct number of backslashes. v = v.replace("\\", "\\\\") if k == 'description' and isinstance(v, str): # Transform instances of \ to \\ as calling write will convert \\ to \. # This will ensure that the output file has the correct number of backslashes. v = v.replace("\\", "\\\\") if k == 'osquery' and isinstance(v, list): # Specifically handle transform.osquery queries for osquery_item in v: if 'query' in osquery_item and isinstance(osquery_item['query'], str): # Transform instances of \ to \\ as calling write will convert \\ to \. # This will ensure that the output file has the correct number of backslashes. osquery_item['query'] = osquery_item['query'].replace("\\", "\\\\") if isinstance(v, dict): bottom[k] = OrderedDict(sorted(v.items())) elif isinstance(v, list): if any([isinstance(value, (dict, list)) for value in v]): bottom[k] = v else: top[k] = v elif k in get_preserved_fmt_fields(): top[k] = NonformattedField(v) else: top[k] = v if query: top.update({'query': "XXxXX"}) if threat_query: top.update({'threat_query': "XXxXX"}) top.update(bottom) top = toml.dumps(OrderedDict({data: top}), encoder=encoder) # we want to preserve the threat_query format, but want to modify it in the context of encoded dump if threat_query: formatted_threat_query = "\nthreat_query = '''\n{}\n'''{}".format(threat_query, '\n\n' if bottom else '') top = top.replace('threat_query = "XXxXX"', formatted_threat_query) # we want to preserve the query format, but want to modify it in the context of encoded dump if query: formatted_query = "\nquery = '''\n{}\n'''{}".format(query, '\n\n' if bottom else '') top = top.replace('query = "XXxXX"', formatted_query) write(top) try: if outfile and not isinstance(outfile, io.IOBase): needs_close = True outfile = open(outfile, 'w') for data in ('metadata', 'transform', 'rule'): _contents = contents.get(data, {}) if not _contents: continue order_rule(_contents) _do_write(data, _contents) finally: if needs_close and hasattr(outfile, "close"): outfile.close()