Add plugin field to adversaries, abilities, and planners (#2345)

* Add plugin field to adversaries, abilities, sources, and planners

Author: argaudreau

app/objects/c_ability.py

@@ -1,5 +1,4 @@
 import collections
-import os
 import uuid
 
 import marshmallow as ma
@@ -8,7 +7,6 @@
 from app.objects.secondclass.c_executor import ExecutorSchema
 from app.objects.secondclass.c_requirement import RequirementSchema
 from app.utility.base_object import BaseObject
-from app.utility.base_service import BaseService
 from app.utility.base_world import AccessSchema
 
 
@@ -27,6 +25,7 @@ class AbilitySchema(ma.Schema):
     additional_info = ma.fields.Dict(keys=ma.fields.String(), values=ma.fields.String())
     access = ma.fields.Nested(AccessSchema, missing=None)
     singleton = ma.fields.Bool(missing=None)
+    plugin = ma.fields.String(missing=None)
 
     @ma.pre_load
     def fix_id(self, data, **_):
@@ -58,7 +57,7 @@ def executors(self):
 
     def __init__(self, ability_id='', name=None, description=None, tactic=None, technique_id=None, technique_name=None,
                  executors=(), requirements=None, privilege=None, repeatable=False, buckets=None, access=None,
-                 additional_info=None, tags=None, singleton=False, **kwargs):
+                 additional_info=None, tags=None, singleton=False, plugin='', **kwargs):
         super().__init__()
         self.ability_id = ability_id if ability_id else str(uuid.uuid4())
         self.tactic = tactic.lower() if tactic else None
