Merge pull request #1780 from mitre/virts-1220

VIRTS-1220/Objectives UI and API Functionality

.gitignore

@@ -27,11 +27,13 @@ data/facts/*yml
 !data/facts/.gitkeep
 data/sources/*
 !data/sources/.gitkeep
+data/objectives/*
+!data/objectives/.gitkeep
 .tox/
 
 # coverage reports
 htmlcov/
 .coverage
 .coverage.*
 *,cover
-_*/
+_*/

app/api/packs/advanced.py

@@ -15,6 +15,7 @@ def __init__(self, services):
 
     async def enable(self):
         self.app_svc.application.router.add_route('GET', '/advanced/sources', self._section_sources)
+        self.app_svc.application.router.add_route('GET', '/advanced/objectives', self._section_objectives)
         self.app_svc.application.router.add_route('GET', '/advanced/planners', self._section_planners)
         self.app_svc.application.router.add_route('GET', '/advanced/contacts', self._section_contacts)
         self.app_svc.application.router.add_route('GET', '/advanced/obfuscators', self._section_obfuscators)
@@ -50,3 +51,9 @@ async def _section_configurations(self, request):
     async def _section_sources(self, request):
         access = await self.auth_svc.get_permissions(request)
         return dict(sources=[s.display for s in await self.data_svc.locate('sources', match=dict(access=tuple(access)))])
+
+    @check_authorization
+    @template('objectives.html')
+    async def _section_objectives(self, request):
+        access = await self.auth_svc.get_permissions(request)
+        return dict(objectives=[o.display for o in await self.data_svc.locate('objectives', match=dict(access=tuple(access)))])

app/api/rest_api.py

@@ -82,6 +82,7 @@ async def rest_core(self, request):
                     adversaries=lambda d: self.rest_svc.persist_adversary(access, d),
                     abilities=lambda d: self.rest_svc.persist_ability(access, d),
                     sources=lambda d: self.rest_svc.persist_source(access, d),
+                    objectives=lambda d: self.rest_svc.persist_objective(access, d),
                     planners=lambda d: self.rest_svc.update_planner(d),
                     agents=lambda d: self.rest_svc.update_agent_data(d),
                     chain=lambda d: self.rest_svc.update_chain_data(d),

app/service/data_svc.py

@@ -4,6 +4,7 @@
 import os.path
 import pickle
 import shutil
+import warnings
 from base64 import b64encode
 
 from app.objects.c_ability import Ability
@@ -139,16 +140,22 @@ async def load_ability_file(self, filename, access):
                                 await self._update_extensions(a)
 
     async def load_adversary_file(self, filename, access):
-        for adv in self.strip_yml(filename):
-            adversary = Adversary.load(adv)
-            adversary.access = access
-            await self.store(adversary)
+        warnings.warn("Function deprecated and will be removed in a future update. Use load_yaml_file", DeprecationWarning)
+        await self.load_yaml_file(Adversary, filename, access)
 
     async def load_source_file(self, filename, access):
+        warnings.warn("Function deprecated and will be removed in a future update. Use load_yaml_file", DeprecationWarning)
+        await self.load_yaml_file(Source, filename, access)
+
+    async def load_objective_file(self, filename, access):
+        warnings.warn("Function deprecated and will be removed in a future update. Use load_yaml_file", DeprecationWarning)
+        await self.load_yaml_file(Objective, filename, access)
+
+    async def load_yaml_file(self, object_class, filename, access):
         for src in self.strip_yml(filename):
-            source = Source.load(src)
-            source.access = access
-            await self.store(source)
+            obj = object_class.load(src)
+            obj.access = access
+            await self.store(obj)
 
     """ PRIVATE """
 
@@ -173,7 +180,7 @@ async def _load(self, plugins=()):
 
     async def _load_adversaries(self, plugin):
         for filename in glob.iglob('%s/adversaries/**/*.yml' % plugin.data_dir, recursive=True):
-            await self.load_adversary_file(filename, plugin.access)
+            await self.load_yaml_file(Adversary, filename, plugin.access)
 
     async def _load_abilities(self, plugin):
         for filename in glob.iglob('%s/abilities/**/*.yml' % plugin.data_dir, recursive=True):
@@ -194,14 +201,11 @@ async def _classify(ability, tactic):
 
     async def _load_sources(self, plugin):
         for filename in glob.iglob('%s/sources/*.yml' % plugin.data_dir, recursive=False):
