From 0801b4c612fbedbd92c3f6c91c61c70e45339209 Mon Sep 17 00:00:00 2001 From: Adam Gaudreau Date: Mon, 1 Nov 2021 15:39:38 -0400 Subject: [PATCH 1/6] Add plugin field to adversaries, abilities, and planners --- app/objects/c_ability.py | 11 +++++------ app/objects/c_adversary.py | 11 +++++------ app/objects/c_planner.py | 11 +++++------ app/service/data_svc.py | 8 +++++--- static/js/core.js | 2 -- templates/BLUE.html | 2 +- templates/RED.html | 2 +- tests/api/v2/handlers/test_abilities_api.py | 5 +++-- tests/api/v2/handlers/test_adversaries_api.py | 6 ++++-- tests/api/v2/handlers/test_planners_api.py | 2 +- tests/services/test_rest_svc.py | 8 ++++---- 11 files changed, 34 insertions(+), 34 deletions(-) diff --git a/app/objects/c_ability.py b/app/objects/c_ability.py index d4dd8cf1c..c718b27b6 100644 --- a/app/objects/c_ability.py +++ b/app/objects/c_ability.py @@ -27,6 +27,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 +59,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=None, **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 +81,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 +105,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)) diff --git a/app/objects/c_adversary.py b/app/objects/c_adversary.py index de8dfab7a..3b1e3079f 100644 --- a/app/objects/c_adversary.py +++ b/app/objects/c_adversary.py @@ -20,6 +20,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 +58,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=None): super().__init__() self.adversary_id = adversary_id if adversary_id else str(uuid.uuid4()) self.name = name @@ -66,6 +67,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 +80,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 +104,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) diff --git a/app/objects/c_planner.py b/app/objects/c_planner.py index 8b0e671c2..7743204f6 100644 --- a/app/objects/c_planner.py +++ b/app/objects/c_planner.py @@ -18,6 +18,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 +35,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=None): super().__init__() self.name = name self.planner_id = planner_id if planner_id else str(uuid.uuid4()) @@ -44,6 +45,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 +55,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 """ PRIVATE """ diff --git a/app/service/data_svc.py b/app/service/data_svc.py index f5a5e3b89..78ef95302 100644 --- a/app/service/data_svc.py +++ b/app/service/data_svc.py @@ -160,6 +160,7 @@ 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 = ab.pop('plugin', filename.split('/')[1]) if tactic and tactic not in filename: self.log.error('Ability=%s has wrong tactic' % id) @@ -167,7 +168,7 @@ async def load_ability_file(self, filename, access): 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 +240,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 = filename.split('/')[1] await self.store(obj) """ PRIVATE """ @@ -338,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=None, **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): diff --git a/static/js/core.js b/static/js/core.js index aeb8a6d15..559fd2377 100644 --- a/static/js/core.js +++ b/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 diff --git a/templates/BLUE.html b/templates/BLUE.html index 9b63af0d5..81802475d 100644 --- a/templates/BLUE.html +++ b/templates/BLUE.html @@ -46,7 +46,7 @@