swapping planner to an object (mitre#642)

app/objects/c_planner.py

@@ -0,0 +1,24 @@
+from app.objects.base_object import BaseObject
+
+
+class Planner(BaseObject):
+
+    @property
+    def unique(self):
+        return self.name
+
+    @property
+    def display(self):
+        return dict(name=self.name, module=self.module, params=self.params)
+
+    def __init__(self, name, module, params):
+        self.name = name
+        self.module = module
+        self.params = params
+
+    def store(self, ram):
+        existing = self.retrieve(ram['planners'], self.unique)
+        if not existing:
+            ram['planners'].append(self)
+            return self.retrieve(ram['planners'], self.unique)
+

app/service/data_svc.py

@@ -3,6 +3,7 @@
 from base64 import b64encode
 from collections import defaultdict
 
+from app.objects.c_planner import Planner
 from app.service.base_service import BaseService
 from app.utility.rule import RuleAction
 
@@ -12,7 +13,7 @@ class DataService(BaseService):
     def __init__(self, dao):
         self.dao = dao
         self.log = self.add_service('data_svc', self)
-        self.ram = dict(agents=[])
+        self.ram = dict(agents=[], planners=[])
 
     async def load_data(self, directory=None, schema='conf/core.sql'):
         """
@@ -28,7 +29,7 @@ async def load_data(self, directory=None, schema='conf/core.sql'):
             await self._load_abilities(directory='%s/abilities' % directory)
             await self._load_adversaries(directory='%s/adversaries' % directory)
             await self._load_facts(directory='%s/facts' % directory)
-            await self._load_planner(directory='%s/planners' % directory)
+            await self._load_planners(directory='%s/planners' % directory)
 
     async def save(self, object_name, object_dict):
         """
@@ -124,8 +125,6 @@ async def explode(self, object_name, criteria=None):
                 return await self._explode_sources(criteria)
             elif object_name == 'result':
                 return await self._explode_results(criteria)
-            elif object_name == 'planner':
-                return await self._explode_planners(criteria)
             elif object_name == 'used':
                 return await self._explode_used(criteria)
             self.log.error('[!] EXPLODE on unknown type: %s' % object_name)
@@ -227,12 +226,6 @@ async def _explode_sources(self, criteria=None):
             s['facts'] = await self.dao.get('core_fact', dict(source_id=s['id']))
         return sources
 
-    async def _explode_planners(self, criteria=None):
-        planners = await self.dao.get('core_planner', criteria=criteria)
-        for p in planners:
-            p['params'] = json.loads(p['params'])
-        return planners
-
     async def _explode_parser(self, criteria=None):
         parsers = await self.dao.get('core_parser', criteria)
         for parser in parsers:
@@ -310,11 +303,14 @@ async def _load_facts(self, directory):
                     rule['source_id'] = source_id
                     await self._create_rule(**rule)
 
-    async def _load_planner(self, directory):
+    async def _load_planners(self, directory):
         for filename in glob.iglob('%s/*.yml' % directory, recursive=False):
             for planner in self.strip_yml(filename):
-                await self.dao.create('core_planner', dict(name=planner.get('name'), module=planner.get('module'),
-                                                           params=json.dumps(planner.get('params'))))
+                await self.store(
+                    Planner(name=planner.get('name'), module=planner.get('module'),
+                            params=json.dumps(planner.get('params')))
+                )
+        self.log.debug('Loaded %s planners' % len(self.ram['planners']))
 
     async def _create_rule(self, fact, source_id, action='DENY', match='.*'):
         try:

app/service/operation_svc.py

@@ -1,8 +1,9 @@
 import asyncio
+import ast
 import traceback
+
 from importlib import import_module
 
-from app.objects.c_agent import Agent
 from app.service.base_service import BaseService
 
 
@@ -61,10 +62,10 @@ async def run(self, op_id):
     """ PRIVATE """
 
     async def _get_planning_module(self, operation):
-        chosen_planner = await self.data_svc.explode('planner', dict(id=operation['planner']))
-        planning_module = import_module(chosen_planner[0]['module'])
-        return getattr(planning_module, 'LogicalPlanner')(operation, self.get_service('planning_svc'),
-                                                          **chosen_planner[0]['params'])
+        chosen_planner = await self.data_svc.locate('planners', match=dict(name=operation['planner']))
+        planning_module = import_module(chosen_planner[0].module)
+        planner_params = ast.literal_eval(chosen_planner[0].params)
+        return getattr(planning_module, 'LogicalPlanner')(operation, self.get_service('planning_svc'), **planner_params)
 
     async def _wait_for_phase_completion(self, operation):
         for member in operation['host_group']:

app/service/reporting_svc.py

@@ -20,9 +20,9 @@ async def generate_operation_report(self, op_id, agent_output=False):
         :return: a JSON report
         """
         op = (await self.data_svc.explode('operation', dict(id=op_id)))[0]
-        planner = (await self.data_svc.explode('planner', criteria=dict(id=op['planner'])))[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, adversary=op['adversary'], jitter=op['jitter'], steps=[])
+                      finish=op['finish'], planner=planner[0].name, adversary=op['adversary'], 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]
@@ -49,7 +49,7 @@ async def generate_operation_report(self, op_id, agent_output=False):
             agents_steps[step['paw']]['steps'].append(step_report)
         report['steps'] = agents_steps
         report['skipped_abilities'] = await self.get_skipped_abilities_by_agent(op_id=op['id'])
-        report.pop('host_group')
+        report['host_group'] = [a.display for a in report['host_group']]
         return report
 
     async def get_skipped_abilities_by_agent(self, op_id):

conf/core.sql

@@ -5,8 +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);
-/** planners **/
-CREATE TABLE if not exists core_planner (id integer primary key AUTOINCREMENT, name text, module text, params json, UNIQUE(name) 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));