ensuring abilities can be updated from the gui (mitre#496)

app/service/data_svc.py

@@ -37,22 +37,7 @@ async def load_abilities(self, directory):
         :return: None
         """
         for filename in glob.iglob('%s/**/*.yml' % directory, recursive=True):
-            for entries in self.strip_yml(filename):
-                for ab in entries:
-                    for pl, executors in ab['platforms'].items():
-                        for name, info in executors.items():
-                            for e in name.split(','):
-                                encoded_test = b64encode(info['command'].strip().encode('utf-8'))
-                                await self.create_ability(ability_id=ab.get('id'), tactic=ab['tactic'].lower(),
-                                                          technique_name=ab['technique']['name'],
-                                                          technique_id=ab['technique']['attack_id'],
-                                                          test=encoded_test.decode(), description=ab.get('description'),
-                                                          executor=e, name=ab['name'], platform=pl,
-                                                          cleanup=b64encode(
-                                                              info['cleanup'].strip().encode(
-                                                                  'utf-8')).decode() if info.get(
-                                                              'cleanup') else None,
-                                                          payload=info.get('payload'), parser=info.get('parser'))
+            await self._save_ability_to_database(filename)
 
     async def load_adversaries(self, directory):
         """
@@ -106,6 +91,7 @@ async def persist_ability(self, ability_id, file_contents):
             f.seek(0)
             f.write(file_contents)
             f.truncate()
+        await self._save_ability_to_database(file_path)
 
     async def persist_adversary(self, i, name, description, phases):
         """
@@ -152,13 +138,17 @@ async def create_ability(self, ability_id, tactic, technique_name, technique_id,
                        technique_id=technique_id, technique_name=technique_name,
                        executor=executor, platform=platform, description=description,
                        cleanup=cleanup)
+        # update
+        unique_criteria = dict(ability_id=ability_id, platform=platform, executor=executor)
+        for entry in await self.dao.get('core_ability', unique_criteria):
+            await self.update('core_ability', 'id', entry['id'], ability)
+            await self.dao.delete('core_parser', dict(ability=entry['id']))
+            await self.dao.delete('core_payload', dict(ability=entry['id']))
+            return await self._save_ability_extras(entry['id'], payload, parser)
+
+        # new
         identifier = await self.dao.create('core_ability', ability)
-        if payload:
-            await self.dao.create('core_payload', dict(ability=identifier, payload=payload))
-        if parser:
-            parser['ability'] = identifier
-            await self.dao.create('core_parser', parser)
-        return identifier
+        return await self._save_ability_extras(identifier, payload, parser)
 
     async def create_adversary(self, i, name, description, phases):
         """
@@ -375,3 +365,40 @@ async def update(self, table, key, value, data):
         :return: None
         """
         await self.dao.update(table, key, value, data)
+
+    """ PRIVATE """
+
+    async def _save_ability_to_database(self, filename):
+        for entries in self.strip_yml(filename):
+            for ab in entries:
+                for pl, executors in ab['platforms'].items():
+                    for name, info in executors.items():
+                        for e in name.split(','):
+                            encoded_test = b64encode(info['command'].strip().encode('utf-8'))
+                            await self.create_ability(ability_id=ab.get('id'), tactic=ab['tactic'].lower(),
+                                                      technique_name=ab['technique']['name'],
+                                                      technique_id=ab['technique']['attack_id'],
+                                                      test=encoded_test.decode(), description=ab.get('description'),
+                                                      executor=e, name=ab['name'], platform=pl,
+                                                      cleanup=b64encode(
+                                                          info['cleanup'].strip().encode(
+                                                              'utf-8')).decode() if info.get(
+                                                          'cleanup') else None,
+                                                      payload=info.get('payload'), parser=info.get('parser'))
+                await self._delete_stale_abilities(ab)
+
+    async def _save_ability_extras(self, identifier, payload, parser):
+        if payload:
+            await self.dao.create('core_payload', dict(ability=identifier, payload=payload))
+        if parser:
+            parser['ability'] = identifier
+            await self.dao.create('core_parser', parser)
+        return identifier
+
+    async def _delete_stale_abilities(self, ability):
+        for saved in await self.dao.get('core_ability', dict(ability_id=ability.get('id'))):
+            for platform, executors in ability['platforms'].items():
+                if platform == saved['platform'] and not saved['executor'] in str(executors.keys()):
+                    await self.dao.delete('core_ability', dict(id=saved['id']))
+            if saved['platform'] not in ability['platforms']:
+                await self.dao.delete('core_ability', dict(id=saved['id']))

conf/core.sql

@@ -1,13 +1,13 @@
 CREATE TABLE if not exists core_result (link_id integer, output text, parsed data);
 CREATE TABLE if not exists core_ability (id integer primary key AUTOINCREMENT, ability_id text, tactic text, technique_name, technique_id text, name text, test text, description text, cleanup text, executor, platform, UNIQUE (ability_id, platform, executor) ON CONFLICT IGNORE);
-CREATE TABLE if not exists core_payload (id integer primary key AUTOINCREMENT, ability integer, payload text, UNIQUE (ability, payload) ON CONFLICT IGNORE);
+CREATE TABLE if not exists core_payload (ability integer, payload text, UNIQUE (ability, payload) ON CONFLICT IGNORE);
 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 (phase integer, adversary_id text, ability_id text, UNIQUE (adversary_id, phase, ability_id));
 CREATE TABLE if not exists core_agent (id integer primary key AUTOINCREMENT, paw text, last_seen date, architecture text, platform text, server text, host_group text, location text, pid integer, ppid integer, trusted integer, last_trusted_seen date, sleep integer);
 CREATE TABLE if not exists core_executor (id integer primary key AUTOINCREMENT, agent_id integer, executor text, preferred integer);
 CREATE TABLE if not exists core_operation (id integer primary key AUTOINCREMENT, name text, host_group text, adversary_id text, jitter text, start date, finish date, phase integer, stealth integer, planner integer, state text, allow_untrusted integer);
 CREATE TABLE if not exists core_chain (id integer primary key AUTOINCREMENT, op_id integer, paw text, ability integer, jitter integer, command text, executor text, cleanup integer, score integer, status integer, decide date, collect date, finish date, UNIQUE(op_id, paw, command));
-CREATE TABLE if not exists core_parser (id integer primary key AUTOINCREMENT, ability integer, name text, property text, script text, UNIQUE(ability, property) ON CONFLICT REPLACE);
+CREATE TABLE if not exists core_parser (ability integer, name text, property text, script text, UNIQUE(ability, property) ON CONFLICT REPLACE);
 CREATE TABLE if not exists core_fact (id integer primary key AUTOINCREMENT, property text, value text, score integer, set_id integer, source_id text, link_id integer);
 CREATE TABLE if not exists core_source (id integer primary key AUTOINCREMENT, name text, UNIQUE(name) ON CONFLICT IGNORE);
 CREATE TABLE if not exists core_source_map (id integer primary key AUTOINCREMENT, op_id integer, source_id integer, UNIQUE(op_id, source_id) ON CONFLICT IGNORE);