Restore (mitre#676)

* adding persistent database

* adding persistent database

.gitignore

@@ -17,3 +17,4 @@ venv/
 .env
 conf/*.yml
 !conf/default.yml
+data/object_store

app/service/data_svc.py

@@ -1,5 +1,8 @@
+import asyncio
 import glob
 import json
+import pickle
+
 from base64 import b64encode
 from collections import defaultdict
 
@@ -20,6 +23,28 @@ def __init__(self, dao):
         self.log = self.add_service('data_svc', self)
         self.ram = dict(agents=[], planners=[], adversaries=[], abilities=[])
 
+    async def save_state(self):
+        """
+        Save RAM database to file
+        :return:
+        """
+        with open('data/object_store', 'wb') as objects:
+            pickle.dump(self.ram, objects)
+
+    async def restore_state(self):
+        """
+        Restore the object database - but wait for YML files to load first
+        :return:
+        """
+        await asyncio.sleep(3)
+        with open('data/object_store', 'rb') as objects:
+            ram = pickle.load(objects)
+            [await self.store(x) for x in ram['agents']]
+            [await self.store(x) for x in ram['planners']]
+            [await self.store(x) for x in ram['abilities']]
+            [await self.store(x) for x in ram['adversaries']]
+        self.log.debug('Restored objects from persistent storage')
+
     async def apply(self, collection):
         """
         Add a new collection to RAM
@@ -228,26 +253,30 @@ async def _create_adversary(self, i, name, description, phases):
         return identifier
 
     async def _load_facts(self, directory):
+        total = 0
         for filename in glob.iglob('%s/*.yml' % directory, recursive=False):
             for source in self.strip_yml(filename):
                 source_id = await self.dao.create('core_source', dict(name=source['name']))
                 for fact in source.get('facts', []):
                     fact['source_id'] = source_id
                     fact['score'] = fact.get('score', 1)
                     await self.save('fact', fact)
-
                 for rule in source.get('rules', []):
                     rule['source_id'] = source_id
                     await self._create_rule(**rule)
+                total += 1
+        self.log.debug('Loaded %s fact sources' % total)
 
     async def _load_planners(self, directory):
+        total = 0
         for filename in glob.iglob('%s/*.yml' % directory, recursive=False):
             for planner in self.strip_yml(filename):
                 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']))
+                total += 1
+        self.log.debug('Loaded %s planners' % total)
 
     async def _create_rule(self, fact, source_id, action='DENY', match='.*'):
         try:
@@ -305,6 +334,7 @@ async def _create_operation(self, name, group, adversary_id, jitter='2/8', sourc
         return op_id
 
     async def _load_adversaries(self, directory):
+        total = 0
         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]
@@ -320,9 +350,11 @@ async def _load_adversaries(self, directory):
                         Adversary(adversary_id=adv['id'], name=adv['name'], description=adv['description'],
                                   phases=phases)
                     )
-        self.log.debug('Loaded %s adversaries' % len(self.ram['adversaries']))
+                    total += 1
+        self.log.debug('Loaded %s adversaries' % total)
 
     async def _load_abilities(self, directory):
+        total = 0
         for filename in glob.iglob('%s/**/*.yml' % directory, recursive=True):
             for entries in self.strip_yml(filename):
                 for ab in entries:
@@ -342,7 +374,8 @@ async def _load_abilities(self, directory):
                                                                'cleanup') else None,
                                                            payload=info.get('payload'), parsers=info.get('parsers', []),
                                                            requirements=ab.get('requirements', []))
-        self.log.debug('Loaded %s abilities' % len(self.ram['abilities']))
+                                total += 1
+        self.log.debug('Loaded %s abilities' % total)
 
     async def _create_ability(self, ability_id, tactic, technique_name, technique_id, name, test, description, executor,
                               platform, cleanup=None, payload=None, parsers=None, requirements=None):

server.py

@@ -2,15 +2,15 @@
 import asyncio
 import logging
 import os
-import sys
 import pathlib
+import sys
 from importlib import import_module
+from subprocess import Popen, DEVNULL
 
 import aiohttp_jinja2
 import jinja2
 import yaml
 from aiohttp import web
-from subprocess import Popen, DEVNULL
 
 from app.database.core_dao import CoreDao
 from app.service.agent_svc import AgentService
@@ -28,6 +28,7 @@ async def background_tasks(app):
     app.loop.create_task(operation_svc.resume())
     app.loop.create_task(data_svc.load_data(directory='data'))
     app.loop.create_task(agent_svc.start_sniffer_untrusted_agents())
+    app.loop.create_task(data_svc.restore_state())
 
 
 def build_plugins(plugs):
@@ -86,7 +87,8 @@ def main(services, host, port, users):
     try:
         loop.run_forever()
     except KeyboardInterrupt:
-        pass
+        loop.run_until_complete(data_svc.save_state())
+        logging.debug('[!] shutting down server...good-bye')
 
 
 if __name__ == '__main__':
@@ -102,7 +104,6 @@ def main(services, host, port, users):
         plugin_modules = build_plugins(cfg['plugins'])
         plugin_svc = PluginService(plugin_modules)
         data_svc = DataService(CoreDao('core.db', memory=cfg['memory']))
-        logging.debug('Using an in-memory database: %s' % cfg['memory'])
         planning_svc = PlanningService()
         parsing_svc = ParsingService()
         reporting_svc = ReportingService()