@@ -80,6 +79,7 @@ def __init__(self, ability_id='', name=None, description=None, tactic=None, tech
         self.additional_info = additional_info or dict()
         self.additional_info.update(**kwargs)
         self.tags = set(tags) if tags else set()
+        self.plugin = plugin
 
     def __getattr__(self, item):
         try:
@@ -103,14 +103,11 @@ def store(self, ram):
         existing.update('buckets', self.buckets)
         existing.update('tags', self.tags)
         existing.update('singleton', self.singleton)
+        existing.update('plugin', self.plugin)
         return existing
 
     async def which_plugin(self):
-        file_svc = BaseService.get_service('file_svc')
-        for plugin in os.listdir('plugins'):
-            if await file_svc.walk_file_path(os.path.join('plugins', plugin, 'data', ''), '%s.yml' % self.ability_id):
-                return plugin
-        return None
+        return self.plugin
 
     def find_executor(self, name, platform):
         return self._executor_map.get(self._make_executor_map_key(name, platform))

app/objects/c_adversary.py

@@ -1,11 +1,9 @@
-import os
 import uuid
 
 import marshmallow as ma
 
 from app.objects.interfaces.i_object import FirstClassObjectInterface
 from app.utility.base_object import BaseObject
-from app.utility.base_service import BaseService
 
 
 DEFAULT_OBJECTIVE_ID = '495a9828-cab1-44dd-a0ca-66e58177d8cc'
@@ -20,6 +18,7 @@ class AdversarySchema(ma.Schema):
     objective = ma.fields.String()
     tags = ma.fields.List(ma.fields.String(), allow_none=True)
     has_repeatable_abilities = ma.fields.Boolean(dump_only=True)
+    plugin = ma.fields.String(missing=None)
 
     @ma.pre_load
     def fix_id(self, adversary, **_):
@@ -57,7 +56,7 @@ class Adversary(FirstClassObjectInterface, BaseObject):
     def unique(self):
         return self.hash('%s' % self.adversary_id)
 
-    def __init__(self, name='', adversary_id='', description='', atomic_ordering=(), objective='', tags=None):
+    def __init__(self, name='', adversary_id='', description='', atomic_ordering=(), objective='', tags=None, plugin=''):
         super().__init__()
         self.adversary_id = adversary_id if adversary_id else str(uuid.uuid4())
         self.name = name
@@ -66,6 +65,7 @@ def __init__(self, name='', adversary_id='', description='', atomic_ordering=(),
         self.objective = objective or DEFAULT_OBJECTIVE_ID
         self.tags = set(tags) if tags else set()
         self.has_repeatable_abilities = False
+        self.plugin = plugin
 
     def store(self, ram):
         existing = self.retrieve(ram['adversaries'], self.unique)
@@ -78,6 +78,7 @@ def store(self, ram):
         existing.update('objective', self.objective)
         existing.update('tags', self.tags)
         existing.update('has_repeatable_abilities', self.check_repeatable_abilities(ram['abilities']))
+        existing.update('plugin', self.plugin)
         return existing
 
     def verify(self, log, abilities, objectives):
@@ -101,11 +102,7 @@ def has_ability(self, ability):
         return False
 
     async def which_plugin(self):
-        file_svc = BaseService.get_service('file_svc')
-        for plugin in os.listdir('plugins'):
-            if await file_svc.walk_file_path(os.path.join('plugins', plugin, 'data', ''), '%s.yml' % self.adversary_id):
-                return plugin
-        return None
+        return self.plugin
 
     def check_repeatable_abilities(self, ability_list):
         return any(ab.repeatable for ab_id in self.atomic_ordering for ab in ability_list if ab.ability_id == ab_id)

app/objects/c_planner.py

@@ -1,11 +1,9 @@
-import os
 import uuid
 
 import marshmallow as ma
 
 from app.objects.interfaces.i_object import FirstClassObjectInterface
 from app.utility.base_object import BaseObject
-from app.utility.base_service import BaseService
 from app.objects.secondclass.c_fact import Fact, FactSchema
 
 
@@ -18,6 +16,7 @@ class PlannerSchema(ma.Schema):
     stopping_conditions = ma.fields.List(ma.fields.Nested(FactSchema()))
     ignore_enforcement_modules = ma.fields.List(ma.fields.String())
     allow_repeatable_abilities = ma.fields.Boolean()
+    plugin = ma.fields.String(missing=None)
 
     @ma.post_load()
     def build_planner(self, data, **kwargs):
@@ -34,7 +33,7 @@ def unique(self):
         return self.hash(self.name)
 
     def __init__(self, name='', planner_id='', module='', params=None, stopping_conditions=None, description=None,
-                 ignore_enforcement_modules=(), allow_repeatable_abilities=False):
+                 ignore_enforcement_modules=(), allow_repeatable_abilities=False, plugin=''):
         super().__init__()
         self.name = name
         self.planner_id = planner_id if planner_id else str(uuid.uuid4())
@@ -44,6 +43,7 @@ def __init__(self, name='', planner_id='', module='', params=None, stopping_cond
         self.stopping_conditions = self._set_stopping_conditions(stopping_conditions)
         self.ignore_enforcement_modules = ignore_enforcement_modules
         self.allow_repeatable_abilities = allow_repeatable_abilities
+        self.plugin = plugin
 
     def store(self, ram):
         existing = self.retrieve(ram['planners'], self.unique)
@@ -53,14 +53,11 @@ def store(self, ram):
         else:
             existing.update('stopping_conditions', self.stopping_conditions)
             existing.update('params', self.params)
+            existing.update('plugin', self.plugin)
         return existing
 
     async def which_plugin(self):
-        file_svc = BaseService.get_service('file_svc')
-        for plugin in os.listdir('plugins'):
-            if await file_svc.walk_file_path(os.path.join('plugins', plugin, 'data', ''), '%s.yml' % self.planner_id):
-                return plugin
-        return None
+        return self.plugin
 
     @staticmethod
     def _set_stopping_conditions(conditions):

app/objects/c_source.py

@@ -33,6 +33,7 @@ class SourceSchema(ma.Schema):
     rules = ma.fields.List(ma.fields.Nested(RuleSchema))
     adjustments = ma.fields.List(ma.fields.Nested(AdjustmentSchema))
     relationships = ma.fields.List(ma.fields.Nested(RelationshipSchema))
+    plugin = ma.fields.String(missing=None)
 
     @ma.pre_load
     def fix_adjustments(self, in_data, **_):
@@ -81,14 +82,15 @@ class Source(FirstClassObjectInterface, BaseObject):
     def unique(self):
         return self.hash('%s' % self.id)
 
-    def __init__(self, name='', id='', facts=(), relationships=(), rules=(), adjustments=()):
+    def __init__(self, name='', id='', facts=(), relationships=(), rules=(), adjustments=(), plugin=''):
         super().__init__()
         self.id = id if id else str(uuid.uuid4())
         self.name = name
         self.facts = facts
         self.rules = rules
         self.adjustments = adjustments
         self.relationships = relationships
+        self.plugin = plugin
 
     def store(self, ram):
         existing = self.retrieve(ram['sources'], self.unique)
@@ -99,4 +101,5 @@ def store(self, ram):
         existing.update('facts', self.facts)
         existing.update('rules', self.rules)
         existing.update('relationships', self.relationships)
+        existing.update('plugin', self.plugin)
         return existing

app/service/data_svc.py

@@ -7,6 +7,7 @@
 import tarfile
 import shutil
 import warnings
+import pathlib
 from importlib import import_module
 
 from app.objects.c_ability import Ability
@@ -160,14 +161,16 @@ async def load_ability_file(self, filename, access):
                 requirements = await self._load_ability_requirements(ab.pop('requirements', []))
                 buckets = ab.pop('buckets', [tactic])
                 ab.pop('access', None)
+                plugin = self._get_plugin_name(filename)
+                ab.pop('plugin', plugin)
 
                 if tactic and tactic not in filename:
                     self.log.error('Ability=%s has wrong tactic' % id)
 
                 await self._create_ability(ability_id=ability_id, name=name, description=description, tactic=tactic,
                                            technique_id=technique_id, technique_name=technique_name,
                                            executors=executors, requirements=requirements, privilege=privilege,
-                                           repeatable=repeatable, buckets=buckets, access=access, singleton=singleton,
+                                           repeatable=repeatable, buckets=buckets, access=access, singleton=singleton, plugin=plugin,
                                            **ab)
 
     async def convert_v0_ability_executor(self, ability_data: dict):
@@ -239,6 +242,7 @@ async def load_yaml_file(self, object_class, filename, access):
         for src in self.strip_yml(filename):
             obj = object_class.load(src)
             obj.access = access
+            obj.plugin = self._get_plugin_name(filename)
             await self.store(obj)
 
     async def _load(self, plugins=()):
@@ -336,11 +340,11 @@ async def _load_data_encoders(self, plugins):
 
     async def _create_ability(self, ability_id, name=None, description=None, tactic=None, technique_id=None,
                               technique_name=None, executors=None, requirements=None, privilege=None,
-                              repeatable=False, buckets=None, access=None, singleton=False, **kwargs):
+                              repeatable=False, buckets=None, access=None, singleton=False, plugin='', **kwargs):
         ability = Ability(ability_id=ability_id, name=name, description=description, tactic=tactic,
                           technique_id=technique_id, technique_name=technique_name, executors=executors,
                           requirements=requirements, privilege=privilege, repeatable=repeatable, buckets=buckets,
-                          access=access, singleton=singleton, **kwargs)
+                          access=access, singleton=singleton, plugin=plugin, **kwargs)
         return await self.store(ability)
 
     async def _prune_non_critical_data(self):
