From 3c16febc96de911e11fb8093446b5cef6b17fe63 Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Fri, 9 Jan 2026 03:19:47 +0000 Subject: [PATCH 01/10] error handling for load_ability_file, update _get_plugin_name --- app/service/data_svc.py | 65 +++++++++++++++++++++++------------------ 1 file changed, 37 insertions(+), 28 deletions(-) diff --git a/app/service/data_svc.py b/app/service/data_svc.py index 3170d706f..ef083d0d4 100644 --- a/app/service/data_svc.py +++ b/app/service/data_svc.py @@ -5,6 +5,7 @@ import os import pickle import tarfile +import re import shutil import warnings import pathlib @@ -152,32 +153,39 @@ async def remove(self, object_name, match): self.log.error('[!] REMOVE: %s' % e) async def load_ability_file(self, filename, access): - for entries in self.strip_yml(filename): - for ab in entries: - ability_id = ab.pop('id', None) - name = ab.pop('name', '') - description = ab.pop('description', '') - tactic = ab.pop('tactic', None) - executors = await self.convert_v0_ability_executor(ab) - technique_id = self.convert_v0_ability_technique_id(ab) - technique_name = self.convert_v0_ability_technique_name(ab) - privilege = ab.pop('privilege', None) - repeatable = ab.pop('repeatable', False) - singleton = ab.pop('singleton', False) - requirements = await self.convert_v0_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' % ability_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, plugin=plugin, - **ab) + try: + for entries in self.strip_yml(filename): + for ab in entries: + if type(ab) is not dict: + self.log.error(f'Malformed ability file {filename}. Expected ability entry to be a dictionary, received {type(ab)} instead.') + continue + ability_id = ab.pop('id', None) + name = ab.pop('name', '') + description = ab.pop('description', '') + tactic = ab.pop('tactic', None) + executors = await self.convert_v0_ability_executor(ab) + technique_id = self.convert_v0_ability_technique_id(ab) + technique_name = self.convert_v0_ability_technique_name(ab) + privilege = ab.pop('privilege', None) + repeatable = ab.pop('repeatable', False) + singleton = ab.pop('singleton', False) + requirements = await self.convert_v0_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.warn(f'Tactic for ability={ability_id} is not in the ability file path {filename}.') + self.log.warn('Please check that the ability is labeled with the correct tactic and is in the correct location.') + + 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, plugin=plugin, + **ab) + except Exception as e: + self.log.exception(f'Failed to load ability file {filename}: {e}') async def convert_v0_ability_executor(self, ability_data: dict): """Checks if ability file follows v0 executor format, otherwise assumes v1 ability formatting.""" @@ -502,8 +510,9 @@ async def _verify_adversary_profiles(self): 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 '' + normalized_path = '/'.join(pathlib.PurePath(filename).parts) + matches = re.findall(r'.*/?plugins/([^/]+)/?.*', normalized_path) + return matches[0] if matches else '' async def get_facts_from_source(self, fact_source_id): fact_sources = await self.locate('sources', match=dict(id=fact_source_id)) From 6a345d8be13c7313ae91d2f8f45930b860370f14 Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Fri, 9 Jan 2026 03:21:05 +0000 Subject: [PATCH 02/10] unit tests for _get_plugin_name --- tests/services/test_data_svc.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/services/test_data_svc.py b/tests/services/test_data_svc.py index bf6528bb7..b8118a431 100644 --- a/tests/services/test_data_svc.py +++ b/tests/services/test_data_svc.py @@ -241,3 +241,22 @@ def _mock_apply_payload_config(config=None, **_): } } mock_apply_config2.assert_called_once_with(name='payloads', config=expected_config_part2) + + def test_get_plugin_name(self, data_svc): + assert 'test' == data_svc._get_plugin_name('plugins/test') + assert 'test' == data_svc._get_plugin_name('plugins/test/') + assert 'test' == data_svc._get_plugin_name('plugins/test/data') + assert 'test' == data_svc._get_plugin_name('plugins/test/data/abilities') + assert 'test' == data_svc._get_plugin_name('plugins/test/data/abilities/collection/123.yml') + assert 'test' == data_svc._get_plugin_name('/full/path/to/plugins/test/data/abilities/collection/123.yml') + assert '' == data_svc._get_plugin_name('test') + assert '' == data_svc._get_plugin_name('plugins') + assert '' == data_svc._get_plugin_name('plugins/') + assert '' == data_svc._get_plugin_name('/full/path/to/plugins') + assert '' == data_svc._get_plugin_name('/full/path/to/plugins/') + assert '' == data_svc._get_plugin_name('plugin/test') + assert '' == data_svc._get_plugin_name('plugin/test/') + assert '' == data_svc._get_plugin_name('plugin/test/data') + assert '' == data_svc._get_plugin_name('plugin/test/data/abilities') + assert '' == data_svc._get_plugin_name('plugin/test/data/abilities/collection/123.yml') + assert '' == data_svc._get_plugin_name('/full/path/to/plugin/test/data/abilities/collection/123.yml') From 6ec1a0c3fcd703651088c050bc6c1041e6f47bd4 Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Fri, 9 Jan 2026 03:21:11 +0000 Subject: [PATCH 03/10] style fix --- app/service/data_svc.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/app/service/data_svc.py b/app/service/data_svc.py index ef083d0d4..c96f656ac 100644 --- a/app/service/data_svc.py +++ b/app/service/data_svc.py @@ -180,10 +180,10 @@ async def load_ability_file(self, filename, access): self.log.warn('Please check that the ability is labeled with the correct tactic and is in the correct location.') 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, plugin=plugin, - **ab) + technique_id=technique_id, technique_name=technique_name, + executors=executors, requirements=requirements, privilege=privilege, + repeatable=repeatable, buckets=buckets, access=access, singleton=singleton, plugin=plugin, + **ab) except Exception as e: self.log.exception(f'Failed to load ability file {filename}: {e}') From 0a96a3289da0c27890331f736ae771ba4035b83c Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Fri, 9 Jan 2026 08:28:00 +0000 Subject: [PATCH 04/10] ensure ability id is string --- app/service/data_svc.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/app/service/data_svc.py b/app/service/data_svc.py index c96f656ac..48abdefc4 100644 --- a/app/service/data_svc.py +++ b/app/service/data_svc.py @@ -160,6 +160,8 @@ async def load_ability_file(self, filename, access): self.log.error(f'Malformed ability file {filename}. Expected ability entry to be a dictionary, received {type(ab)} instead.') continue ability_id = ab.pop('id', None) + if ability_id != None and type(ability_id) is not str: + ability_id = str(ability_id) name = ab.pop('name', '') description = ab.pop('description', '') tactic = ab.pop('tactic', None) From b737a0e536ce277a99b1c7dedf8f10ed6e12987d Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Fri, 9 Jan 2026 08:28:37 +0000 Subject: [PATCH 05/10] add unit test for load_ability_file --- tests/services/test_data_svc.py | 101 ++++++++++++++++++++++++++++++++ 1 file changed, 101 insertions(+) diff --git a/tests/services/test_data_svc.py b/tests/services/test_data_svc.py index b8118a431..17ae8c98a 100644 --- a/tests/services/test_data_svc.py +++ b/tests/services/test_data_svc.py @@ -77,10 +77,81 @@ } +ABILITY_YAMLS = { + 'plugins/testing/data/discovery/764efa883dda1e11db47671c4a3bbd9e.yml': [yaml.safe_load(''' +--- + +- id: 764efa883dda1e11db47671c4a3bbd9e + name: Find deletable dirs (per user) + description: Discover all directories containing deletable files by user + tactic: discovery + technique: + attack_id: T1082 + name: System Information Discovery + platforms: + darwin: + sh: + command: | + testcommand + linux: + sh: + command: | + testcommand +''')], + 'plugins/testing/data/discovery/848aa201-4b00-4f08-ae3a-3e84dfb5065c.yml': [yaml.safe_load(''' +--- + +- id: 848aa201-4b00-4f08-ae3a-3e84dfb5065c + name: Find deletable dirs (per user) + description: Discover all directories containing deletable files by user + tactic: discovery + technique: + attack_id: T1082 + name: System Information Discovery + platforms: + darwin: + sh: + command: | + testcommand + linux: + sh: + command: | + testcommand +''')], + 'plugins/testing/data/discovery/101.yml': [yaml.safe_load(''' +--- + +- id: 101 + name: Find deletable dirs (per user) + description: Discover all directories containing deletable files by user + tactic: discovery + technique: + attack_id: T1082 + name: System Information Discovery + platforms: + darwin: + sh: + command: | + testcommand + linux: + sh: + command: | + testcommand +''')], + 'plugins/testing/data/discovery/102.yml': [yaml.safe_load(''' +malformed +''')], +} + + def strip_payload_yaml(path): return PAYLOAD_CONFIG_YAMLS.get(path, []) +def strip_ability_yaml(path): + return ABILITY_YAMLS.get(path, []) + + def async_mock_return(to_return): mock_future = asyncio.Future() mock_future.set_result(to_return) @@ -242,6 +313,36 @@ def _mock_apply_payload_config(config=None, **_): } mock_apply_config2.assert_called_once_with(name='payloads', config=expected_config_part2) + @mock.patch.object(BaseWorld, 'strip_yml', wraps=strip_ability_yaml) + async def test_load_ability_file(self, event_loop, data_svc): + with patch.object(DataService, '_create_ability', return_value=None) as mock_create_ability: + await data_svc.load_ability_file('plugins/testing/data/discovery/764efa883dda1e11db47671c4a3bbd9e.yml', BaseWorld.Access.RED) + mock_create_ability.assert_called_once_with(ability_id='764efa883dda1e11db47671c4a3bbd9e', name='Find deletable dirs (per user)', + description='Discover all directories containing deletable files by user', + tactic='discovery', technique_id='T1082', technique_name='System Information Discovery', + executors=mock.ANY, requirements=[], privilege=None, + repeatable=False, buckets=['discovery'], access=BaseWorld.Access.RED, singleton=False, plugin='testing') + + with patch.object(DataService, '_create_ability', return_value=None) as mock_create_ability: + await data_svc.load_ability_file('plugins/testing/data/discovery/848aa201-4b00-4f08-ae3a-3e84dfb5065c.yml', BaseWorld.Access.RED) + mock_create_ability.assert_called_once_with(ability_id='848aa201-4b00-4f08-ae3a-3e84dfb5065c', name='Find deletable dirs (per user)', + description='Discover all directories containing deletable files by user', + tactic='discovery', technique_id='T1082', technique_name='System Information Discovery', + executors=mock.ANY, requirements=[], privilege=None, + repeatable=False, buckets=['discovery'], access=BaseWorld.Access.RED, singleton=False, plugin='testing') + + with patch.object(DataService, '_create_ability', return_value=None) as mock_create_ability: + await data_svc.load_ability_file('plugins/testing/data/discovery/101.yml', BaseWorld.Access.RED) + mock_create_ability.assert_called_once_with(ability_id='101', name='Find deletable dirs (per user)', + description='Discover all directories containing deletable files by user', + tactic='discovery', technique_id='T1082', technique_name='System Information Discovery', + executors=mock.ANY, requirements=[], privilege=None, + repeatable=False, buckets=['discovery'], access=BaseWorld.Access.RED, singleton=False, plugin='testing') + + with patch.object(DataService, '_create_ability', return_value=None) as mock_create_ability: + await data_svc.load_ability_file('plugins/testing/data/discovery/102.yml', BaseWorld.Access.RED) + mock_create_ability.assert_not_called() + def test_get_plugin_name(self, data_svc): assert 'test' == data_svc._get_plugin_name('plugins/test') assert 'test' == data_svc._get_plugin_name('plugins/test/') From d535756ac3b863e15b86dfbc299bfdca7bb05add Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Fri, 9 Jan 2026 09:08:16 +0000 Subject: [PATCH 06/10] compare executors --- app/objects/secondclass/c_executor.py | 12 ++++++++++++ tests/services/test_data_svc.py | 16 +++++++++++++--- 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/app/objects/secondclass/c_executor.py b/app/objects/secondclass/c_executor.py index 870a921d2..fbb38519e 100644 --- a/app/objects/secondclass/c_executor.py +++ b/app/objects/secondclass/c_executor.py @@ -83,6 +83,18 @@ def __getattr__(self, item): def replace_cleanup(self, command, payload): return command.replace(self.RESERVED['payload'], payload) + def __eq__(self, other): + """Overrides the default eq implementation""" + if isinstance(other, Executor): + return self.name == other.name and self.platform == other.platform and \ + self.command == other.command and self.code == other.code and \ + self.language == other.language and self.build_target == other.build_target and \ + self.payloads == other.payloads and self.uploads == other.uploads and \ + self.timeout == other.timeout and self.parsers == other.parsers and \ + self.cleanup == other.cleanup and self.variations == other.variations and \ + self.additional_info == other.additional_info + return False + def get_variations(data): variations = [] diff --git a/tests/services/test_data_svc.py b/tests/services/test_data_svc.py index 17ae8c98a..377b2e4f2 100644 --- a/tests/services/test_data_svc.py +++ b/tests/services/test_data_svc.py @@ -315,12 +315,22 @@ def _mock_apply_payload_config(config=None, **_): @mock.patch.object(BaseWorld, 'strip_yml', wraps=strip_ability_yaml) async def test_load_ability_file(self, event_loop, data_svc): + want_executors = [ + Executor(name='sh', platform='darwin', command='testcommand', + code=None, language=None, build_target=None, + payloads=None, uploads=None, timeout=60, + parsers=[], cleanup=None, variations=[]), + Executor(name='sh', platform='linux', command='testcommand', + code=None, language=None, build_target=None, + payloads=None, uploads=None, timeout=60, + parsers=[], cleanup=None, variations=[]) + ] with patch.object(DataService, '_create_ability', return_value=None) as mock_create_ability: await data_svc.load_ability_file('plugins/testing/data/discovery/764efa883dda1e11db47671c4a3bbd9e.yml', BaseWorld.Access.RED) mock_create_ability.assert_called_once_with(ability_id='764efa883dda1e11db47671c4a3bbd9e', name='Find deletable dirs (per user)', description='Discover all directories containing deletable files by user', tactic='discovery', technique_id='T1082', technique_name='System Information Discovery', - executors=mock.ANY, requirements=[], privilege=None, + executors=want_executors, requirements=[], privilege=None, repeatable=False, buckets=['discovery'], access=BaseWorld.Access.RED, singleton=False, plugin='testing') with patch.object(DataService, '_create_ability', return_value=None) as mock_create_ability: @@ -328,7 +338,7 @@ async def test_load_ability_file(self, event_loop, data_svc): mock_create_ability.assert_called_once_with(ability_id='848aa201-4b00-4f08-ae3a-3e84dfb5065c', name='Find deletable dirs (per user)', description='Discover all directories containing deletable files by user', tactic='discovery', technique_id='T1082', technique_name='System Information Discovery', - executors=mock.ANY, requirements=[], privilege=None, + executors=want_executors, requirements=[], privilege=None, repeatable=False, buckets=['discovery'], access=BaseWorld.Access.RED, singleton=False, plugin='testing') with patch.object(DataService, '_create_ability', return_value=None) as mock_create_ability: @@ -336,7 +346,7 @@ async def test_load_ability_file(self, event_loop, data_svc): mock_create_ability.assert_called_once_with(ability_id='101', name='Find deletable dirs (per user)', description='Discover all directories containing deletable files by user', tactic='discovery', technique_id='T1082', technique_name='System Information Discovery', - executors=mock.ANY, requirements=[], privilege=None, + executors=want_executors, requirements=[], privilege=None, repeatable=False, buckets=['discovery'], access=BaseWorld.Access.RED, singleton=False, plugin='testing') with patch.object(DataService, '_create_ability', return_value=None) as mock_create_ability: From 5ded2159ce86f634bcda064cbb0d6666dfc537c3 Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Fri, 9 Jan 2026 09:08:52 +0000 Subject: [PATCH 07/10] proper none comparison --- app/service/data_svc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app/service/data_svc.py b/app/service/data_svc.py index 48abdefc4..aacb7f441 100644 --- a/app/service/data_svc.py +++ b/app/service/data_svc.py @@ -160,7 +160,7 @@ async def load_ability_file(self, filename, access): self.log.error(f'Malformed ability file {filename}. Expected ability entry to be a dictionary, received {type(ab)} instead.') continue ability_id = ab.pop('id', None) - if ability_id != None and type(ability_id) is not str: + if ability_id is not None and type(ability_id) is not str: ability_id = str(ability_id) name = ab.pop('name', '') description = ab.pop('description', '') From 8f3c3ff86020015b4624cfce123878a999067459 Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Tue, 27 Jan 2026 14:29:50 +0000 Subject: [PATCH 08/10] simplify eq method --- app/objects/secondclass/c_executor.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/app/objects/secondclass/c_executor.py b/app/objects/secondclass/c_executor.py index fbb38519e..4a5970ea7 100644 --- a/app/objects/secondclass/c_executor.py +++ b/app/objects/secondclass/c_executor.py @@ -85,15 +85,13 @@ def replace_cleanup(self, command, payload): def __eq__(self, other): """Overrides the default eq implementation""" - if isinstance(other, Executor): - return self.name == other.name and self.platform == other.platform and \ - self.command == other.command and self.code == other.code and \ - self.language == other.language and self.build_target == other.build_target and \ - self.payloads == other.payloads and self.uploads == other.uploads and \ - self.timeout == other.timeout and self.parsers == other.parsers and \ - self.cleanup == other.cleanup and self.variations == other.variations and \ - self.additional_info == other.additional_info - return False + return isinstance(other, Executor) and self.name == other.name and self.platform == other.platform and \ + self.command == other.command and self.code == other.code and \ + self.language == other.language and self.build_target == other.build_target and \ + self.payloads == other.payloads and self.uploads == other.uploads and \ + self.timeout == other.timeout and self.parsers == other.parsers and \ + self.cleanup == other.cleanup and self.variations == other.variations and \ + self.additional_info == other.additional_info def get_variations(data): From 3467c3e082dbdf71170fcf0925d6af3d725be9ac Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Tue, 27 Jan 2026 15:50:17 +0000 Subject: [PATCH 09/10] test additional lines --- tests/services/test_data_svc.py | 26 ++++++++++++++++++++------ 1 file changed, 20 insertions(+), 6 deletions(-) diff --git a/tests/services/test_data_svc.py b/tests/services/test_data_svc.py index 16a7b97a8..dbcbda5e0 100644 --- a/tests/services/test_data_svc.py +++ b/tests/services/test_data_svc.py @@ -1,5 +1,6 @@ import glob import json +import logging import yaml from unittest import mock @@ -123,7 +124,7 @@ - id: 101 name: Find deletable dirs (per user) description: Discover all directories containing deletable files by user - tactic: discovery + tactic: purposefullywrongtactic technique: attack_id: T1082 name: System Information Discovery @@ -306,8 +307,9 @@ def _mock_apply_payload_config(config=None, **_): } mock_apply_config2.assert_called_once_with(name='payloads', config=expected_config_part2) + @mock.patch.object(logging.Logger, 'warn') @mock.patch.object(BaseWorld, 'strip_yml', wraps=strip_ability_yaml) - async def test_load_ability_file(self, event_loop, data_svc): + async def test_load_ability_file(self, mock_strip_yml, mock_warn, data_svc): want_executors = [ Executor(name='sh', platform='darwin', command='testcommand', code=None, language=None, build_target=None, @@ -336,15 +338,27 @@ async def test_load_ability_file(self, event_loop, data_svc): with patch.object(DataService, '_create_ability', return_value=None) as mock_create_ability: await data_svc.load_ability_file('plugins/testing/data/discovery/101.yml', BaseWorld.Access.RED) + mock_warn.assert_any_call('Tactic for ability=101 is not in the ability file path plugins/testing/data/discovery/101.yml.') + mock_warn.assert_called_with('Please check that the ability is labeled with the correct tactic and is in the correct location.') mock_create_ability.assert_called_once_with(ability_id='101', name='Find deletable dirs (per user)', description='Discover all directories containing deletable files by user', - tactic='discovery', technique_id='T1082', technique_name='System Information Discovery', + tactic='purposefullywrongtactic', technique_id='T1082', technique_name='System Information Discovery', executors=want_executors, requirements=[], privilege=None, - repeatable=False, buckets=['discovery'], access=BaseWorld.Access.RED, singleton=False, plugin='testing') + repeatable=False, buckets=['purposefullywrongtactic'], access=BaseWorld.Access.RED, singleton=False, plugin='testing') with patch.object(DataService, '_create_ability', return_value=None) as mock_create_ability: - await data_svc.load_ability_file('plugins/testing/data/discovery/102.yml', BaseWorld.Access.RED) - mock_create_ability.assert_not_called() + with patch.object(logging.Logger, 'error') as mock_error: + await data_svc.load_ability_file('plugins/testing/data/discovery/102.yml', BaseWorld.Access.RED) + mock_create_ability.assert_not_called() + assert mock_error.called + + # Test exception + with patch.object(DataService, '_create_ability', side_effect=Exception('mockexception')): + with patch.object(logging.Logger, 'exception') as mock_exception: + await data_svc.load_ability_file('plugins/testing/data/discovery/101.yml', BaseWorld.Access.RED) + mock_exception.assert_called_once_with(mock.ANY) + assert 'Failed to load ability file plugins/testing/data/discovery/101.yml' in mock_exception.call_args.args[0] + def test_get_plugin_name(self, data_svc): assert 'test' == data_svc._get_plugin_name('plugins/test') From 52f920ecc5a7123816e5caf750690566a5211fba Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Tue, 27 Jan 2026 16:01:55 +0000 Subject: [PATCH 10/10] address sonarcloud concerns --- app/service/data_svc.py | 10 ++++++---- tests/services/test_data_svc.py | 1 - 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/app/service/data_svc.py b/app/service/data_svc.py index aacb7f441..d5e9d8ccf 100644 --- a/app/service/data_svc.py +++ b/app/service/data_svc.py @@ -5,7 +5,6 @@ import os import pickle import tarfile -import re import shutil import warnings import pathlib @@ -512,9 +511,12 @@ async def _verify_adversary_profiles(self): adv.verify(log=self.log, abilities=self.ram['abilities'], objectives=self.ram['objectives']) def _get_plugin_name(self, filename): - normalized_path = '/'.join(pathlib.PurePath(filename).parts) - matches = re.findall(r'.*/?plugins/([^/]+)/?.*', normalized_path) - return matches[0] if matches else '' + path_components = pathlib.PurePath(filename).parts + num_parts = len(path_components) + for i, part in enumerate(path_components): + if part == 'plugins' and i < num_parts - 1: + return path_components[i + 1] + return '' async def get_facts_from_source(self, fact_source_id): fact_sources = await self.locate('sources', match=dict(id=fact_source_id)) diff --git a/tests/services/test_data_svc.py b/tests/services/test_data_svc.py index dbcbda5e0..f3b3fde16 100644 --- a/tests/services/test_data_svc.py +++ b/tests/services/test_data_svc.py @@ -359,7 +359,6 @@ async def test_load_ability_file(self, mock_strip_yml, mock_warn, data_svc): mock_exception.assert_called_once_with(mock.ANY) assert 'Failed to load ability file plugins/testing/data/discovery/101.yml' in mock_exception.call_args.args[0] - def test_get_plugin_name(self, data_svc): assert 'test' == data_svc._get_plugin_name('plugins/test') assert 'test' == data_svc._get_plugin_name('plugins/test/')