diff --git a/app/objects/c_ability.py b/app/objects/c_ability.py index d4dd8cf1c..51283cda2 100644 --- a/app/objects/c_ability.py +++ b/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)) diff --git a/app/objects/c_adversary.py b/app/objects/c_adversary.py index de8dfab7a..32c5114b5 100644 --- a/app/objects/c_adversary.py +++ b/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) diff --git a/app/objects/c_planner.py b/app/objects/c_planner.py index 0d371f864..15d6e062f 100644 --- a/app/objects/c_planner.py +++ b/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): diff --git a/app/objects/c_source.py b/app/objects/c_source.py index 755553074..f8dd33624 100644 --- a/app/objects/c_source.py +++ b/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,7 +82,7 @@ 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 @@ -89,6 +90,7 @@ def __init__(self, name='', id='', facts=(), relationships=(), rules=(), adjustm 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 diff --git a/app/service/data_svc.py b/app/service/data_svc.py index 6243dfafd..5434ed1cf 100644 --- a/app/service/data_svc.py +++ b/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,6 +161,8 @@ 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) @@ -167,7 +170,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 +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 '' 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 d0febd1bd..80cb01ce0 100644 --- a/templates/BLUE.html +++ b/templates/BLUE.html @@ -46,7 +46,7 @@