@@ -410,3 +414,7 @@ async def _verify_default_objective_exists(self):
     async def _verify_adversary_profiles(self):
         for adv in await self.locate('adversaries'):
             adv.verify(log=self.log, abilities=self.ram['abilities'], objectives=self.ram['objectives'])
+
+    def _get_plugin_name(self, filename):
+        plugin_path = pathlib.PurePath(filename).parts
+        return plugin_path[1] if 'plugins' in plugin_path else ''

static/js/core.js

@@ -28,8 +28,6 @@ function alpineCore() {
             if (tabName === 'fieldmanual') {
                 restRequest('GET', null, (data) => { this.setTabContent({ name: tabName, contentID: `tab-${tabName}`, address: address }, data); }, address);
                 return;
-            } else if (tabName === 'stockpile' || tabName === 'atomic') {
-                return;
             }
 
             // If tab is already open, jump to it

templates/BLUE.html

@@ -46,7 +46,7 @@
                     <ul class="menu-list">
                         {% for plugin in plugins | sort(attribute='name') %}
                         <li>
-                            <a href="#home" x-on:click="addTab('{{ plugin.name }}', '{{ plugin.address }}')" class="nav-item" x-bind:class="{ 'disabled': '{{ plugin.name }}' === 'atomic' || '{{ plugin.name }}' === 'stockpile' }">
+                            <a href="#home" x-on:click="addTab('{{ plugin.name }}', '{{ plugin.address }}')" class="nav-item">
                                 {{ plugin.name }}
                                     <template x-if="`{{plugin.name | e}}` === 'fieldmanual'">
                                         <sup><i class="fas fa-external-link-alt pl-1 is-size-7"></i></sup>

templates/RED.html

@@ -46,7 +46,7 @@
                     <ul class="menu-list">
                         {% for plugin in plugins | sort(attribute='name') %}
                             <li>
-                                <a href="#home" x-on:click="addTab('{{ plugin.name }}', '{{ plugin.address }}')" class="nav-item" x-bind:class="{ 'disabled': '{{ plugin.name }}' === 'atomic' || '{{ plugin.name }}' === 'stockpile' }">
+                                <a href="#home" x-on:click="addTab('{{ plugin.name }}', '{{ plugin.address }}')" class="nav-item">
                                     {{ plugin.name }}
                                     <template x-if="`{{plugin.name | e}}` === 'fieldmanual'">
                                         <sup><i class="fas fa-external-link-alt pl-1 is-size-7"></i></sup>

tests/api/v2/handlers/test_abilities_api.py

@@ -23,14 +23,15 @@ def new_ability_payload():
             'privilege': '',
             'repeatable': False,
             'requirements': [],
-            'singleton': False
+            'singleton': False,
+            'plugin': ''
             }
 
 
 @pytest.fixture
 def updated_ability_payload(test_ability):
     ability_data = test_ability.schema.dump(test_ability)
