| from app.objects.secondclass.c_fact import Fact, OriginType |
| from app.objects.secondclass.c_relationship import Relationship |
| from app.objects.secondclass.c_rule import Rule |
| from app.objects.secondclass.c_link import Link |
| |
| |
| class TestKnowledgeService: |
| |
| async def test_no_duplicate_fact(self, knowledge_svc, fire_event_mock): |
| await knowledge_svc.add_fact(Fact(trait='test', value='demo', score=1, |
| collected_by=['thin_air'], technique_id='T1234')) |
| await knowledge_svc.add_fact(Fact(trait='test', value='demo', score=1, |
| collected_by=['thin_air'], technique_id='T1234')) |
| facts = await knowledge_svc.get_facts(dict(trait='test')) |
| assert len(facts) == 1 |
| |
| async def test_no_duplicate_rules(self, knowledge_svc): |
| await knowledge_svc.add_rule(Rule(action='BLOCK', trait='a.c', match='.*')) |
| await knowledge_svc.add_rule(Rule(action='BLOCK', trait='a.c', match='.*')) |
| rules = await knowledge_svc.get_rules(dict(trait='a.c')) |
| assert len(rules) == 1 |
| |
| async def test_no_duplicate_relationship(self, knowledge_svc): |
| dummy = Fact(trait='test', value='demo', score=1, collected_by=['thin_air'], technique_id='T1234') |
| await knowledge_svc.add_relationship(Relationship(source=dummy, edge='potato', target=dummy)) |
| await knowledge_svc.add_relationship(Relationship(source=dummy, edge='potato', target=dummy)) |
| relationships = await knowledge_svc.get_relationships(dict(edge='potato')) |
| assert len(relationships) == 1 |
| |
| async def test_remove_fact(self, knowledge_svc, event_svc, fire_event_mock): |
| await knowledge_svc.add_fact(Fact(trait='rtest', value='rdemo', score=1, |
| collected_by=['thin_air'], technique_id='T1234'), |
| constraints=dict(test_field='test_value')) |
| await knowledge_svc.add_fact(Fact(trait='ktest', value='rdemo', score=1, |
| collected_by=['thin_air'], technique_id='T1234')) |
| await knowledge_svc.delete_fact(dict(trait='rtest')) |
| facts = await knowledge_svc.get_facts(dict(value='rdemo')) |
| assert len(facts) == 1 |
| assert len(knowledge_svc._KnowledgeService__loaded_knowledge_module.fact_ram['constraints']) == 0 |
| |
| async def test_remove_rules(self, knowledge_svc): |
| await knowledge_svc.add_rule(Rule(action='rBLOCK', trait='ra.c', match='.*'), |
| constraints=dict(test_field='test_value')) |
| await knowledge_svc.delete_rule(dict(trait='ra.c')) |
| rules = await knowledge_svc.get_rules(dict(trait='ra.c')) |
| assert len(rules) == 0 |
| assert len(knowledge_svc._KnowledgeService__loaded_knowledge_module.fact_ram['constraints']) == 0 |
| |
| async def test_remove_relationship(self, knowledge_svc): |
| dummy = Fact(trait='rtest', value='rdemo', score=1, collected_by=['thin_air'], technique_id='T1234') |
| await knowledge_svc.add_relationship(Relationship(source=dummy, edge='rpotato', target=dummy), |
| constraints=dict(test_field='test_value')) |
| await knowledge_svc.delete_relationship(dict(edge='rpotato')) |
| relationships = await knowledge_svc.get_relationships(dict(edge='rpotato')) |
| assert len(relationships) == 0 |
| assert len(knowledge_svc._KnowledgeService__loaded_knowledge_module.fact_ram['constraints']) == 0 |
| |
| async def test_update_fact(self, knowledge_svc, fire_event_mock): |
| await knowledge_svc.add_fact(Fact(trait='utest', value='udemo', score=1, |
| collected_by=['thin_air'], technique_id='T1234')) |
| await knowledge_svc.update_fact(criteria=dict(trait='utest'), |
| updates=dict(trait='utest2', value='udemo2')) |
| facts = await knowledge_svc.get_facts(dict(trait='utest2')) |
| assert len(facts) == 1 |
| assert facts[0].value == 'udemo2' |
| |
| async def test_update_relationship(self, knowledge_svc): |
| dummy = Fact(trait='utest', value='udemo', score=1, collected_by=['thin_air'], technique_id='T1234') |
| dummy2 = Fact(trait='utest2', value='udemo2', score=1, collected_by=['thin_air'], technique_id='T4321') |
| await knowledge_svc.add_relationship(Relationship(source=dummy, edge='upotato', target=dummy)) |
| await knowledge_svc.update_relationship(criteria=dict(edge='upotato'), |
| updates=dict(source=dummy2, edge='ubacon')) |
| relationships = await knowledge_svc.get_relationships(dict(edge='ubacon')) |
| assert len(relationships) == 1 |
| assert relationships[0].source == dummy2 |
| |
| async def test_retrieve_fact(self, knowledge_svc, fire_event_mock): |
| await knowledge_svc.add_fact(Fact(trait='ttestA', value='tdemoB', score=24, |
| collected_by=['thin_airA'], technique_id='T1234')) |
| await knowledge_svc.add_fact(Fact(trait='ttestB', value='tdemoA', score=42, |
| collected_by=['thin_airB'], technique_id='T4321')) |
| facts = await knowledge_svc.get_facts(dict(trait='ttestB')) |
| assert len(facts) == 1 |
| readable = facts[0].display |
| assert readable['value'] == 'tdemoA' |
| assert readable['score'] == 42 |
| |
| async def test_retrieve_relationship(self, knowledge_svc): |
| dummy = Fact(trait='ttest', value='tdemo', score=1, collected_by=['thin_air'], technique_id='T1234') |
| dummy2 = Fact(trait='ttest2', value='tdemo2', score=1, collected_by=['thin_air'], technique_id='T1234') |
| await knowledge_svc.add_relationship(Relationship(source=dummy, edge='tpotato', target=dummy2)) |
| await knowledge_svc.add_relationship(Relationship(source=dummy2, edge='tpotato', target=dummy)) |
| relationships = await knowledge_svc.get_relationships(dict(edge='tpotato')) |
| assert len(relationships) == 2 |
| specific = await knowledge_svc.get_relationships(dict(source=dummy)) |
| assert len(specific) == 1 |
| readable = specific[0].display |
| assert readable['edge'] == 'tpotato' |
| assert readable['target'].trait == 'ttest2' |
| |
| async def test_retrieve_rule(self, knowledge_svc): |
| await knowledge_svc.add_rule(Rule(action='tBLOCK', trait='ta.d', match='4.5.*')) |
| await knowledge_svc.add_rule(Rule(action='tALLOW', trait='ta.d', match='*.5.*')) |
| rules = await knowledge_svc.get_rules(dict(trait='ta.d')) |
| assert len(rules) == 2 |
| |
| fuzzy1 = await knowledge_svc.get_rules(dict(trait='ta.d', match='4.5.6')) |
| assert len(fuzzy1) == 2 |
| fuzzy2 = await knowledge_svc.get_rules(dict(trait='ta.d', match='6.5.4')) |
| assert len(fuzzy2) == 1 |
| assert fuzzy2[0].action == 'tALLOW' |
| fuzzy3 = await knowledge_svc.get_rules(dict(trait='ta.d', match='5.*')) |
| assert len(fuzzy3) == 2 |
| |
| async def test_fact_origin(self, knowledge_svc, ability, executor): |
| texecutor = executor(name='sh', platform='darwin', command='mkdir test', cleanup='rm -rf test') |
| tability = ability(ability_id='123', executors=[texecutor], repeatable=True, buckets=['test']) |
| link = Link.load(dict(command='', paw='n1234', ability=tability, executor=next(tability.executors), status=0, |
| id='ganymede')) |
| type1_fact = Fact(trait='t1', value='d1', score=1, collected_by=['thin_air'], technique_id='T1234', |
| links=[link.id], origin_type=OriginType.LEARNED) |
| type2_fact = Fact(trait='t2', value='d2', score=1, collected_by=['thin_air'], technique_id='T1234', |
| links=[link.id], origin_type=OriginType.LEARNED) |
| type3_fact = Fact(trait='t3', value='d3', score=1, collected_by=['tiny_lightning_bolts_running_through_sand'], |
| technique_id='T1234', origin_type=OriginType.SEEDED, source="Europa") |
| await knowledge_svc.add_fact(type1_fact) |
| await knowledge_svc.add_fact(type2_fact) |
| await knowledge_svc.add_fact(type3_fact) |
| origin_1, type_1 = await knowledge_svc.get_fact_origin(type1_fact) |
| origin_2, type_2 = await knowledge_svc.get_fact_origin(type2_fact.trait) |
| origin_3, type_3 = await knowledge_svc.get_fact_origin(type3_fact.trait) |
| assert origin_1 == link.id |
| assert origin_2 == link.id |
| assert origin_3 == 'Europa' |
| assert type_1 == OriginType.LEARNED |
| assert type_2 == OriginType.LEARNED |
| assert type_3 == OriginType.SEEDED |
