swapping adversary to c_object (mitre#665)

app/objects/c_adversary.py

@@ -0,0 +1,24 @@
+from app.objects.base_object import BaseObject
+
+
+class Adversary(BaseObject):
+
+    @property
+    def unique(self):
+        return self.hash('%s' % self.adversary_id)
+
+    @property
+    def display(self):
+        return dict(adversary_id=self.adversary_id, name=self.name, description=self.description, phases=self.phases)
+
+    def __init__(self, adversary_id, name, description, phases):
+        self.adversary_id = adversary_id
+        self.name = name
+        self.description = description
+        self.phases = phases
+
+    def store(self, ram):
+        existing = self.retrieve(ram['adversaries'], self.unique)
+        if not existing:
+            ram['adversaries'].append(self)
+            return self.retrieve(ram['adversaries'], self.unique)

app/service/agent_svc.py

@@ -118,7 +118,6 @@ async def perform_action(self, link: typing.Dict) -> int:
             else:
                 await asyncio.sleep(30)
                 operation = (await self.data_svc.get('operation', dict(id=op_id)))[0]
-        link.pop('adversary_map_id')
         return await self.data_svc.save('link', link)
 
     @staticmethod

app/service/data_svc.py

@@ -3,6 +3,7 @@
 from base64 import b64encode
 from collections import defaultdict
 
+from app.objects.c_adversary import Adversary
 from app.objects.c_planner import Planner
 from app.service.base_service import BaseService
 from app.utility.rule import RuleAction
@@ -13,7 +14,7 @@ class DataService(BaseService):
     def __init__(self, dao):
         self.dao = dao
         self.log = self.add_service('data_svc', self)
-        self.ram = dict(agents=[], planners=[])
+        self.ram = dict(agents=[], planners=[], adversaries=[])
 
     async def apply(self, collection):
         """
@@ -51,8 +52,6 @@ async def save(self, object_name, object_dict):
                 return await self._create_operation(**object_dict)
             elif object_name == 'link':
                 return await self._create_link(object_dict)
-            elif object_name == 'adversary':
-                return await self._create_adversary(**object_dict)
             elif object_name == 'ability':
                 return await self._create_ability(**object_dict)
             elif object_name == 'relationship':
@@ -123,8 +122,6 @@ async def explode(self, object_name, criteria=None):
                 return await self._explode_operation(criteria)
             elif object_name == 'chain':
                 return await self._explode_chain(criteria)
-            elif object_name == 'adversary':
-                return await self._explode_adversaries(criteria)
             elif object_name == 'ability':
                 return await self._explode_abilities(criteria)
             elif object_name == 'parser':
@@ -187,22 +184,11 @@ async def _explode_abilities(self, criteria=None):
                 r['enforcements'] = (await self.dao.get('core_requirement_map', dict(requirement_id=r['id'])))[0]
         return abilities
 
-    async def _explode_adversaries(self, criteria=None):
-        adversaries = await self.dao.get('core_adversary', criteria)
-        for adv in adversaries:
-            phases = defaultdict(list)
-            for t in await self.dao.get('core_adversary_map', dict(adversary_id=adv['adversary_id'])):
-                for ability in await self._explode_abilities(dict(ability_id=t['ability_id'])):
-                    ability['adversary_map_id'] = t['id']
-                    phases[t['phase']].append(ability)
-            adv['phases'] = dict(phases)
-        return adversaries
-
     async def _explode_operation(self, criteria=None):
         operations = await self.dao.get('core_operation', criteria)
         for op in operations:
             op['chain'] = sorted(await self._explode_chain(criteria=dict(op_id=op['id'])), key=lambda k: k['id'])
-            adversaries = await self._explode_adversaries(dict(id=op['adversary_id']))
+            adversaries = await self.locate('adversaries', match=dict(adversary_id=op['adversary_id']))
             op['adversary'] = adversaries[0]
             op['host_group'] = await self.locate('agents', match=dict(group=op['host_group']))
             sources = await self.dao.get('core_source_map', dict(op_id=op['id']))
@@ -289,15 +275,6 @@ async def _load_abilities(self, directory):
         for filename in glob.iglob('%s/**/*.yml' % directory, recursive=True):
             await self._write_ability(filename)
 
-    async def _load_adversaries(self, directory):
-        for filename in glob.iglob('%s/*.yml' % directory, recursive=True):
-            for adv in self.strip_yml(filename):
-                phases = [dict(phase=k, id=i) for k, v in adv.get('phases', dict()).items() for i in v]
-                for pack in [await self._add_adversary_packs(p) for p in adv.get('packs', [])]:
-                    phases += pack
-                if adv.get('visible', True):
-                    await self._create_adversary(adv['id'], adv['name'], adv['description'], phases)
-
     async def _load_facts(self, directory):
         for filename in glob.iglob('%s/*.yml' % directory, recursive=False):
             for source in self.strip_yml(filename):
@@ -404,3 +381,20 @@ async def _create_operation(self, name, group, adversary_id, jitter='2/8', sourc
         for s_id in [s for s in sources if s]:
             await self.dao.create('core_source_map', dict(op_id=op_id, source_id=s_id))
         return op_id
+
+    async def _load_adversaries(self, directory):
+        for filename in glob.iglob('%s/*.yml' % directory, recursive=True):
+            for adv in self.strip_yml(filename):
+                phases = [dict(phase=k, id=i) for k, v in adv.get('phases', dict()).items() for i in v]
+                for pack in [await self._add_adversary_packs(p) for p in adv.get('packs', [])]:
+                    phases += pack
+                if adv.get('visible', True):
+                    pp = defaultdict(list)
+                    for phase in phases:
+                        for ability in await self._explode_abilities(dict(ability_id=phase['id'])):
+                            pp[phase['phase']].append(ability)
+                    phases = dict(pp)
+                    await self.store(
+                        Adversary(adversary_id=adv['id'], name=adv['name'], description=adv['description'], phases=phases)
+                    )
+        self.log.debug('Loaded %s adversaries' % len(self.ram['adversaries']))

app/service/operation_svc.py

@@ -50,7 +50,7 @@ async def run(self, op_id):
         operation = await self.data_svc.explode('operation', dict(id=op_id))
         try:
             planner = await self._get_planning_module(operation[0])
-            for phase in operation[0]['adversary']['phases']:
+            for phase in operation[0]['adversary'].phases:
                 await planner.execute(phase)
                 await self._wait_for_phase_completion(operation[0])
                 await self.data_svc.update('operation', key='id', value=op_id, data=dict(phase=phase))

app/service/planning_svc.py

@@ -27,7 +27,7 @@ async def select_links(self, operation, agent, phase):
         if (not agent.trusted) and (not operation['allow_untrusted']):
             self.log.debug('Agent %s untrusted: no link created' % agent.paw)
             return []
-        phase_abilities = [i for p, v in operation['adversary']['phases'].items() if p <= phase for i in v]
+        phase_abilities = [i for p, v in operation['adversary'].phases.items() if p <= phase for i in v]
         phase_abilities = sorted(phase_abilities, key=lambda i: i['id'])
         link_status = await self._default_link_status(operation)
 
@@ -36,7 +36,7 @@ async def select_links(self, operation, agent, phase):
             links.append(
                 dict(op_id=operation['id'], paw=agent.paw, ability=a['id'], command=a['test'], score=0,
                      status=link_status, decide=datetime.now(), executor=a['executor'],
-                     jitter=self.jitter(operation['jitter']), adversary_map_id=a['adversary_map_id']))
+                     jitter=self.jitter(operation['jitter'])))
         ability_requirements = {ab['id']: ab.get('requirements', []) for ab in phase_abilities}
         links[:] = await self._trim_links(operation, links, agent, ability_requirements)
         return await self._sort_links(links)
@@ -68,7 +68,7 @@ async def _sort_links(links):
         """
         sort links by their score then by the order they are defined in an adversary profile
         """
-        return sorted(links, key=lambda k: (-k['score'], k['adversary_map_id']))
+        return sorted(links, key=lambda k: (-k['score']))
 
     async def _trim_links(self, operation, links, agent, ability_requirements=None):
         host_already_ran = [l['command'] for l in operation['chain'] if l['paw'] == agent.paw]

app/service/reporting_svc.py

@@ -22,7 +22,7 @@ async def generate_operation_report(self, op_id, agent_output=False):
         op = (await self.data_svc.explode('operation', dict(id=op_id)))[0]
         planner = await self.data_svc.locate('planners', match=dict(name=op['planner']))
         report = dict(name=op['name'], id=op['id'], host_group=op['host_group'], start=op['start'], facts=op['facts'],
-                      finish=op['finish'], planner=planner[0].name, adversary=op['adversary'], jitter=op['jitter'], steps=[])
+                      finish=op['finish'], planner=planner[0].name, adversary=op['adversary'].display, jitter=op['jitter'], steps=[])
         agents_steps = {a.paw: {'steps': []} for a in op['host_group']}
         for step in op['chain']:
             ability = (await self.data_svc.explode('ability', criteria=dict(id=step['ability'])))[0]
@@ -87,8 +87,8 @@ async def write_report(report):
 
     @staticmethod
     async def _get_all_possible_abilities_by_agent(hosts, adversary):
-        return {a.paw: {'all_abilities': [ab for p in adversary['phases']
-                                          for ab in adversary['phases'][p]]} for a in hosts}
+        return {a.paw: {'all_abilities': [ab for p in adversary.phases
+                                          for ab in adversary.phases[p]]} for a in hosts}
 
     async def _get_operation_data(self, op_id):
         operation = (await self.get_service('data_svc').explode('operation', criteria=dict(id=op_id)))[0]

conf/core.sql

@@ -5,9 +5,6 @@ CREATE TABLE if not exists core_parser (id integer primary key AUTOINCREMENT, ab
 CREATE TABLE if not exists core_parser_map (parser_id integer, source text, edge text, target text, UNIQUE(parser_id, source, edge, target) ON CONFLICT IGNORE);
 CREATE TABLE if not exists core_requirement (id integer primary key AUTOINCREMENT, ability integer, module text, UNIQUE(ability, module) ON CONFLICT REPLACE);
 CREATE TABLE if not exists core_requirement_map (requirement_id integer, source text, edge text, target text, UNIQUE(requirement_id, source, edge, target) ON CONFLICT IGNORE);
-/** adversaries **/
-CREATE TABLE if not exists core_adversary (id integer primary key AUTOINCREMENT, adversary_id text, name text, description text, UNIQUE (adversary_id));
-CREATE TABLE if not exists core_adversary_map (id integer primary key AUTOINCREMENT, phase integer, adversary_id text, ability_id text, UNIQUE (adversary_id, phase, ability_id));
 /** facts **/
 CREATE TABLE if not exists core_fact (id integer primary key AUTOINCREMENT, property text, value text, score integer, source_id text, link_id integer DEFAULT 0);
 CREATE TABLE if not exists core_source (id integer primary key AUTOINCREMENT, name text, UNIQUE(name) ON CONFLICT IGNORE);