-    ability_data.update(dict(name='an updated test ability', tactic='defense-evasion'))
+    ability_data.update(dict(name='an updated test ability', tactic='defense-evasion', plugin=''))
     return ability_data
 
 
@@ -39,15 +40,15 @@ def replaced_ability_payload(test_ability):
     ability_data = test_ability.schema.dump(test_ability)
     test_executor_linux = Executor(name='sh', platform='linux', command='whoami')
     ability_data.update(dict(name='replaced test ability', tactic='collection', technique_name='discovery',
-                             technique_id='2', executors=[ExecutorSchema().dump(test_executor_linux)]))
+                             technique_id='2', executors=[ExecutorSchema().dump(test_executor_linux)], plugin=''))
     return ability_data
 
 
 @pytest.fixture
 def test_ability(loop, api_v2_client, executor):
     executor_linux = executor(name='sh', platform='linux')
     ability = Ability(ability_id='123', name='Test Ability', executors=[executor_linux],
-                      technique_name='collection', technique_id='1', description='', privilege='', tactic='discovery')
+                      technique_name='collection', technique_id='1', description='', privilege='', tactic='discovery', plugin='testplugin')
     loop.run_until_complete(BaseService.get_service('data_svc').store(ability))
     return ability
 

tests/api/v2/handlers/test_adversaries_api.py

@@ -32,7 +32,8 @@ def new_adversary_payload():
         'adversary_id': '456',
         'objective': '495a9828-cab1-44dd-a0ca-66e58177d8cc',
         'tags': [],