-            await self.load_source_file(filename, plugin.access)
+            await self.load_yaml_file(Source, filename, plugin.access)
 
     async def _load_objectives(self, plugin):
         for filename in glob.iglob('%s/objectives/*.yml' % plugin.data_dir, recursive=False):
-            for src in self.strip_yml(filename):
-                objective = Objective.load(src)
-                objective.access = plugin.access
-                await self.store(objective)
+            await self.load_yaml_file(Objective, filename, plugin.access)
 
     async def _load_payloads(self, plugin):
         for filename in glob.iglob('%s/payloads/*.yml' % plugin.data_dir, recursive=False):
@@ -216,10 +220,7 @@ async def _load_payloads(self, plugin):
 
     async def _load_planners(self, plugin):
         for filename in glob.iglob('%s/planners/*.yml' % plugin.data_dir, recursive=False):
-            for planner in self.strip_yml(filename):
-                planner = Planner.load(planner)
-                planner.access = plugin.access
-                await self.store(planner)
+            await self.load_yaml_file(Planner, filename, plugin.access)
 
     async def _load_extensions(self):
         for entry in self._app_configuration['payloads']['extensions']:

app/service/file_svc.py

@@ -47,8 +47,8 @@ async def get_file(self, headers):
             display_name = file_path.replace('.xored', '')
         return file_path, contents, display_name
 
-    async def save_file(self, filename, payload, target_dir):
-        self._save(os.path.join(target_dir, filename), payload)
+    async def save_file(self, filename, payload, target_dir, encrypt=True):
+        self._save(os.path.join(target_dir, filename), payload, encrypt)
 
     async def create_exfil_sub_directory(self, dir_name):
         path = os.path.join(self.get_config('exfil_dir'), dir_name)
@@ -143,8 +143,8 @@ def get_payload_name_from_uuid(self, payload):
 
     """ PRIVATE """
 
-    def _save(self, filename, content):
-        if self.encryptor and self.encrypt_output:
+    def _save(self, filename, content, encrypt=True):
+        if encrypt and (self.encryptor and self.encrypt_output):
             content = bytes(FILE_ENCRYPTION_FLAG, 'utf-8') + self.encryptor.encrypt(content)
         with open(filename, 'wb') as f:
             f.write(content)

app/service/rest_svc.py

@@ -10,7 +10,9 @@
 from aiohttp import web
 
 from app.objects.c_adversary import Adversary
+from app.objects.c_objective import Objective
 from app.objects.c_operation import Operation
+from app.objects.c_source import Source
 from app.objects.c_schedule import Schedule
 from app.objects.secondclass.c_fact import Fact
 from app.service.interfaces.i_rest_svc import RestServiceInterface
@@ -73,6 +75,19 @@ async def persist_source(self, access, data):
             r.extend(await self._persist_source(access, source))
         return r
 
+    async def persist_objective(self, access, data):
+        """Persist objectives. Accepts single objective or a bulk set of objectives.
+        For bulk, supply dict of form {"bulk": [{objective}, ...]}.
+        """
+        if data.get('bulk', False):
+            data = data['bulk']
+        else:
+            data = [data]
+        r = []
+        for obj in data:
+            r.extend(await self._persist_objective(access, obj))
+        return r
+
     async def delete_agent(self, data):
         await self.get_service('data_svc').remove('agents', data)
         return 'Delete action completed'
@@ -358,17 +373,12 @@ async def _persist_adversary(self, access, adv):
             # new
             file_path = 'data/adversaries/%s.yml' % adv['id']
             allowed = self._get_allowed_from_access(access)
-            adv['objective'] = adv.get('objective',
-                                       (await self._services.get('data_svc').locate('objectives', match=dict(name='default')))[0])
+            adv['objective'] = adv.get('objective', obj_default)
             final = adv
         # verfiy objective is valid
         if len(await self.get_service('data_svc').locate('objectives', match=dict(id=final['objective']))) == 0:
             final['objective'] = obj_default.id
-        with open(file_path, 'w+') as f:
-            f.seek(0)
-            f.write(yaml.dump(final))
-            f.truncate()
-        await self._services.get('data_svc').load_adversary_file(file_path, allowed)
+        await self._save_and_refresh_item(file_path, Adversary, final, allowed)
         return [a.display for a in await self._services.get('data_svc').locate('adversaries', dict(adversary_id=final["id"]))]
 
     async def _persist_ability(self, access, ab):
