Files
splunk-security_content/bin/required_fields_cleaner.py
2022-05-23 17:00:20 -07:00

452 lines
17 KiB
Python

import collections
import yaml
import sys
import re
import json
import pprint
import argparse
import pathlib
#find datamodel=SOMETHING, allowing for whitespace on either side of =
DATAMODEL_PATTERN = r"datamodel\s*=?\s*\S*"
QUOTATIONS_PATTERN = r'''(["'])(?:(?=(\\?))\2.)*?\1'''
KEY_VALUE_PATTERN = r"[a-zA-Z0-9_]+\.[a-zA-Z0-9_]+"
SMALL_INDENT = ' -'
def dictPrint(d:dict):
printer = pprint.PrettyPrinter(indent=3)
printer.pprint(d)
def parse_datamodel(datamodel_name, object)->dict:
submodules = {}
if "objects" in object:
raise(Exception(("nested objects!")))
#print(object.keys())
#print(object['calculations'])
#input("W")
allFields = []
if 'fields' in object:
allFields += object['fields']
if 'calculations' in object:
for calc_dict in object['calculations']:
if 'outputFields' in calc_dict:
allFields += calc_dict['outputFields']
for field in allFields:
fieldName = field['fieldName']
if fieldName in submodules:
raise(Exception(f"Duplicate Field {field} in {datamodel_name}"))
submodel_fieldname = f"{object['objectName']}.{fieldName}"
model_submodel_fieldname = f"{datamodel_name}.{submodel_fieldname}"
submodules[fieldName] = {"field_name": fieldName, "submodel.fieldname": submodel_fieldname, "datamodel.submodel.field_name": model_submodel_fieldname}
return submodules
def load_datamodels_from_directory(datamodels_directory:pathlib.PosixPath)->dict:
print("Loading datamodel templates...", end='', flush=True)
all_models = {}
datamodel_filenames = list(datamodels_directory.glob("*.json"))
#Iterate through all of the YML files (*.json) in the directory that we passed.
for datamodel_filename in datamodel_filenames:
#Load the YML File
with open(datamodel_filename, "r") as model_stream:
json_datamodel = json.load(model_stream)
#Get the name of the model
model_name = json_datamodel['modelName']
model_fields = {}
#Load all the submodels from the YML file. They are called 'objects'
for submodel in json_datamodel['objects']:
model_fields[submodel['objectName']] = parse_datamodel(model_name, submodel)
all_models[model_name] = model_fields
print(f"[{len(datamodel_filenames):4d}] datamodel templates loaded")
return all_models
def get_datamodels(search:str, filename:str)->dict:
#print(search)
matches = re.findall(DATAMODEL_PATTERN, search)
#get the models and strip out the whitespace
#If we want different dot-separated versions
models_and_submodels = [model.replace(" ","=").split("=")[1].strip() for model in matches]
results = {}
for model_and_submodel in models_and_submodels:
if model_and_submodel.count(".") != 1:
print(f"datamodel {model_and_submodel} is not in model.submodel format in {filename}")
else:
model, submodel = model_and_submodel.split(".")
if model in results:
results[model].append(submodel)
else:
results[model]=[submodel]
#If we only care about the top level
#models_only = [model.split(".")[0].strip() for model in models_and_submodels]
return results
def remove_datamodel_submodel_from_fields(defined_datamodels:dict ,fields:set[str])->set[str]:
trimmed_fields = set()
for field in fields:
parts = field.split(".")
if len(parts) == 2:
if (parts[0] not in defined_datamodels) or (parts[1] not in defined_datamodels[parts[0]]):
trimmed_fields.add(field)
return trimmed_fields
def get_datamodel_fields(defined_datamodels:dict, search:str)->set[str]:
quoted_text_removed = re.sub(QUOTATIONS_PATTERN, "", search)
all_fields = set(re.findall(KEY_VALUE_PATTERN, quoted_text_removed))
return remove_datamodel_submodel_from_fields(defined_datamodels, all_fields)
def check_for_presence_of_fields(filename:str, container:dict, field_names:list[str])->bool:
missing_fields = []
for field_name in field_names:
if field_name not in container:
missing_field = True
missing_fields.append(field_name)
if len(missing_fields) > 0:
#One of more fields that we were required to find were not found
print(f"Detection file {filename} is missing the following keys: [{sorted(missing_fields)}]")
return False
return True
def verify_dataset_field(filename:str, detection_file_data:dict)->tuple[bool,bool,dict]:
#Always validate that the same dataset is linked in the detection as in the test file
#Update the filename to refer to the test instead of the detection
test_filename = filename.replace(".yml", ".test.yml",1).replace("/detections/","/tests/",1)
test_file_data_files = set()
try:
with open(test_filename, "r") as test_stream:
test_datamodel = yaml.safe_load(test_stream)
except Exception as e:
print(f"Error while parsing {test_filename}: {str(e)}")
return (False,False,detection_file_data)
#Get all of the links to test datasets in the test file
if 'tests' in test_datamodel:
for test in test_datamodel['tests']:
if 'attack_data' in test:
for data_file in test['attack_data']:
if 'data' in data_file:
test_file_data_files.add(data_file['data'])
else:
print(f"'data' field not found in attack_data for {test_filename}: {dictPrint(data_file)}")
return (False, False,detection_file_data)
else:
print(f"'attack_data' not found in the test file {test_filename}")
return (False, False,detection_file_data)
else:
print(f"'tests' field not found in attack_data for {test_filename}")
return(False, False,detection_file_data)
#Get all of the datasets that are referenced in the detection YML file
if 'dataset' in detection_file_data['tags']:
detection_file_data_files = set(detection_file_data['tags']['dataset'])
else:
detection_file_data_files = set()
#Check to see if the test and detection ymls have the same datasets
diff = test_file_data_files.symmetric_difference(detection_file_data_files)
if len(diff) > 0:
print(f"{len(diff)} Error(s) for {filename:}")
for data_file in (test_file_data_files - detection_file_data_files):
print(f"\tTest file references file NOT detection file:\n\t{SMALL_INDENT} {data_file}")
for data_file in (detection_file_data_files - test_file_data_files):
print(f"\tDetection file references file NOT included in the test file:\n\t{SMALL_INDENT} {data_file}")
print("")
return (False,False,{})
#Success, no updates need to be made!
return (True, False, detection_file_data)
def validate_datamodels(filename:str, search_datamodels:dict, defined_datamodels:dict)->tuple[bool,dict]:
errors = []
submodels_with_fields = {}
for search_datamodel in search_datamodels:
if search_datamodel not in defined_datamodels:
errors.append(f"Failed to find the datamodel {search_datamodel} in predefined datamodels {defined_datamodels.keys()}")
continue
for submodel in search_datamodels[search_datamodel]:
if submodel not in defined_datamodels[search_datamodel]:
errors.append(f"Failed to find {'.'.join([search_datamodel,submodel])} in {search_datamodel}: {defined_datamodels[search_datamodel].keys()} ")
elif submodel not in submodels_with_fields:
submodels_with_fields[submodel] = defined_datamodels[search_datamodel][submodel]
if len(errors) == 0:
return (True, submodels_with_fields)
else:
#There was at least one error, print it out
print(f"Error(s) validating the search datamodel for {filename}:")
for error in errors:
print(f"{SMALL_INDENT} {error}")
return (False, {})
def validate_fields(filename:str, submodels_with_fields:dict, search_submodel_fields:set)->bool:
errors = []
for field in search_submodel_fields:
tokens = field.split(".")
if len(tokens) <= 1:
errors.append(f"Failed to find a submodel and a model for {field}")
elif len(tokens) == 2:
submodel, field = tokens
try:
if field not in submodels_with_fields[submodel]:
errors.append(f"[{submodel}] does not contain field [{field}]")
except Exception as e:
errors.append(f"Problem resolving [{submodel}.{field}]: [{str(e)}] is not a valid submodel")
else:
errors.append(f"Found more than just a submodel and field for {field}")
if len(errors) == 0:
return True
else:
#There was at least one error, print it out
print(f"Error(s) validating the datamodel fields in the search for {filename}:")
for error in errors:
print(f"{SMALL_INDENT} {error}")
return False
def update_datamodels(filename:str, defined_datamodels:dict, detection_file_data: dict)->tuple[bool,bool,dict]:
if not check_for_presence_of_fields(filename, detection_file_data, ["search","tags"]):
return (False,False,{})
if not check_for_presence_of_fields(filename, detection_file_data["tags"], ["required_fields"]):
return (False,False,{})
if not check_for_presence_of_fields(filename, detection_file_data, ["datamodel"]):
return (False,False,{})
#Pull the datamodel(s) from the search
search_datamodels = get_datamodels(detection_file_data['search'],filename)
#Validate that the datamodel(s) we pulled from the search exist
success, submodels_with_fields = validate_datamodels(filename, search_datamodels, defined_datamodels)
if success is False:
return (success, False, {})
if len(detection_file_data["datamodel"]) == 0 and len(search_datamodels) == 0:
if "|tstats" in detection_file_data["search"] or "| tstats" in detection_file_data["search"]:
print(f"{SMALL_INDENT} Error - {filename} contains tstats but no datamodels. Raw search:\n\t{detection_file_data['search']}")
return (False,False,{})
else:
#print(f"{filename} is a search that does not contain any datamodels. required_fields will not be validated or updated")
return (True,False,detection_file_data)
#Pull all the the datamodel fields from the search
search_submodel_fields = get_datamodel_fields(defined_datamodels, detection_file_data['search'])
#Validate the fields we pulled from the search
success = validate_fields(filename, submodels_with_fields, search_submodel_fields)
if success is False:
return (success, False, detection_file_data)
search_datamodel_and_submodel_set = set()
for model in search_datamodels:
for submodel in search_datamodels[model]:
search_datamodel_and_submodel_set.add(f"{model}.{submodel}")
search_datamodel_and_submodel_list = sorted(list(search_datamodel_and_submodel_set))
if search_datamodel_and_submodel_list != detection_file_data['datamodel']:
print(f"Updated datamodel(s) in {filename}:\n\t {detection_file_data['datamodel']} --> {search_datamodel_and_submodel_list}")
detection_file_data['datamodel'] = search_datamodel_and_submodel_list
updated = True
else:
#No updates to the datamodel field
updated = False
#print(f"{filename}:\n\t{datamodel_fields}")
#sys.exit(0)
return (True, updated, detection_file_data)
def validate_detection(filename:str, defined_datamodels:dict)->tuple[bool,bool, dict]:
#Use this variable to determine whether or not the yaml data is updated and should
#be rewritten to the file
dataset_updates = False
#load the yaml file
try:
with open(filename, "r") as detection_file:
detection_file_data = yaml.safe_load(detection_file)
except Exception as e:
print(f"Error processing file [{filename}]: {str(e)}")
return (False,False,{})
#Ensure the 'type' field exists
if not check_for_presence_of_fields(filename, detection_file_data, ['type']):
raise(Exception("YML missing field name 'type'"))
if detection_file_data['type'] not in ["Anomaly", "Hunting", "TTP" ]:
#This is not one of the detection types we want to update
return (True,False,{})
if not check_for_presence_of_fields(filename, detection_file_data, ['search']):
raise(Exception("YML missing field name 'search'"))
if not check_for_presence_of_fields(filename, detection_file_data, ['tags']):
raise(Exception("YML missing field name 'tags'"))
if not check_for_presence_of_fields(filename, detection_file_data, ['datamodel']):
raise(Exception("YML missing field name 'datamodel'"))
if not check_for_presence_of_fields(filename, detection_file_data['tags'], ['required_fields']):
raise(Exception("YML missing field name ['tags']['required_fields']"))
#Always validate that the same dataset is linked in the detection as in the test file
success, dataset_updates, detection_file_data = verify_dataset_field(filename, detection_file_data)
if success is False:
return (success, False, {})
success, datamodel_updates, detection_file_data = update_datamodels(filename, defined_datamodels, detection_file_data)
if success is False:
return (success, False, {})
#Updates requires that one, or both, of the dataset and datamodels are updated
updates = dataset_updates | datamodel_updates
return (success, updates, detection_file_data)
def get_detection_filenames(detection_directory:pathlib.PosixPath, sort=True)->list[str]:
detection_filenames = [str(p) for p in detection_directory.glob("**/[!ssa___]*.yml") if "/deprecated/" not in str(p) and "/experimental/" not in str(p)]
print("Enumerating detections........", end='', flush=True)
print(f"[{len(detection_filenames):4d}] detections found")
if sort:
return sorted(detection_filenames)
else:
return detection_filenames
def output_updated_detection(detection_filename:str, detection_data:dict)->bool:
try:
with open(detection_filename, "w") as updated_detection:
yaml.dump(detection_data, updated_detection, sort_keys=False)
print(f"\t{detection_filename} changes written to disk")
return True
except Exception as e:
print(f"Error writing {detection_filename} to disk: {str(e)}")
return False
def main():
parser =argparse.ArgumentParser(prog="fields_extractor_and_updater",
description="This tool parses the search field and uses it to "
"determine and update the values in the datamodel "
"and required_fields portions of the detection yml.")
parser.add_argument("mode",
choices=["check", "update"],
help="Determines whether detections should be checked or updated. "
"Check will print the results, but update will update "
"(and overrwrite) files that require changes.")
parser.add_argument("-d",
"--detection_directory",
type=pathlib.Path,
required=False,
default = "../detections/",
help="Root directory (or filename) to update. "
"Please note that this should NOT be a regular expression")
parser.add_argument("-dm",
"--datamodel_directory",
type=pathlib.Path,
required=False,
default="base_datamodels",
help="Root directory containing the standard datamodels. "
"Please note that this should NOT be a regular expression")
args = parser.parse_args()
#Load all the datamodels to validate against
defined_datamodels = load_datamodels_from_directory(args.datamodel_directory)
#Get all the files that we will process
detection_filenames = get_detection_filenames(args.detection_directory)
all_success = True
#Update/check each of the files
for detection_filename in detection_filenames:
try:
success, dataset_has_updates, updated_dataset = validate_detection(detection_filename, defined_datamodels)
all_success &= success #accumulate any errors that may occur here
if success and dataset_has_updates and (args.mode == "update"):
output_updated_detection(detection_filename, updated_dataset)
except Exception as e:
print(f"Error processing {detection_filename}: {str(e)}")
if all_success is True:
print("All required fields pass!")
sys.exit(0)
else:
print("At least one error occurred - check the output for details")
sys.exit(1)
if __name__ == "__main__":
main()