-        'atomic_ordering': []
+        'atomic_ordering': [],
+        'plugin': ''
     }
 
 
@@ -49,7 +50,8 @@ def test_adversary(loop):
                           'adversary_id': '123',
                           'objective': '495a9828-cab1-44dd-a0ca-66e58177d8cc',
                           'tags': [],
-                          'atomic_ordering': []}
+                          'atomic_ordering': [],
+                          'plugin': ''}
     test_adversary = AdversarySchema().load(expected_adversary)
     loop.run_until_complete(BaseService.get_service('data_svc').store(test_adversary))
     return test_adversary

tests/api/v2/handlers/test_planners_api.py

@@ -8,7 +8,7 @@
 
 @pytest.fixture
 def test_planner(loop, api_v2_client):
-    planner = Planner(name="test planner", planner_id="123", description="a test planner")
+    planner = Planner(name="test planner", planner_id="123", description="a test planner", plugin="planner")
     loop.run_until_complete(BaseService.get_service('data_svc').store(planner))
     return planner
 

tests/api/v2/handlers/test_sources_api.py

@@ -34,7 +34,8 @@ def new_source_payload():
         'name': 'new test source',
         'facts': [fact],
         'rules': [rule.schema.dump(rule)],
-        'relationships': [relationship]
+        'relationships': [relationship],
+        'plugin': ''
     }
     return source
 
@@ -80,6 +81,7 @@ def expected_updated_source_dump(updated_source_payload, mocker, mock_time):
         source = SourceSchema().load(updated_source_payload)
         dumped_obj = source.display_schema.dump(source)
         dumped_obj['relationships'][0]['unique'] = mock.ANY
+        dumped_obj['plugin'] = ''
         return dumped_obj
 
 

tests/services/test_rest_svc.py

@@ -75,10 +75,10 @@ def test_delete_operation(self, loop, rest_svc, data_svc):
                               'adversary': {'description': 'an empty adversary profile', 'name': 'ad-hoc',
                                             'adversary_id': 'ad-hoc', 'atomic_ordering': [],
                                             'objective': '495a9828-cab1-44dd-a0ca-66e58177d8cc',
-                                            'tags': [], 'has_repeatable_abilities': False}, 'state': 'finished',
+                                            'tags': [], 'has_repeatable_abilities': False, 'plugin': None}, 'state': 'finished',
                               'planner': {'name': 'test', 'description': None, 'module': 'test',
                                           'stopping_conditions': [], 'params': {}, 'allow_repeatable_abilities': False,
-                                          'ignore_enforcement_modules': [], 'id': '123'}, 'jitter': '2/8',
+                                          'ignore_enforcement_modules': [], 'id': '123', 'plugin': ''}, 'jitter': '2/8',
                               'host_group': [{'trusted': True, 'architecture': 'unknown', 'watchdog': 0,
                                               'contact': 'unknown', 'username': 'unknown', 'links': [], 'sleep_max': 8,
                                               'exe_name': 'unknown', 'executors': ['pwsh', 'psh'], 'ppid': 0,
@@ -147,11 +147,11 @@ def test_create_operation(self, loop, rest_svc, data_svc):
         want = {'name': 'Test',
                 'adversary': {'description': 'an empty adversary profile', 'name': 'ad-hoc', 'adversary_id': 'ad-hoc',
                               'atomic_ordering': [], 'objective': '495a9828-cab1-44dd-a0ca-66e58177d8cc', 'tags': [],
-                              'has_repeatable_abilities': False},
+                              'has_repeatable_abilities': False, 'plugin': None},
                 'state': 'finished',
                 'planner': {'name': 'test', 'description': None, 'module': 'test', 'stopping_conditions': [],
                             'params': {},
-                            'ignore_enforcement_modules': [], 'id': '123', 'allow_repeatable_abilities': False},
+                            'ignore_enforcement_modules': [], 'id': '123', 'allow_repeatable_abilities': False, 'plugin': ''},
                 'jitter': '2/8',
                 'group': '',
                 'source': '',