@@ -412,35 +422,38 @@ async def _persist_ability(self, access, ab):
             file_path = '%s/%s.yml' % (d, new_ability['id'])
             allowed = self._get_allowed_from_access(access)
             final = new_ability
-        with open(file_path, 'w+') as f:
-            f.seek(0)
-            f.write(yaml.dump([final]))
+        await self.get_service('file_svc').save_file(file_path, yaml.dump([final], encoding='utf-8'), '', encrypt=False)
         await self.get_service('data_svc').remove('abilities', dict(ability_id=final['id']))
         await self.get_service('data_svc').load_ability_file(file_path, allowed)
         await self._restore_exec_timeouts(final['id'], new_ability_exec_timeouts)
         return [a.display for a in
                 await self.get_service('data_svc').locate('abilities', dict(ability_id=final['id']))]
 
     async def _persist_source(self, access, source):
-        if not source.get('id') or not source['id']:
-            source['id'] = str(uuid.uuid4())
-        _, file_path = await self.get_service('file_svc').find_file_path('%s.yml' % source['id'], location='data')
+        return await self._persist_item(access, 'sources', Source, source)
+
+    async def _persist_objective(self, access, objective):
+        return await self._persist_item(access, 'objectives', Objective, objective)
+
+    async def _persist_item(self, access, object_class_name, object_class, item):
+        if not item.get('id') or not item['id']:
+            item['id'] = str(uuid.uuid4())
+        _, file_path = await self.get_service('file_svc').find_file_path('%s.yml' % item['id'], location='data')
         if file_path:
-            # exists
-            current_source = dict(self.strip_yml(file_path)[0])
-            allowed = (await self.get_service('data_svc').locate('sources', dict(id=source['id'])))[0].access
-            current_source.update(source)
-            final = source
+            current_item = dict(self.strip_yml(file_path)[0])
+            allowed = (await self.get_service('data_svc').locate(object_class_name, dict(id=item['id'])))[0].access
+            current_item.update(item)
+            final = item
         else:
-            # new
-            file_path = 'data/sources/%s.yml' % source['id']
+            file_path = 'data/%s/%s.yml' % (object_class_name, item['id'])
             allowed = self._get_allowed_from_access(access)
-            final = source
-        with open(file_path, 'w+') as f:
-            f.seek(0)
-            f.write(yaml.dump(final))
-        await self._services.get('data_svc').load_source_file(file_path, allowed)
-        return [s.display for s in await self._services.get('data_svc').locate('sources', dict(id=final['id']))]
+            final = item
+        await self._save_and_refresh_item(file_path, object_class, final, allowed)
+        return [i.display for i in await self.get_service('data_svc').locate(object_class_name, dict(id=final['id']))]
+
+    async def _save_and_refresh_item(self, file_path, object_class, final, allowed):
+        await self.get_service('file_svc').save_file(file_path, yaml.dump(final, encoding='utf-8'), '', encrypt=False)
+        await self.get_service('data_svc').load_yaml_file(object_class, file_path, allowed)
 
     async def _prep_new_ability(self, ab):
         """Take an ability dict, supplied by frontend, extract executor timeouts,

server.py

@@ -29,6 +29,7 @@ def setup_logger(level=logging.DEBUG):
             continue
         else:
             logging.getLogger(logger_name).setLevel(100)
+    logging.captureWarnings(True)
 
 
 async def start_server():

static/css/basic.css

@@ -156,6 +156,14 @@
     text-align: left;
     padding-bottom: 20px;
 }
+.pane-header-editable input,
+.fillable-table input {
+    background-color: inherit;
+    border: none;
+    color: var(--font-color);
+    text-align: left;
+    padding: 0;
+}
 .op-selected {
     visibility: hidden;
 }

templates/RED.html

@@ -38,6 +38,7 @@
             <p>Advanced</p>
         </div>
         <a onclick="viewSection('sources', '/advanced/sources')">sources</a>
+        <a onclick="viewSection('objectives', '/advanced/objectives')">objectives</a>
         <a onclick="viewSection('planners', '/advanced/planners')">planners</a>
         <a onclick="viewSection('contacts', '/advanced/contacts')">contacts</a>
         <a onclick="viewSection('obfuscators', '/advanced/obfuscators')">obfuscators</a>