def test_apply_diff_merge_detetion_points(): original_root = RootAnalysis() original_observable = original_root.add_observable("test", "test") modified_root = original_root.copy() modified_observable = modified_root.get_observable(original_observable) modified_observable.add_detection_point("test") target_root = original_root.copy() target_observable = target_root.get_observable(original_observable) assert not target_observable.has_detection_points() target_observable.apply_diff_merge(original_observable, modified_observable) assert target_observable.has_detection_points # exists before but not after original_root = RootAnalysis() original_observable = original_root.add_observable("test", "test") modified_root = original_root.copy() modified_observable = modified_root.get_observable(original_observable) target_root = original_root.copy() target_observable = target_root.get_observable(original_observable) original_observable.add_detection_point("test") assert not target_observable.has_detection_points() target_observable.apply_diff_merge(original_observable, modified_observable) assert not target_observable.has_detection_points()
def from_dict(value: dict, system: "ace.system.ACESystem") -> "AnalysisRequest": assert isinstance(value, dict) data = AnalysisRequestModel(**value) root = None if isinstance(data.root, RootAnalysisModel): root = RootAnalysis.from_dict(data.root.dict(), system=system) observable = None if data.observable: observable = Observable.from_dict(data.observable.dict(), root) observable = root.get_observable(observable) type = None if data.type: type = AnalysisModuleType.from_dict(data.type.dict()) ar = AnalysisRequest(system, root, observable, type) ar.id = data.id # ar.dependency_analysis = json_data["dependency_analysis"] ar.status = data.status ar.owner = data.owner if data.original_root: ar.original_root = RootAnalysis.from_dict( data.original_root.dict(), system) if data.modified_root: ar.modified_root = RootAnalysis.from_dict( data.modified_root.dict(), system) return ar
def test_apply_diff_merge_tags(): original_root = RootAnalysis() original_observable = original_root.add_observable("test", "test") modified_root = original_root.copy() modified_observable = modified_root.get_observable(original_observable) modified_observable.add_tag("test") target_root = original_root.copy() target_observable = target_root.add_observable("test", "test") assert not target_observable.tags target_observable.apply_diff_merge(original_observable, modified_observable) assert target_observable.tags assert target_observable.tags[0] == "test" # exists before but not after original_root = RootAnalysis() original_observable = original_root.add_observable("test", "test") modified_root = original_root.copy() modified_observable = modified_root.get_observable(original_observable) target_root = original_root.copy() target_observable = target_root.add_observable("test", "test") original_observable.add_tag("test") assert not target_observable.tags target_observable.apply_diff_merge(original_observable, modified_observable) assert not target_observable.tags
def test_root_all_detection_points(): amt = AnalysisModuleType("test", "") root = RootAnalysis() observable = root.add_observable("test", "test") analysis = observable.add_analysis(type=amt) observable_2 = analysis.add_observable("test", "test2") root.add_detection_point("test") assert root.all_detection_points == [DetectionPoint("test")] observable.add_detection_point("test") assert root.all_detection_points == [ DetectionPoint("test"), DetectionPoint("test") ] analysis.add_detection_point("test") assert root.all_detection_points == [ DetectionPoint("test"), DetectionPoint("test"), DetectionPoint("test") ] observable_2.add_detection_point("test") assert root.all_detection_points == [ DetectionPoint("test"), DetectionPoint("test"), DetectionPoint("test"), DetectionPoint("test"), ]
def test_observable_properties(): amt = AnalysisModuleType("test", "") root = RootAnalysis() observable = root.add_observable("test", "test") analysis = observable.add_analysis(type=amt) assert observable.all_analysis == [analysis] assert observable.children == [analysis]
def test_root_all(): amt = AnalysisModuleType("test", "") root = RootAnalysis() observable = root.add_observable("test", "test") analysis = observable.add_analysis(type=amt) observable_2 = analysis.add_observable("test", "test2") assert sorted(root.all) == sorted( [root, analysis, observable, observable_2])
def test_add_duplicate_analysis(): amt = AnalysisModuleType("test", "") root = RootAnalysis() observable = root.add_observable("test", "test") analysis = observable.add_analysis(type=amt) # adding the same analysis twice will fail with pytest.raises(ValueError): observable.add_analysis(type=amt)
def test_search_down(): amt = AnalysisModuleType("test", "") root = RootAnalysis() observable = root.add_observable("test", "test") analysis = observable.add_analysis(type=amt) child = analysis.add_observable("test", "child") observable_2 = root.add_observable("test", "other") assert search_down(child, lambda obj: obj == analysis) == analysis assert search_down(child, lambda obj: obj == observable) == observable assert search_down(child, lambda obj: obj == root) == root assert search_down(child, lambda obj: False) is None
def test_root_copy(): root = RootAnalysis() observable = root.add_observable("test", "test") amt = AnalysisModuleType("test", "") analysis = observable.add_analysis(type=amt, details={"test": "test"}) root_copy = root.copy() observable_copy = root_copy.get_observable(observable) assert observable_copy == observable assert not (observable_copy is observable) analysis_copy = observable_copy.get_analysis(amt) assert analysis_copy == analysis assert not (analysis_copy is analysis)
def test_relationships(): root = RootAnalysis() observable_1 = root.add_observable("test", "test1") observable_2 = root.add_observable("test", "test2") observable_1.add_relationship("test", observable_2) assert observable_1.has_relationship("test") assert observable_1.has_relationship("test", observable_2) assert observable_1.get_relationships_by_type("test") == [observable_2] assert observable_1.get_relationship_by_type("test") == observable_2 assert observable_1.get_relationships_by_type("unknown") == [] assert observable_1.get_relationship_by_type("unknown") is None
def test_analysis_properties(): amt = AnalysisModuleType("test", "") root = RootAnalysis() observable = root.add_observable("test", "test") analysis = observable.add_analysis(type=amt, details={"test": "test"}) observable_2 = analysis.add_observable("test2", "test2") assert analysis.uuid assert analysis.root == root assert analysis.type == amt assert analysis.observable == observable assert analysis.observables == [observable_2] assert analysis._details == {"test": "test"} assert analysis.children == [observable_2] assert analysis.observable_types == ["test2"]
def test_add_excluded_analysis(): observable = RootAnalysis().add_observable("test", "test") assert not observable.excluded_analysis assert not observable.is_excluded("test") observable.exclude_analysis("test") assert observable.excluded_analysis assert "test" in observable.excluded_analysis assert observable.is_excluded("test") assert observable.is_excluded(AnalysisModuleType("test", "")) observable.exclude_analysis(AnalysisModuleType("other", "")) assert "other" in observable.excluded_analysis
def test_root_get_observable(): amt = AnalysisModuleType("test", "") root = RootAnalysis() observable = root.add_observable("test", "test") # get by uuid assert root.get_observable(observable.uuid) == observable # get by identity assert root.get_observable(observable) == observable # get by new object assert root.get_observable(RootAnalysis().add_observable( "test", "test")) == observable # get invalid object assert root.get_observable("") is None assert root.get_observable(RootAnalysis().add_observable("test", "blah")) is None
async def test_EVENT_CACHE_HIT(system): handler = TestEventHandler() await system.register_event_handler(EVENT_CACHE_HIT, handler) amt = AnalysisModuleType("test", "", cache_ttl=60) await system.register_analysis_module_type(amt) root = system.new_root() observable = root.add_observable("test", "test") root_request = root.create_analysis_request() await system.process_analysis_request(root_request) request = await system.get_next_analysis_request("owner", amt, 0) request.initialize_result() request.modified_observable.add_analysis(type=amt, details={"test": "test"}) await system.process_analysis_request(request) assert handler.event is None root = system.new_root() observable = root.add_observable("test", "test") root_request = root.create_analysis_request() await system.process_analysis_request(root_request) await handler.wait() assert handler.event.name == EVENT_CACHE_HIT event_root = RootAnalysis.from_dict(handler.event.args[0], system) assert event_root.uuid == root.uuid and event_root.version is not None assert handler.event.args[1]["type"] == observable.type assert handler.event.args[1]["value"] == observable.value assert isinstance(AnalysisRequest.from_dict(handler.event.args[2], system), AnalysisRequest)
def test_root_analysis_serialization(): root = RootAnalysis( tool="test", tool_instance="test", alert_type="test", desc="test", event_time=datetime.datetime.now(), name="test", analysis_mode="test", queue="test", instructions="test", ) amt = AnalysisModuleType("test", "") observable = root.add_observable("test", "test") analysis = observable.add_analysis(type=amt, details={"test": "test"}) root.add_detection_point("test") new_root = RootAnalysis.from_dict(root.to_dict()) assert root == new_root assert root.tool == new_root.tool assert root.tool_instance == new_root.tool assert root.alert_type == new_root.alert_type assert root.description == new_root.description assert root.event_time == new_root.event_time assert root.name == new_root.name assert root.analysis_mode == new_root.analysis_mode assert root.queue == new_root.queue assert root.instructions == new_root.instructions assert root.detections == new_root.detections # the observable property for the root should always be None assert root.observable is None assert len(root.observables) == 1 new_root = RootAnalysis.from_json(root.to_json()) assert root == new_root assert root.tool == new_root.tool assert root.tool_instance == new_root.tool assert root.alert_type == new_root.alert_type assert root.description == new_root.description assert root.event_time == new_root.event_time assert root.name == new_root.name assert root.analysis_mode == new_root.analysis_mode assert root.queue == new_root.queue assert root.instructions == new_root.instructions # the observable property for the root should always be None assert root.observable is None assert len(root.observables) == 1
def test_observable_serialization(): root = RootAnalysis() o_time = utc_now() target = root.add_observable("test", "other") o1 = root.add_observable( "test", "test", time=o_time, context="text context", directives=["directive1", "directive2"], limited_analysis=["limit1", "limit2"], excluded_analysis=["excluded1", "excluded2"], requested_analysis=["requested1", "requested2"], ) o1.add_relationship("test", target) root = RootAnalysis.from_dict(root.to_model().dict()) o2 = root.get_observable(o1) # should be two separate instances assert id(o1) != id(o2) assert o1.type == o2.type assert o1.value == o2.value assert o1.time == o2.time assert o1.context == o2.context assert o1.directives == o2.directives assert o1.limited_analysis == o2.limited_analysis assert o1.excluded_analysis == o2.excluded_analysis assert o1.requested_analysis == o2.requested_analysis assert o1.relationships == o2.relationships
def test_root_analysis_get_observables_by_type(): amt = AnalysisModuleType("test", "") root = RootAnalysis() observable_1 = root.add_observable("test", "test_1") observable_2 = root.add_observable("test", "test_2") observable_3 = root.add_observable("test_3", "test_3") analysis = observable_3.add_analysis(type=amt) observable_4 = analysis.add_observable("test_4", "test_4") assert root.get_observables_by_type("test") == [observable_1, observable_2] assert root.get_observables_by_type("test_3") == [observable_3] assert root.get_observables_by_type("test_4") == [observable_4] assert root.get_observables_by_type("unknown") == [] assert root.get_observable_by_type("test") in [observable_1, observable_2] assert root.get_observable_by_type("test_3") == observable_3 assert root.get_observable_by_type("test_4") == observable_4 assert root.get_observable_by_type("unknown") is None
def test_has_observable(): root = RootAnalysis() observable = root.add_observable("test", "test") assert root.has_observable("test", "test") assert not root.has_observable("test", "t3st") assert root.has_observable(Observable("test", "test")) assert not root.has_observable(Observable("t3st", "test")) assert not root.has_observable(Observable("test", "t3st"))
async def i_update_root_analysis(self, root: RootAnalysis) -> bool: # when we update we also update the version new_version = str(uuid.uuid4()) async with self.get_db() as db: result = await db.execute( update(RootAnalysisTracking).values( version=new_version, json_data=root.to_json(exclude_analysis_details=True) ) # so the version has to match for the update to work .where(and_(RootAnalysisTracking.uuid == root.uuid, RootAnalysisTracking.version == root.version)) ) await db.commit() if result.rowcount == 0: # if the version doesn't match then the update fails return False root.version = new_version return True
async def i_track_root_analysis(self, root: RootAnalysis) -> bool: """Tracks the given root to the given RootAnalysis uuid.""" version = root.version if version is None: version = str(uuid.uuid4()) try: async with self.get_db() as db: await db.execute( insert(RootAnalysisTracking).values( uuid=root.uuid, version=version, json_data=root.to_json(exclude_analysis_details=True) ) ) await db.commit() root.version = version return True except sqlalchemy.exc.IntegrityError: return False
def test_add_link(): root = RootAnalysis() observable_1 = root.add_observable("test", "test1") observable_2 = root.add_observable("test", "test2") observable_3 = root.add_observable("test", "test3") observable_2.add_link(observable_1) assert observable_2.links == [observable_1] observable_2.add_link(observable_3) assert sorted(observable_2.links) == sorted([observable_1, observable_3]) observable_2.add_tag("test") assert observable_2.has_tag("test") assert observable_1.has_tag("test") assert observable_3.has_tag("test") with pytest.raises(ValueError): # cannot do this since 2 already points to 1 observable_1.add_link(observable_2)
def test_analysis_find_observables(): amt = AnalysisModuleType("test", "") root = RootAnalysis() observable = root.add_observable("test", "no find") analysis = observable.add_analysis(type=amt) o1 = analysis.add_observable("test", "test_1") o2 = analysis.add_observable("test", "test_2") o_all = sorted([o1, o2]) # search by type, single observable assert analysis.find_observable("test") in o_all # search by type, multi observable assert sorted(analysis.find_observables("test")) == o_all # search by lambda, single observable assert analysis.find_observable(lambda o: o.type == "test") in o_all # search by lambda, multi observable assert sorted( analysis.find_observables(lambda o: o.type == "test")) == o_all
def test_root_eq(): # two different uuids assert RootAnalysis() != RootAnalysis() root = RootAnalysis() # same uuids assert root == root.copy() # invalid compare assert root != object() # same uuid different version root = RootAnalysis() modified_root = root.copy() modified_root.version = str(uuid.uuid4()) assert root != modified_root # same uuid same version root.version = modified_root.version assert root == modified_root
async def get_root_analysis(self, uuid: str) -> Union[RootAnalysis, None]: assert isinstance(uuid, str) async with self.get_client() as client: response = await client.get(f"/analysis_tracking/root/{uuid}") if response.status_code == 404: return None _raise_exception_on_error(response) return RootAnalysis.from_dict(response.json())
async def test_EVENT_ANALYSIS_ROOT_EXPIRED(system): handler = TestEventHandler() await system.register_event_handler(EVENT_ANALYSIS_ROOT_EXPIRED, handler) root = system.new_root(expires=True) await root.submit() await handler.wait() assert handler.event.name == EVENT_ANALYSIS_ROOT_EXPIRED event_root = RootAnalysis.from_dict(handler.event.args, system) assert event_root.uuid == root.uuid and event_root.version is not None
def test_add_limited_analysis(): observable = RootAnalysis().add_observable("test", "test") assert not observable.limited_analysis observable.limit_analysis("test") assert observable.limited_analysis assert "test" in observable.limited_analysis observable.limit_analysis(AnalysisModuleType("other", "")) assert "other" in observable.limited_analysis
async def test_analysis_completed(system): await system.register_analysis_module_type( amt := AnalysisModuleType("test", "test", ["test"])) root = system.new_root() observable = root.add_observable("test", "test") assert not root.analysis_completed(observable, amt) observable.add_analysis(Analysis(type=amt, details=TEST_DETAILS)) assert root.analysis_completed(observable, amt) # unknown observable with pytest.raises(UnknownObservableError): root.analysis_completed(RootAnalysis().add_observable("test", "blah"), amt)
async def i_get_root_analysis(self, uuid: str) -> Union[RootAnalysisTracking, None]: """Returns the root for the given uuid or None if it does not exist..""" async with self.get_db() as db: result = ( await db.execute(select(RootAnalysisTracking).where(RootAnalysisTracking.uuid == uuid)) ).one_or_none() if not result: return None # we keep a copy of the actual value of the version property in the database # (the JSON would have a copy of the previous value) # also see update_root_analysis root = RootAnalysis.from_json(result[0].json_data) root.version = result[0].version return root
async def test_EVENT_ANALYSIS_ROOT_MODIFIED(system): handler = TestEventHandler() await system.register_event_handler(EVENT_ANALYSIS_ROOT_MODIFIED, handler) root = system.new_root() await system.track_root_analysis(root) assert handler.event is None handler = TestEventHandler() await system.register_event_handler(EVENT_ANALYSIS_ROOT_MODIFIED, handler) # already tracked so should fire as modified await system.track_root_analysis(root) await handler.wait() assert handler.event.name == EVENT_ANALYSIS_ROOT_MODIFIED assert RootAnalysis.from_dict(handler.event.args, system) == root
async def test_analysis_tracked(system): await system.register_analysis_module_type( amt := AnalysisModuleType("test", "test", ["test"])) root = system.new_root() observable = root.add_observable("test", "test") assert not root.analysis_tracked(observable, amt) ar = observable.create_analysis_request(amt) observable.track_analysis_request(ar) assert root.analysis_tracked(observable, amt) # invalid observable with pytest.raises(UnknownObservableError): root.analysis_tracked(RootAnalysis().add_observable("test", "blah"), amt)