import json from pathlib import Path import pytest from CommonServerPython import CommandResults from InvestigationSummaryParse import KillChain, Result, Source, parse_command TEST_DATA_DIR = Path(__file__).parent / "test_data" def _list_to_context(command_results: list[CommandResults]): return [result.to_context() for result in command_results] def _dump_test_file(file_name: str, content: dict): (TEST_DATA_DIR / file_name).write_text(json.dumps(content)) def _load_test_file(file_name: str): return json.loads((TEST_DATA_DIR / file_name).read_text()) def test_empty(): """ Given an empty context When calling parse_command Then make sure the result is the default """ assert _list_to_context(parse_command(context={})) == _load_test_file("empty_context.json") @pytest.mark.parametrize( "mocked_context,result,sources,search_value,test_index", [ ({}, Result.NOT_DETECTED, set(), "", 0), ({"incident": {"mitretacticname": "hello"}}, Result.SUSPICIOUS, {Source.EDR}, "hello", 1), ( { "incident": {"mitretacticname": "hello"}, "csfalconx": {"resource": {"sandbox": {"mitre_attacks": {"tactic": "hello"}}}}, }, Result.SUSPICIOUS, {Source.SANDBOX, Source.EDR}, "hello", 2, ), ({"MITREATTACK": [{"value": "hello"}]}, Result.SUSPICIOUS, {Source.EDR}, "hello", 3), ], ) def test_kill_chain(mocked_context: dict, result: Result, sources: set[Source], search_value: str, test_index: int): """ Given a class inheriting from KillChain When instantiating Then check the result of KillChain._parse_context """ class NewKillChain(KillChain): def __init__(self, context: dict): super().__init__(name="dummy", tactic="tactic", context=context, search_value=search_value) test_object = NewKillChain(mocked_context) assert test_object.result == result assert test_object.sources == sources assert test_object.to_context() == _load_test_file(f"kill_chain_{test_index}.json")