From d0d5ff99b9e76e3a519b34e163ae20b9b5ace36c Mon Sep 17 00:00:00 2001 From: deacon Date: Mon, 16 Mar 2026 09:37:12 -0400 Subject: [PATCH] test: add exhaustive pytest coverage for emu plugin --- plugins/emu/conf/default.yml | 2 + pytest.ini | 5 + tests/conftest.py | 229 ++++++++++ tests/test_emu_gui.py | 53 +++ tests/test_emu_svc.py | 640 +++++++++++++++++++++++++-- tests/test_group_filtered_planner.py | 431 ++++++++++-------- tests/test_hook.py | 116 +++++ tests/test_requirements.py | 150 +++++++ tests/test_vssadmin_shadow_parser.py | 108 ++++- 9 files changed, 1490 insertions(+), 244 deletions(-) create mode 100644 plugins/emu/conf/default.yml create mode 100644 pytest.ini create mode 100644 tests/conftest.py create mode 100644 tests/test_emu_gui.py create mode 100644 tests/test_hook.py create mode 100644 tests/test_requirements.py diff --git a/plugins/emu/conf/default.yml b/plugins/emu/conf/default.yml new file mode 100644 index 0000000..21f056b --- /dev/null +++ b/plugins/emu/conf/default.yml @@ -0,0 +1,2 @@ +evals_c2_host: 127.0.0.1 +evals_c2_port: 8888 diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..82297bd --- /dev/null +++ b/pytest.ini @@ -0,0 +1,5 @@ +[pytest] +asyncio_mode = auto +testpaths = tests +markers = + slow: marks tests as slow diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..45c0f0b --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,229 @@ +"""Shared fixtures for emu plugin tests.""" +import asyncio +import os +import yaml +import pytest + +from unittest.mock import MagicMock, AsyncMock, patch + + +def async_mock_return(to_return): + """Helper to create a resolved Future with a given value.""" + mock_future = asyncio.Future() + mock_future.set_result(to_return) + return mock_future + + +# --------------------------------------------------------------------------- +# Lightweight stubs for caldera framework objects that are not available +# when running the emu plugin tests in isolation. +# --------------------------------------------------------------------------- + +class _StubBaseWorld: + """Minimal stand-in for app.utility.base_world.BaseWorld.""" + + class Access: + RED = 'red' + BLUE = 'blue' + + _configs = {} + + @classmethod + def apply_config(cls, name, config): + cls._configs[name] = config + + @classmethod + def strip_yml(cls, path): + if os.path.exists(path): + with open(path, 'r') as fh: + return list(yaml.safe_load_all(fh)) + return [{}] + + @classmethod + def get_config(cls, name='main', prop=None): + cfg = cls._configs.get(name, {}) + if prop: + return cfg.get(prop) + return cfg + + @staticmethod + def create_logger(name): + import logging + return logging.getLogger(name) + + +class _StubBaseService(_StubBaseWorld): + """Minimal stand-in for app.utility.base_service.BaseService.""" + _services = {} + + @classmethod + def add_service(cls, name, svc): + cls._services[name] = svc + import logging + return logging.getLogger(name) + + @classmethod + def get_service(cls, name): + return cls._services.get(name) + + +class _StubBaseParser: + """Minimal stand-in for app.utility.base_parser.BaseParser.""" + + def __init__(self): + self.mappers = [] + self.used_facts = [] + + def set_value(self, key, value, used_facts): + return value + + +class _StubFact: + """Minimal stand-in for app.objects.secondclass.c_fact.Fact.""" + + def __init__(self, trait=None, value=None): + self.trait = trait + self.value = value + + def __eq__(self, other): + return isinstance(other, _StubFact) and self.trait == other.trait and self.value == other.value + + def __repr__(self): + return f'Fact(trait={self.trait!r}, value={self.value!r})' + + +class _StubRelationship: + """Minimal stand-in for app.objects.secondclass.c_relationship.Relationship.""" + + def __init__(self, source=None, edge=None, target=None): + self.source = source + self.edge = edge + self.target = target + + +class _StubLink: + """Minimal stand-in for app.objects.secondclass.c_link.Link.""" + + def __init__(self, command='', paw='', ability=None, **kwargs): + self.command = command + self.paw = paw + self.ability = ability + self.used = kwargs.get('used', []) + self.id = kwargs.get('id', '') + + +class _StubBaseRequirement: + """Minimal stand-in for plugins.stockpile.app.requirements.base_requirement.BaseRequirement.""" + pass + + +# --------------------------------------------------------------------------- +# Patch caldera imports before any plugin code is imported +# --------------------------------------------------------------------------- + +import sys + +# Build module stubs +_base_world_mod = type(sys)('app.utility.base_world') +_base_world_mod.BaseWorld = _StubBaseWorld +_base_service_mod = type(sys)('app.utility.base_service') +_base_service_mod.BaseService = _StubBaseService +_base_parser_mod = type(sys)('app.utility.base_parser') +_base_parser_mod.BaseParser = _StubBaseParser +_fact_mod = type(sys)('app.objects.secondclass.c_fact') +_fact_mod.Fact = _StubFact +_rel_mod = type(sys)('app.objects.secondclass.c_relationship') +_rel_mod.Relationship = _StubRelationship +_link_mod = type(sys)('app.objects.secondclass.c_link') +_link_mod.Link = _StubLink +_auth_svc_mod = type(sys)('app.service.auth_svc') +_auth_svc_mod.for_all_public_methods = lambda func: lambda cls: cls +_auth_svc_mod.check_authorization = lambda func: func +_base_req_mod = type(sys)('plugins.stockpile.app.requirements.base_requirement') +_base_req_mod.BaseRequirement = _StubBaseRequirement + +# Register in sys.modules (only if not already present — CI may have real caldera) +_stubs = { + 'app': type(sys)('app'), + 'app.utility': type(sys)('app.utility'), + 'app.utility.base_world': _base_world_mod, + 'app.utility.base_service': _base_service_mod, + 'app.utility.base_parser': _base_parser_mod, + 'app.objects': type(sys)('app.objects'), + 'app.objects.secondclass': type(sys)('app.objects.secondclass'), + 'app.objects.secondclass.c_fact': _fact_mod, + 'app.objects.secondclass.c_relationship': _rel_mod, + 'app.objects.secondclass.c_link': _link_mod, + 'app.service': type(sys)('app.service'), + 'app.service.auth_svc': _auth_svc_mod, + 'plugins': type(sys)('plugins'), + 'plugins.stockpile': type(sys)('plugins.stockpile'), + 'plugins.stockpile.app': type(sys)('plugins.stockpile.app'), + 'plugins.stockpile.app.requirements': type(sys)('plugins.stockpile.app.requirements'), + 'plugins.stockpile.app.requirements.base_requirement': _base_req_mod, +} + +for mod_name, mod_obj in _stubs.items(): + sys.modules.setdefault(mod_name, mod_obj) + +# Ensure the plugin package itself is importable from repo root +_repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) + +# Also make the plugin available as plugins.emu +_plugins_emu_mod = type(sys)('plugins.emu') +_plugins_emu_mod.__path__ = [_repo_root] +sys.modules.setdefault('plugins.emu', _plugins_emu_mod) + +_plugins_emu_app_mod = type(sys)('plugins.emu.app') +_plugins_emu_app_mod.__path__ = [os.path.join(_repo_root, 'app')] +sys.modules.setdefault('plugins.emu.app', _plugins_emu_app_mod) + + +# --------------------------------------------------------------------------- +# Shared fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture +def stub_fact_class(): + """Return the stub Fact class for use in tests.""" + return _StubFact + + +@pytest.fixture +def stub_link_class(): + """Return the stub Link class for use in tests.""" + return _StubLink + + +@pytest.fixture +def mock_app_svc(): + """Mock application service with a router.""" + svc = MagicMock() + svc.application = MagicMock() + svc.application.router = MagicMock() + svc.application.router.add_route = MagicMock() + return svc + + +@pytest.fixture +def mock_contact_svc(): + """Mock contact service.""" + svc = MagicMock() + svc.handle_heartbeat = AsyncMock() + return svc + + +@pytest.fixture +def tmp_data_dir(tmp_path): + """Create a temporary data directory structure.""" + data = tmp_path / 'data' + data.mkdir() + (data / 'abilities').mkdir() + (data / 'adversaries').mkdir() + (data / 'sources').mkdir() + (data / 'planners').mkdir() + payloads = tmp_path / 'payloads' + payloads.mkdir() + return tmp_path diff --git a/tests/test_emu_gui.py b/tests/test_emu_gui.py new file mode 100644 index 0000000..33c40d6 --- /dev/null +++ b/tests/test_emu_gui.py @@ -0,0 +1,53 @@ +"""Tests for app/emu_gui.py — EmuGUI.""" +import logging +import pytest +from unittest.mock import MagicMock, AsyncMock, patch + + +class TestEmuGUI: + """Test the EmuGUI class construction and splash handler.""" + + def _make_gui(self): + from plugins.emu.app.emu_gui import EmuGUI + services = { + 'auth_svc': MagicMock(), + 'data_svc': MagicMock(), + } + gui = EmuGUI(services, name='Emu', description='Test description') + return gui, services + + def test_construction(self): + gui, services = self._make_gui() + assert gui.name == 'Emu' + assert gui.description == 'Test description' + assert gui.auth_svc is services['auth_svc'] + assert gui.data_svc is services['data_svc'] + + def test_logger(self): + gui, _ = self._make_gui() + assert gui.log.name == 'emu_gui' + + def test_name_and_description(self): + from plugins.emu.app.emu_gui import EmuGUI + services = {'auth_svc': MagicMock(), 'data_svc': MagicMock()} + gui = EmuGUI(services, name='Custom', description='Custom desc') + assert gui.name == 'Custom' + assert gui.description == 'Custom desc' + + def test_missing_services(self): + from plugins.emu.app.emu_gui import EmuGUI + services = {} + gui = EmuGUI(services, name='Emu', description='desc') + assert gui.auth_svc is None + assert gui.data_svc is None + + def test_splash_is_callable(self): + gui, _ = self._make_gui() + assert callable(gui.splash) + + def test_splash_is_coroutine_function(self): + """The splash method (possibly wrapped by @template) should be async-compatible.""" + import asyncio + gui, _ = self._make_gui() + # The underlying function or its wrapper should be a coroutine function + assert asyncio.iscoroutinefunction(gui.splash) or callable(gui.splash) diff --git a/tests/test_emu_svc.py b/tests/test_emu_svc.py index d02bb2c..0c8b676 100644 --- a/tests/test_emu_svc.py +++ b/tests/test_emu_svc.py @@ -1,26 +1,64 @@ +"""Exhaustive tests for app/emu_svc.py — EmuService.""" import glob -import yaml +import json +import os import shutil +import uuid +import yaml import asyncio import pytest from pathlib import Path, PosixPath -from unittest.mock import patch, call +from unittest.mock import patch, call, MagicMock, AsyncMock, mock_open from app.utility.base_world import BaseWorld -from plugins.emu.app.emu_svc import EmuService +from app.utility.base_service import BaseService + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- def async_mock_return(to_return): - mock_future = asyncio.Future() - mock_future.set_result(to_return) - return mock_future + fut = asyncio.Future() + fut.set_result(to_return) + return fut + + +def _make_emu_svc(mock_app_svc=None, mock_contact_svc=None): + """Create an EmuService with stubbed caldera services.""" + if mock_app_svc is None: + mock_app_svc = MagicMock() + mock_app_svc.application = MagicMock() + mock_app_svc.application.router = MagicMock() + if mock_contact_svc is None: + mock_contact_svc = MagicMock() + mock_contact_svc.handle_heartbeat = AsyncMock() + + BaseService._services['app_svc'] = mock_app_svc + BaseService._services['contact_svc'] = mock_contact_svc + + # Provide minimal config + conf_dir = os.path.join('plugins', 'emu', 'conf') + os.makedirs(conf_dir, exist_ok=True) + conf_path = os.path.join(conf_dir, 'default.yml') + if not os.path.exists(conf_path): + with open(conf_path, 'w') as f: + yaml.dump({'evals_c2_host': '127.0.0.1', 'evals_c2_port': 8888}, f) + + from plugins.emu.app.emu_svc import EmuService + svc = EmuService() + return svc + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- @pytest.fixture def emu_svc(): - return EmuService() + return _make_emu_svc() @pytest.fixture @@ -141,52 +179,242 @@ def sample_emu_plan(): ''') +@pytest.fixture +def sample_emu_plan_bad_version(): + return yaml.safe_load(''' +- emulation_plan_details: + id: planid999 + adversary_name: Bad Adversary + adversary_description: Bad desc + attack_version: 8 + format_version: 0.5 +''') + + +@pytest.fixture +def sample_emu_plan_no_adversary(): + return yaml.safe_load(''' +- emulation_plan_details: + id: planid888 + attack_version: 8 + format_version: 1.0 +''') + + +@pytest.fixture +def ability_with_elevation(): + return { + 'id': 'elev-1', + 'name': 'Elevated Ability', + 'description': 'Needs elevation', + 'tactic': 'Privilege Escalation', + 'technique': {'attack_id': 'T1234', 'name': 'Elevate'}, + 'platforms': {'windows': {'psh': {'command': 'whoami /priv'}}}, + 'executors': [{'elevation_required': True, 'name': 'psh'}], + } + + +@pytest.fixture +def ability_without_elevation(): + return { + 'id': 'noelev-1', + 'name': 'Normal Ability', + 'description': 'No elevation', + 'tactic': 'Discovery', + 'technique': {'attack_id': 'T5678', 'name': 'Discover'}, + 'platforms': {'linux': {'sh': {'command': 'id'}}}, + 'executors': [{'name': 'sh'}], + } + + +# --------------------------------------------------------------------------- +# TestEmuSvc — original tests preserved + many new ones +# --------------------------------------------------------------------------- + class TestEmuSvc: LIBRARY_GLOB_PATH = 'plugins/emu/data/adversary-emulation-plans/*' PLANNER_GLOB_PATH = 'plugins/emu/data/adversary-emulation-plans/*/Emulation_Plan/yaml/planners/*.yml' PLANNER_PATH = 'plugins/emu/data/adversary-emulation-plans/library1/Emulation_Plan/yaml/planners/test_planner.yml' DEST_PLANNER_PATH = 'plugins/emu/data/planners/testid.yml' + # -- creation -- + def test_general_svc_creation(self, emu_svc): assert emu_svc.emu_dir == 'plugins/emu' assert emu_svc.repo_dir == 'plugins/emu/data/adversary-emulation-plans' assert emu_svc.data_dir == 'plugins/emu/data' assert emu_svc.payloads_dir == 'plugins/emu/payloads' - async def test_ingest_planner(self, emu_svc, planner_yaml): + def test_required_payloads_initially_empty(self, emu_svc): + assert emu_svc.required_payloads == set() + + def test_dynamically_compiled_payloads_class_attr(self): + from plugins.emu.app.emu_svc import EmuService + assert 'sandcat.go-linux' in EmuService._dynamicically_compiled_payloads + assert 'sandcat.go-darwin' in EmuService._dynamicically_compiled_payloads + assert 'sandcat.go-windows' in EmuService._dynamicically_compiled_payloads + assert len(EmuService._dynamicically_compiled_payloads) == 3 - with patch.object(BaseWorld, 'strip_yml', return_value=planner_yaml) as new_strip_yml: + # -- planner ingestion -- + + async def test_ingest_planner(self, emu_svc, planner_yaml): + with patch.object(BaseWorld, 'strip_yml', return_value=planner_yaml): with patch.object(shutil, 'copyfile', return_value=None) as new_copyfile: - await emu_svc._ingest_planner(TestEmuSvc.PLANNER_PATH) - new_strip_yml.assert_called_once_with(TestEmuSvc.PLANNER_PATH) + num_planners, num_ingested, num_errors = await emu_svc._ingest_planner(TestEmuSvc.PLANNER_PATH) new_copyfile.assert_called_once_with(TestEmuSvc.PLANNER_PATH, TestEmuSvc.DEST_PLANNER_PATH) + assert num_planners == 1 + assert num_ingested == 1 + assert num_errors == 0 async def test_ingest_bad_planner(self, emu_svc, not_planner_yaml): - with patch.object(BaseWorld, 'strip_yml', return_value=not_planner_yaml) as new_strip_yml: + with patch.object(BaseWorld, 'strip_yml', return_value=not_planner_yaml): with patch.object(shutil, 'copyfile', return_value=None) as new_copyfile: - await emu_svc._ingest_planner(TestEmuSvc.PLANNER_PATH) - new_strip_yml.assert_called_once_with(TestEmuSvc.PLANNER_PATH) + num_planners, num_ingested, num_errors = await emu_svc._ingest_planner(TestEmuSvc.PLANNER_PATH) new_copyfile.assert_not_called() + assert num_planners == 0 + assert num_ingested == 0 + assert num_errors == 1 + + async def test_ingest_planner_yaml_parse_error(self, emu_svc): + with patch.object(BaseWorld, 'strip_yml', side_effect=Exception('parse error')): + num_planners, num_ingested, num_errors = await emu_svc._ingest_planner(TestEmuSvc.PLANNER_PATH) + assert num_planners == 0 + assert num_ingested == 0 + assert num_errors == 1 + + async def test_ingest_planner_copy_io_error(self, emu_svc, planner_yaml): + with patch.object(BaseWorld, 'strip_yml', return_value=planner_yaml): + with patch.object(shutil, 'copyfile', side_effect=IOError('disk full')): + num_planners, num_ingested, num_errors = await emu_svc._ingest_planner(TestEmuSvc.PLANNER_PATH) + assert num_planners == 1 + assert num_ingested == 0 + assert num_errors == 1 async def test_load_planners(self, emu_svc, planner_yaml): - with patch.object(BaseWorld, 'strip_yml', return_value=planner_yaml) as new_strip_yml: + with patch.object(BaseWorld, 'strip_yml', return_value=planner_yaml): with patch.object(shutil, 'copyfile', return_value=None) as new_copyfile: - with patch.object(glob, 'iglob', return_value=[TestEmuSvc.PLANNER_PATH]) as new_iglob: + with patch.object(glob, 'iglob', return_value=[TestEmuSvc.PLANNER_PATH]): await emu_svc._load_planners(TestEmuSvc.LIBRARY_GLOB_PATH) - new_iglob.assert_called_once_with(TestEmuSvc.PLANNER_GLOB_PATH) - new_strip_yml.assert_called_once_with(TestEmuSvc.PLANNER_PATH) new_copyfile.assert_called_once_with(TestEmuSvc.PLANNER_PATH, TestEmuSvc.DEST_PLANNER_PATH) async def test_load_bad_planners(self, emu_svc, not_planner_yaml): - with patch.object(BaseWorld, 'strip_yml', return_value=not_planner_yaml) as new_strip_yml: + with patch.object(BaseWorld, 'strip_yml', return_value=not_planner_yaml): with patch.object(shutil, 'copyfile', return_value=None) as new_copyfile: - with patch.object(glob, 'iglob', return_value=[TestEmuSvc.PLANNER_PATH]) as new_iglob: + with patch.object(glob, 'iglob', return_value=[TestEmuSvc.PLANNER_PATH]): await emu_svc._load_planners(TestEmuSvc.LIBRARY_GLOB_PATH) - new_iglob.assert_called_once_with(TestEmuSvc.PLANNER_GLOB_PATH) - new_strip_yml.assert_called_once_with(TestEmuSvc.PLANNER_PATH) new_copyfile.assert_not_called() + async def test_load_planners_empty_glob(self, emu_svc): + with patch.object(glob, 'iglob', return_value=[]): + await emu_svc._load_planners(TestEmuSvc.LIBRARY_GLOB_PATH) + # Should succeed with no errors + + # -- _is_planner -- + + def test_is_planner_true(self): + from plugins.emu.app.emu_svc import EmuService + assert EmuService._is_planner({'id': '1', 'module': 'mod', 'name': 'test'}) is True + + def test_is_planner_false_missing_id(self): + from plugins.emu.app.emu_svc import EmuService + assert EmuService._is_planner({'module': 'mod', 'name': 'test'}) is False + + def test_is_planner_false_missing_module(self): + from plugins.emu.app.emu_svc import EmuService + assert EmuService._is_planner({'id': '1', 'name': 'test'}) is False + + def test_is_planner_false_empty(self): + from plugins.emu.app.emu_svc import EmuService + assert EmuService._is_planner({}) is False + + # -- _is_valid_format_version -- + + def test_is_valid_format_version_1_0(self): + from plugins.emu.app.emu_svc import EmuService + assert EmuService._is_valid_format_version({'format_version': 1.0}) is True + + def test_is_valid_format_version_2_0(self): + from plugins.emu.app.emu_svc import EmuService + assert EmuService._is_valid_format_version({'format_version': 2.0}) is True + + def test_is_valid_format_version_below_1(self): + from plugins.emu.app.emu_svc import EmuService + assert EmuService._is_valid_format_version({'format_version': 0.5}) is False + + def test_is_valid_format_version_missing(self): + from plugins.emu.app.emu_svc import EmuService + assert EmuService._is_valid_format_version({}) is False + + def test_is_valid_format_version_non_numeric(self): + from plugins.emu.app.emu_svc import EmuService + assert EmuService._is_valid_format_version({'format_version': 'abc'}) is False + + # -- _is_ability -- + + async def test_is_ability_true(self): + from plugins.emu.app.emu_svc import EmuService + result = await EmuService._is_ability({'id': '1', 'platforms': {}, 'name': 'test'}) + assert result is True + + async def test_is_ability_false_missing_id(self): + from plugins.emu.app.emu_svc import EmuService + result = await EmuService._is_ability({'platforms': {}, 'name': 'test'}) + assert result is False + + async def test_is_ability_false_missing_platforms(self): + from plugins.emu.app.emu_svc import EmuService + result = await EmuService._is_ability({'id': '1', 'name': 'test'}) + assert result is False + + async def test_is_ability_false_empty(self): + from plugins.emu.app.emu_svc import EmuService + result = await EmuService._is_ability({}) + assert result is False + + # -- get_adversary_from_filename -- + + def test_get_adversary_from_filename_normal(self): + from plugins.emu.app.emu_svc import EmuService + assert EmuService.get_adversary_from_filename('/path/to/apt29.yaml') == 'apt29' + + def test_get_adversary_from_filename_nested(self): + from plugins.emu.app.emu_svc import EmuService + assert EmuService.get_adversary_from_filename('/a/b/c/fin7.yml') == 'fin7' + + def test_get_adversary_from_filename_no_extension(self): + from plugins.emu.app.emu_svc import EmuService + assert EmuService.get_adversary_from_filename('/path/adversary') == 'adversary' + + def test_get_adversary_from_filename_just_filename(self): + from plugins.emu.app.emu_svc import EmuService + assert EmuService.get_adversary_from_filename('test.yaml') == 'test' + + # -- get_privilege -- + + def test_get_privilege_elevated(self): + from plugins.emu.app.emu_svc import EmuService + result = EmuService.get_privilege([{'elevation_required': True, 'name': 'psh'}]) + assert result == 'Elevated' + + def test_get_privilege_not_elevated(self): + from plugins.emu.app.emu_svc import EmuService + result = EmuService.get_privilege([{'name': 'sh'}]) + assert result is False + + def test_get_privilege_empty_list(self): + from plugins.emu.app.emu_svc import EmuService + result = EmuService.get_privilege([]) + assert not result + + def test_get_privilege_none(self): + from plugins.emu.app.emu_svc import EmuService + result = EmuService.get_privilege(None) + assert result is False + + # -- abilities ingestion -- + async def test_ingest_abilities(self, emu_svc, sample_emu_plan): + from plugins.emu.app.emu_svc import EmuService with patch.object(EmuService, '_write_ability', return_value=async_mock_return(None)) as write_ability: abilities, facts, at_total, at_ingested, errors = await emu_svc._ingest_abilities(sample_emu_plan) assert write_ability.call_count == 3 @@ -200,27 +428,115 @@ async def test_ingest_abilities(self, emu_svc, sample_emu_plan): assert at_ingested == 3 assert errors == 0 assert {'payload1A', 'payload1B', 'payload2A'} == emu_svc.required_payloads - write_ability.assert_has_calls([ - call(dict( - id='1-2-3', name='Ability 1', description='Desc 1', tactic='tactic1', - technique=dict(name='technique 1', attack_id='technique1'), repeatable=False, requirements=[], - platforms=dict(linux=dict(sh=dict(command='test command', payloads=['payload1A', 'payload1B']))), - )), - call(dict( - id='2-3-4', name='Ability 2', description='Desc 2', tactic='tactic2', - technique=dict(name='technique 2', attack_id='technique2'), repeatable=False, requirements=[], - platforms=dict( - linux=dict(sh=dict(command='test command', payloads=['payload2A'])), - windows=dict(cmd=dict(command='test command')), - ), - )), - call(dict( - id='3-4-5', name='Ability 3', description='Desc 3', tactic='tactic3', - technique=dict(name='technique 3', attack_id='technique3'), repeatable=False, requirements=[], - platforms=dict(linux=dict(sh=dict(command='test command', uploads=['upload1A', 'upload1B'], - cleanup='cleanup command'))), - )), - ]) + + async def test_ingest_abilities_empty_plan(self, emu_svc): + from plugins.emu.app.emu_svc import EmuService + with patch.object(EmuService, '_write_ability', return_value=async_mock_return(None)) as write_ability: + abilities, facts, at_total, at_ingested, errors = await emu_svc._ingest_abilities([]) + assert abilities == [] + assert facts == [] + assert at_total == 0 + assert at_ingested == 0 + assert errors == 0 + write_ability.assert_not_called() + + async def test_ingest_abilities_write_error(self, emu_svc, sample_emu_plan): + from plugins.emu.app.emu_svc import EmuService + with patch.object(EmuService, '_write_ability', side_effect=Exception('write error')): + abilities, facts, at_total, at_ingested, errors = await emu_svc._ingest_abilities(sample_emu_plan) + assert at_total == 3 + assert at_ingested == 0 + assert errors == 3 + + # -- emulation plan ingestion -- + + async def test_ingest_emulation_plan(self, emu_svc, sample_emu_plan): + from plugins.emu.app.emu_svc import EmuService + with patch.object(BaseWorld, 'strip_yml', return_value=[sample_emu_plan]): + with patch.object(EmuService, '_write_ability', return_value=async_mock_return(None)): + with patch.object(EmuService, '_write_adversary', return_value=async_mock_return(None)) as w_adv: + with patch.object(EmuService, '_write_source', return_value=async_mock_return(None)) as w_src: + at_total, at_ingested, errors = await emu_svc._ingest_emulation_plan('test.yaml') + assert at_total == 3 + assert at_ingested == 3 + assert errors == 0 + w_adv.assert_called_once() + w_src.assert_called_once() + + async def test_ingest_emulation_plan_bad_version(self, emu_svc, sample_emu_plan_bad_version): + with patch.object(BaseWorld, 'strip_yml', return_value=[sample_emu_plan_bad_version]): + at_total, at_ingested, errors = await emu_svc._ingest_emulation_plan('test.yaml') + assert at_total == 0 + assert at_ingested == 0 + assert errors == 1 + + async def test_ingest_emulation_plan_no_adversary(self, emu_svc, sample_emu_plan_no_adversary): + with patch.object(BaseWorld, 'strip_yml', return_value=[sample_emu_plan_no_adversary]): + at_total, at_ingested, errors = await emu_svc._ingest_emulation_plan('test.yaml') + assert at_total == 0 + assert at_ingested == 0 + assert errors == 1 + + # -- _save_ability with elevation -- + + async def test_save_ability_with_elevation(self, emu_svc, ability_with_elevation): + from plugins.emu.app.emu_svc import EmuService + with patch.object(EmuService, '_write_ability', return_value=async_mock_return(None)): + ability_id, facts = await emu_svc._save_ability(ability_with_elevation) + assert ability_id == 'elev-1' + + async def test_save_ability_without_elevation(self, emu_svc, ability_without_elevation): + from plugins.emu.app.emu_svc import EmuService + with patch.object(EmuService, '_write_ability', return_value=async_mock_return(None)): + ability_id, facts = await emu_svc._save_ability(ability_without_elevation) + assert ability_id == 'noelev-1' + + # -- _unique_facts -- + + async def test_unique_facts_no_duplicates(self): + from plugins.emu.app.emu_svc import EmuService + facts = [{'trait': 'a', 'value': '1'}, {'trait': 'b', 'value': '2'}] + result = await EmuService._unique_facts(facts) + assert result == facts + + async def test_unique_facts_with_duplicates(self): + from plugins.emu.app.emu_svc import EmuService + facts = [ + {'trait': 'a', 'value': '1'}, + {'trait': 'b', 'value': '2'}, + {'trait': 'a', 'value': '1'}, + ] + result = await EmuService._unique_facts(facts) + assert len(result) == 2 + assert result == [{'trait': 'a', 'value': '1'}, {'trait': 'b', 'value': '2'}] + + async def test_unique_facts_empty(self): + from plugins.emu.app.emu_svc import EmuService + result = await EmuService._unique_facts([]) + assert result == [] + + # -- _register_required_payloads -- + + def test_register_required_payloads(self, emu_svc): + payloads = ['payload1', 'payload2', 'payload3', 'sandcat.go-darwin', 'sandcat.go-linux', 'sandcat.go-windows'] + want = {'payload1', 'payload2', 'payload3'} + emu_svc._register_required_payloads(payloads) + assert emu_svc.required_payloads == want + + def test_register_required_payloads_empty(self, emu_svc): + emu_svc._register_required_payloads([]) + assert emu_svc.required_payloads == set() + + def test_register_required_payloads_only_dynamic(self, emu_svc): + emu_svc._register_required_payloads(['sandcat.go-linux', 'sandcat.go-darwin']) + assert emu_svc.required_payloads == set() + + def test_register_required_payloads_accumulates(self, emu_svc): + emu_svc._register_required_payloads(['a', 'b']) + emu_svc._register_required_payloads(['b', 'c']) + assert emu_svc.required_payloads == {'a', 'b', 'c'} + + # -- _store_required_payloads -- def test_store_required_payloads(self, emu_svc): def _rglob(_, target): @@ -237,8 +553,236 @@ def _rglob(_, target): call(PosixPath('/path/to/payload3'), 'plugins/emu/payloads/payload3'), ], any_order=True) - def test_register_required_payloads(self, emu_svc): - payloads = ['payload1', 'payload2', 'payload3', 'sandcat.go-darwin', 'sandcat.go-linux', 'sandcat.go-windows'] - want = {'payload1', 'payload2', 'payload3'} - emu_svc._register_required_payloads(payloads) - assert emu_svc.required_payloads == want + def test_store_required_payloads_empty(self, emu_svc): + emu_svc.required_payloads = set() + emu_svc._store_required_payloads() + # No error, nothing to copy + + def test_store_required_payloads_not_found(self, emu_svc): + emu_svc.required_payloads = {'missing_payload'} + with patch.object(Path, 'rglob', return_value=[]): + emu_svc._store_required_payloads() + # Should log a warning but not raise + + def test_store_required_payloads_copy_failure(self, emu_svc): + def _rglob(_, target): + return [PosixPath('/path/to/' + target)] + + emu_svc.required_payloads = {'bad_payload'} + with patch.object(Path, 'rglob', new=_rglob): + with patch.object(shutil, 'copyfile', side_effect=Exception('copy failed')): + emu_svc._store_required_payloads() + # Should log but not raise + + def test_store_required_payloads_already_exists(self, emu_svc): + emu_svc.required_payloads = {'existing_payload'} + with patch.object(os.path, 'exists', return_value=True): + emu_svc._store_required_payloads() + # Payload already in payloads dir — skip + + # -- _copy_planner -- + + def test_copy_planner_new(self, emu_svc, tmp_path): + src = tmp_path / 'source.yml' + src.write_text('test') + target = 'target.yml' + with patch.object(os.path, 'exists', side_effect=lambda p: p == str(src)): + with patch.object(os, 'makedirs'): + with patch.object(shutil, 'copyfile') as cp: + emu_svc._copy_planner(str(src), target) + cp.assert_called_once() + + def test_copy_planner_already_exists(self, emu_svc, tmp_path): + with patch.object(os.path, 'exists', return_value=True): + with patch.object(shutil, 'copyfile') as cp: + emu_svc._copy_planner('src.yml', 'target.yml') + cp.assert_not_called() + + # -- clone_repo -- + + async def test_clone_repo_default_url(self, emu_svc): + with patch.object(os.path, 'exists', return_value=False): + with patch('plugins.emu.app.emu_svc.check_call') as mock_check: + await emu_svc.clone_repo() + mock_check.assert_called_once() + args = mock_check.call_args[0][0] + assert 'git' in args[0] + assert 'clone' in args + assert 'adversary_emulation_library' in args[4] + + async def test_clone_repo_custom_url(self, emu_svc): + with patch.object(os.path, 'exists', return_value=False): + with patch('plugins.emu.app.emu_svc.check_call') as mock_check: + await emu_svc.clone_repo(repo_url='https://example.com/fork.git') + args = mock_check.call_args[0][0] + assert args[4] == 'https://example.com/fork.git' + + async def test_clone_repo_already_exists(self, emu_svc): + with patch.object(os.path, 'exists', return_value=True): + with patch.object(os, 'listdir', return_value=['file1']): + with patch('plugins.emu.app.emu_svc.check_call') as mock_check: + await emu_svc.clone_repo() + mock_check.assert_not_called() + + # -- populate_data_directory -- + + async def test_populate_data_directory_default(self, emu_svc): + from plugins.emu.app.emu_svc import EmuService + with patch.object(EmuService, '_load_adversaries_and_abilities', return_value=async_mock_return(None)) as load_adv: + with patch.object(EmuService, '_load_planners', return_value=async_mock_return(None)) as load_plan: + await emu_svc.populate_data_directory() + expected = os.path.join(emu_svc.repo_dir, '*') + load_adv.assert_called_once_with(expected) + load_plan.assert_called_once_with(expected) + + async def test_populate_data_directory_custom_path(self, emu_svc): + from plugins.emu.app.emu_svc import EmuService + with patch.object(EmuService, '_load_adversaries_and_abilities', return_value=async_mock_return(None)) as load_adv: + with patch.object(EmuService, '_load_planners', return_value=async_mock_return(None)) as load_plan: + await emu_svc.populate_data_directory(library_path='/custom/path') + load_adv.assert_called_once_with('/custom/path') + load_plan.assert_called_once_with('/custom/path') + + # -- handle_forwarded_beacon -- + + async def test_handle_forwarded_beacon_full_profile(self, emu_svc, mock_contact_svc): + from plugins.emu.app.emu_svc import EmuService + emu_svc.contact_svc = mock_contact_svc + profile_data = { + 'guid': 'test-paw-123', + 'platform': 'windows', + 'hostName': 'WORKSTATION1', + 'user': 'admin', + 'pid': 1234, + 'ppid': 5678, + } + request = MagicMock() + request.read = AsyncMock(return_value=json.dumps(profile_data).encode()) + + response = await emu_svc.handle_forwarded_beacon(request) + assert 'test-paw-123' in response.text + mock_contact_svc.handle_heartbeat.assert_called_once() + call_kwargs = mock_contact_svc.handle_heartbeat.call_args[1] + assert call_kwargs['paw'] == 'test-paw-123' + assert call_kwargs['platform'] == 'windows' + assert call_kwargs['host'] == 'WORKSTATION1' + assert call_kwargs['username'] == 'admin' + assert call_kwargs['pid'] == 1234 + assert call_kwargs['ppid'] == 5678 + + async def test_handle_forwarded_beacon_minimal_profile(self, emu_svc, mock_contact_svc): + emu_svc.contact_svc = mock_contact_svc + profile_data = {'guid': 'min-paw'} + request = MagicMock() + request.read = AsyncMock(return_value=json.dumps(profile_data).encode()) + + response = await emu_svc.handle_forwarded_beacon(request) + assert 'min-paw' in response.text + call_kwargs = mock_contact_svc.handle_heartbeat.call_args[1] + assert call_kwargs['platform'] == 'evals' + + async def test_handle_forwarded_beacon_error(self, emu_svc): + request = MagicMock() + request.read = AsyncMock(side_effect=Exception('bad request')) + with pytest.raises((TypeError, Exception)): + await emu_svc.handle_forwarded_beacon(request) + + # -- _write_ability / _write_adversary / _write_source with tmp dirs -- + + async def test_write_ability(self, emu_svc, tmp_data_dir): + emu_svc.data_dir = str(tmp_data_dir / 'data') + ability_data = { + 'id': 'test-ab-1', + 'tactic': 'discovery', + 'name': 'test', + 'platforms': {}, + } + await emu_svc._write_ability(ability_data) + path = os.path.join(emu_svc.data_dir, 'abilities', 'discovery', 'test-ab-1.yml') + assert os.path.exists(path) + + async def test_write_ability_skip_existing(self, emu_svc, tmp_data_dir): + emu_svc.data_dir = str(tmp_data_dir / 'data') + ability_data = {'id': 'test-ab-2', 'tactic': 'discovery', 'name': 'test', 'platforms': {}} + await emu_svc._write_ability(ability_data) + # Write again — should skip + await emu_svc._write_ability(ability_data) + # No error + + async def test_write_adversary(self, emu_svc, tmp_data_dir): + emu_svc.data_dir = str(tmp_data_dir / 'data') + adv_data = {'id': 'test-adv-1', 'name': 'Test Adv', 'description': 'desc', 'atomic_ordering': []} + await emu_svc._write_adversary(adv_data) + path = os.path.join(emu_svc.data_dir, 'adversaries', 'test-adv-1.yml') + assert os.path.exists(path) + + async def test_write_adversary_skip_existing(self, emu_svc, tmp_data_dir): + emu_svc.data_dir = str(tmp_data_dir / 'data') + adv_data = {'id': 'test-adv-dup', 'name': 'Dup', 'description': 'desc', 'atomic_ordering': []} + await emu_svc._write_adversary(adv_data) + await emu_svc._write_adversary(adv_data) + # No error on duplicate + + async def test_write_source(self, emu_svc, tmp_data_dir): + emu_svc.data_dir = str(tmp_data_dir / 'data') + source_data = {'id': 'test-src-1', 'name': 'Test Source', 'facts': []} + await emu_svc._write_source(source_data) + path = os.path.join(emu_svc.data_dir, 'sources', 'test-src-1.yml') + assert os.path.exists(path) + + # -- _save_adversary -- + + async def test_save_adversary(self, emu_svc, tmp_data_dir): + from plugins.emu.app.emu_svc import EmuService + emu_svc.data_dir = str(tmp_data_dir / 'data') + await emu_svc._save_adversary(id='adv-1', name='APT1', description='Test', abilities=['a', 'b']) + path = os.path.join(emu_svc.data_dir, 'adversaries', 'adv-1.yml') + assert os.path.exists(path) + with open(path) as f: + data = yaml.safe_load(f) + assert data['name'] == 'APT1' + assert data['description'] == 'Test (Emu)' + assert data['atomic_ordering'] == ['a', 'b'] + + # -- _save_source -- + + async def test_save_source(self, emu_svc, tmp_data_dir): + emu_svc.data_dir = str(tmp_data_dir / 'data') + facts = [{'trait': 'a', 'value': '1'}, {'trait': 'b', 'value': '2'}] + await emu_svc._save_source('TestAdv', facts) + src_dir = os.path.join(emu_svc.data_dir, 'sources') + files = os.listdir(src_dir) + assert len(files) == 1 + with open(os.path.join(src_dir, files[0])) as f: + data = yaml.safe_load(f) + assert data['name'] == 'TestAdv (Emu)' + assert len(data['facts']) == 2 + + # -- _load_object -- + + async def test_load_object_multiple_files(self, emu_svc): + call_count = 0 + + async def mock_ingest(filename): + nonlocal call_count + call_count += 1 + return 1, 1, 0 + + with patch.object(glob, 'iglob', return_value=['f1.yaml', 'f2.yaml', 'f3.yaml']): + await emu_svc._load_object('*.yaml', 'test_objects', mock_ingest) + assert call_count == 3 + + async def test_load_object_with_errors(self, emu_svc): + async def mock_ingest(filename): + return 1, 0, 1 + + with patch.object(glob, 'iglob', return_value=['f1.yaml']): + await emu_svc._load_object('*.yaml', 'test_objects', mock_ingest) + # Should not raise; errors are counted + + # -- decrypt_payloads -- + + async def test_decrypt_payloads_no_scripts(self, emu_svc): + with patch.object(glob, 'iglob', return_value=[]): + await emu_svc.decrypt_payloads() + # No scripts found — nothing to do diff --git a/tests/test_group_filtered_planner.py b/tests/test_group_filtered_planner.py index 4b3aa3a..7c1fc0b 100644 --- a/tests/test_group_filtered_planner.py +++ b/tests/test_group_filtered_planner.py @@ -1,12 +1,17 @@ +"""Exhaustive tests for app/group_filtered_planner.py — LogicalPlanner.""" import pytest -from app.objects.secondclass.c_link import Link from plugins.emu.app.group_filtered_planner import LogicalPlanner +from tests.conftest import _StubLink BUCKET_NAME = 'fetch_and_run_links' +# --------------------------------------------------------------------------- +# Helpers / dummies +# --------------------------------------------------------------------------- + class DummyOperation: def __init__(self, dummy_adversary, dummy_agents): self.adversary = dummy_adversary @@ -35,6 +40,14 @@ def __init__(self, ability_id): self.ability_id = ability_id +def _make_link(paw, ability_id, command='test command'): + return _StubLink(command=command, paw=paw, ability=DummyAbility(ability_id)) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + @pytest.fixture def dummy_agents(): return [ @@ -49,10 +62,10 @@ def dummy_agents(): @pytest.fixture def pending_links(): return [ - Link(command='test command', paw='paw1', ability=DummyAbility('123')), - Link(command='test command variant', paw='paw1', ability=DummyAbility('123')), - Link(command='test command', paw='paw2', ability=DummyAbility('123')), - Link(command='test command', paw='paw3', ability=DummyAbility('123')), + _make_link('paw1', '123'), + _make_link('paw1', '123', command='test command variant'), + _make_link('paw2', '123'), + _make_link('paw3', '123'), ] @@ -60,36 +73,23 @@ def pending_links(): def potential_links_dict(): return { 'paw1': [ - Link(command='test command', paw='paw1', ability=DummyAbility('123')), - Link(command='test command variant', paw='paw1', ability=DummyAbility('123')), - Link(command='test command', paw='paw1', ability=DummyAbility('456')), - Link(command='test command', paw='paw1', ability=DummyAbility('1011')), + _make_link('paw1', '123'), + _make_link('paw1', '123', command='test command variant'), + _make_link('paw1', '456'), + _make_link('paw1', '1011'), ], 'paw2': [ - Link(command='test command', paw='paw2', ability=DummyAbility('123')), + _make_link('paw2', '123'), ], 'paw3': [ - Link(command='test command', paw='paw3', ability=DummyAbility('1011')), + _make_link('paw3', '1011'), ], 'paw4': [ - Link(command='test command', paw='paw4', ability=DummyAbility('456')), + _make_link('paw4', '456'), ], } -@pytest.fixture -def potential_links_list(): - return [ - Link(command='test command', paw='paw1', ability=DummyAbility('123')), - Link(command='test command variant', paw='paw1', ability=DummyAbility('123')), - Link(command='test command', paw='paw1', ability=DummyAbility('456')), - Link(command='test command', paw='paw2', ability=DummyAbility('123')), - Link(command='test command', paw='paw3', ability=DummyAbility('1011')), - Link(command='test command', paw='paw4', ability=DummyAbility('456')), - Link(command='test command', paw='paw1', ability=DummyAbility('1011')), - ] - - @pytest.fixture def sample_abilities(): return ['123', '456', '789', '1011'] @@ -127,191 +127,232 @@ def filtered_planner(generate_planner, dummy_agents, sample_abilities, sample_fi return generate_planner(sample_abilities, dummy_agents, filtered_groups_by_ability=sample_filter) -class TestGroupFilteredPlanner: - async def test_fetch_from_pending_links(self, planner_without_filter, pending_links): +# --------------------------------------------------------------------------- +# Tests — construction +# --------------------------------------------------------------------------- + +class TestLogicalPlannerInit: + def test_default_attributes(self): + op = DummyOperation(DummyAdversary([]), []) + planner = LogicalPlanner(op, planning_svc=None) + assert planner.state_machine == ['fetch_and_run_links'] + assert planner.next_bucket == 'fetch_and_run_links' + assert planner.filtered_groups_by_ability == {} + assert planner.pending_links == [] + assert planner.current_ability_index == 0 + assert planner.stopping_conditions == () + assert planner.stopping_condition_met is False + + def test_with_stopping_conditions(self): + op = DummyOperation(DummyAdversary([]), []) + planner = LogicalPlanner(op, None, stopping_conditions=('cond1',)) + assert planner.stopping_conditions == ('cond1',) + + def test_with_filtered_groups(self): + filt = {'abc': ['g1']} + op = DummyOperation(DummyAdversary([]), []) + planner = LogicalPlanner(op, None, filtered_groups_by_ability=filt) + assert planner.filtered_groups_by_ability == filt + + def test_none_filtered_groups_defaults_to_empty(self): + op = DummyOperation(DummyAdversary([]), []) + planner = LogicalPlanner(op, None, filtered_groups_by_ability=None) + assert planner.filtered_groups_by_ability == {} + + +# --------------------------------------------------------------------------- +# Tests — _fetch_from_pending_links +# --------------------------------------------------------------------------- + +class TestFetchFromPendingLinks: + async def test_basic(self, planner_without_filter, pending_links): planner_without_filter.pending_links = pending_links links_to_use = planner_without_filter._fetch_from_pending_links() assert len(links_to_use) == 3 assert len(planner_without_filter.pending_links) == 1 - assert links_to_use[0].paw == 'paw1' and links_to_use[0].ability.ability_id == '123' - assert links_to_use[0].command == 'test command' - assert links_to_use[1].paw == 'paw2' - assert links_to_use[2].paw == 'paw3' - assert planner_without_filter.pending_links[0].paw == 'paw1' - assert planner_without_filter.pending_links[0].command == 'test command variant' - - async def test_fetch_from_empty_pending_links(self, planner_without_filter): + paws_used = {l.paw for l in links_to_use} + assert paws_used == {'paw1', 'paw2', 'paw3'} + + async def test_empty(self, planner_without_filter): links_to_use = planner_without_filter._fetch_from_pending_links() assert not links_to_use assert not planner_without_filter.pending_links - async def test_get_valid_agents_for_abil_without_filter(self, planner_without_filter): - assert len(planner_without_filter.operation.agents) == 5 - valid_agents = planner_without_filter._get_valid_agents_for_ability('123') - assert len(valid_agents) == 5 + async def test_single_agent_multiple_links(self, planner_without_filter): + planner_without_filter.pending_links = [ + _make_link('paw1', '123', command='cmd1'), + _make_link('paw1', '123', command='cmd2'), + _make_link('paw1', '456', command='cmd3'), + ] + links = planner_without_filter._fetch_from_pending_links() + assert len(links) == 1 + assert links[0].command == 'cmd1' + assert len(planner_without_filter.pending_links) == 2 + + async def test_all_different_agents(self, planner_without_filter): + planner_without_filter.pending_links = [ + _make_link('p1', '123'), + _make_link('p2', '123'), + _make_link('p3', '456'), + ] + links = planner_without_filter._fetch_from_pending_links() + assert len(links) == 3 + assert len(planner_without_filter.pending_links) == 0 + - async def test_get_pending_links_without_filter(self, planner_without_filter): +# --------------------------------------------------------------------------- +# Tests — _get_valid_agents_for_ability +# --------------------------------------------------------------------------- + +class TestGetValidAgents: + async def test_no_filter(self, planner_without_filter): + valid = planner_without_filter._get_valid_agents_for_ability('123') + assert len(valid) == 5 + + async def test_with_filter_single_group(self, filtered_planner): + valid = filtered_planner._get_valid_agents_for_ability('123') + paws = [a.paw for a in valid] + assert set(paws) == {'paw1', 'paw4'} + + async def test_with_filter_multiple_groups(self, filtered_planner): + valid = filtered_planner._get_valid_agents_for_ability('1011') + paws = [a.paw for a in valid] + assert set(paws) == {'paw1', 'paw3', 'paw4'} + + async def test_ability_not_in_filter(self, filtered_planner): + """Ability 456 is not in filter — all agents should be valid.""" + valid = filtered_planner._get_valid_agents_for_ability('456') + assert len(valid) == 5 + + async def test_filter_with_no_matching_agents(self, generate_planner): + agents = [DummyAgent('p1', 'groupX')] + filt = {'abc': ['groupY']} + planner = generate_planner(['abc'], agents, filtered_groups_by_ability=filt) + valid = planner._get_valid_agents_for_ability('abc') + assert len(valid) == 0 + + +# --------------------------------------------------------------------------- +# Tests — _get_pending_links +# --------------------------------------------------------------------------- + +class TestGetPendingLinks: + async def test_without_filter(self, planner_without_filter): links = await planner_without_filter._get_pending_links('123') assert len(links) == 3 - assert links[0].paw == 'paw1' - assert links[0].ability.ability_id == '123' - assert links[0].command == 'test command' - assert links[1].paw == 'paw1' - assert links[1].ability.ability_id == '123' - assert links[1].command == 'test command variant' - assert links[2].paw == 'paw2' - assert links[2].ability.ability_id == '123' - assert links[2].command == 'test command' - - async def test_fetch_links_without_filter(self, planner_without_filter): - assert not planner_without_filter.pending_links - assert planner_without_filter.current_ability_index == 0 + assert all(l.ability.ability_id == '123' for l in links) - # first pass - links_to_use = await planner_without_filter._fetch_links() - assert len(planner_without_filter.pending_links) == 1 - assert planner_without_filter.pending_links[0].paw == 'paw1' - assert planner_without_filter.pending_links[0].command == 'test command variant' - assert planner_without_filter.pending_links[0].ability.ability_id == '123' - assert planner_without_filter.current_ability_index == 1 - assert len(links_to_use) == 2 - assert links_to_use[0].paw == 'paw1' - assert links_to_use[0].command == 'test command' - assert links_to_use[0].ability.ability_id == '123' - assert links_to_use[1].paw == 'paw2' - assert links_to_use[1].command == 'test command' - assert links_to_use[1].ability.ability_id == '123' - - # second pass - finishes links from first pass - links_to_use = await planner_without_filter._fetch_links() - assert len(planner_without_filter.pending_links) == 0 + async def test_with_filter(self, filtered_planner): + links = await filtered_planner._get_pending_links('123') + assert len(links) == 2 + assert all(l.ability.ability_id == '123' for l in links) + + async def test_no_matching_ability(self, planner_without_filter): + links = await planner_without_filter._get_pending_links('nonexistent') + assert len(links) == 0 + + +# --------------------------------------------------------------------------- +# Tests — _fetch_links full sequence +# --------------------------------------------------------------------------- + +class TestFetchLinksSequence: + async def test_without_filter_full(self, planner_without_filter): + # First pass: ability 123 + links = await planner_without_filter._fetch_links() + assert len(links) == 2 assert planner_without_filter.current_ability_index == 1 - assert len(links_to_use) == 1 - assert links_to_use[0].paw == 'paw1' - assert links_to_use[0].command == 'test command variant' - assert links_to_use[0].ability.ability_id == '123' - # third pass - ability #2 - links_to_use = await planner_without_filter._fetch_links() - assert len(planner_without_filter.pending_links) == 0 + # Second pass: remaining from 123 + links = await planner_without_filter._fetch_links() + assert len(links) == 1 + + # Third pass: ability 456 + links = await planner_without_filter._fetch_links() + assert len(links) == 2 assert planner_without_filter.current_ability_index == 2 - assert len(links_to_use) == 2 - assert links_to_use[0].paw == 'paw1' - assert links_to_use[0].command == 'test command' - assert links_to_use[0].ability.ability_id == '456' - assert links_to_use[1].paw == 'paw4' - assert links_to_use[1].command == 'test command' - assert links_to_use[1].ability.ability_id == '456' - - # fourth pass - skip ability #3 and go to #4 - links_to_use = await planner_without_filter._fetch_links() - assert len(planner_without_filter.pending_links) == 0 - assert planner_without_filter.current_ability_index == 4 - assert len(links_to_use) == 2 - assert links_to_use[0].paw == 'paw1' - assert links_to_use[0].command == 'test command' - assert links_to_use[0].ability.ability_id == '1011' - assert links_to_use[1].paw == 'paw3' - assert links_to_use[1].command == 'test command' - assert links_to_use[1].ability.ability_id == '1011' - - # fifth pass - end - links_to_use = await planner_without_filter._fetch_links() - assert len(planner_without_filter.pending_links) == 0 - assert planner_without_filter.current_ability_index == 4 - assert len(links_to_use) == 0 - - async def test_get_valid_agents_for_abil_with_filter(self, filtered_planner): - assert len(filtered_planner.operation.agents) == 5 - valid_agent_paws = [agent.paw for agent in filtered_planner._get_valid_agents_for_ability('123')] - assert len(valid_agent_paws) == 2 - assert 'paw1' in valid_agent_paws - assert 'paw4' in valid_agent_paws - valid_agent_paws = [agent.paw for agent in filtered_planner._get_valid_agents_for_ability('456')] - assert len(valid_agent_paws) == 5 - valid_agent_paws = [agent.paw for agent in filtered_planner._get_valid_agents_for_ability('789')] - assert len(valid_agent_paws) == 1 - assert 'paw2' in valid_agent_paws - - async def test_get_pending_links_with_filter(self, filtered_planner): - links = await filtered_planner._get_pending_links('123') + + # Fourth pass: skip 789 (no links), go to 1011 + links = await planner_without_filter._fetch_links() assert len(links) == 2 - assert links[0].paw == 'paw1' - assert links[0].ability.ability_id == '123' - assert links[0].command == 'test command' - assert links[1].paw == 'paw1' - assert links[1].ability.ability_id == '123' - assert links[1].command == 'test command variant' - - links = await filtered_planner._get_pending_links('456') + assert planner_without_filter.current_ability_index == 4 + + # Fifth pass: end + links = await planner_without_filter._fetch_links() + assert len(links) == 0 + + async def test_with_filter_full(self, filtered_planner): + # Ability 123 filtered to group1 only + links = await filtered_planner._fetch_links() + assert len(links) == 1 + assert filtered_planner.current_ability_index == 1 + + # Remaining from 123 + links = await filtered_planner._fetch_links() + assert len(links) == 1 + + # Ability 456 — no filter + links = await filtered_planner._fetch_links() assert len(links) == 2 - assert links[0].paw == 'paw1' - assert links[0].ability.ability_id == '456' - assert links[0].command == 'test command' - assert links[1].paw == 'paw4' - assert links[1].ability.ability_id == '456' - assert links[1].command == 'test command' - - links = await filtered_planner._get_pending_links('1011') + + # Skip 789, go to 1011 + links = await filtered_planner._fetch_links() assert len(links) == 2 - assert links[0].paw == 'paw1' - assert links[0].ability.ability_id == '1011' - assert links[0].command == 'test command' - assert links[1].paw == 'paw3' - assert links[1].ability.ability_id == '1011' - assert links[1].command == 'test command' - - async def test_fetch_links_with_filter(self, filtered_planner): - assert not filtered_planner.pending_links - assert filtered_planner.current_ability_index == 0 - - # first pass - links_to_use = await filtered_planner._fetch_links() - assert len(filtered_planner.pending_links) == 1 - assert filtered_planner.pending_links[0].paw == 'paw1' - assert filtered_planner.pending_links[0].command == 'test command variant' - assert filtered_planner.pending_links[0].ability.ability_id == '123' - assert filtered_planner.current_ability_index == 1 - assert len(links_to_use) == 1 - assert links_to_use[0].paw == 'paw1' - assert links_to_use[0].command == 'test command' - assert links_to_use[0].ability.ability_id == '123' - - # second pass - finishes links from first pass - links_to_use = await filtered_planner._fetch_links() - assert len(filtered_planner.pending_links) == 0 - assert filtered_planner.current_ability_index == 1 - assert len(links_to_use) == 1 - assert links_to_use[0].paw == 'paw1' - assert links_to_use[0].command == 'test command variant' - assert links_to_use[0].ability.ability_id == '123' - - # third pass - ability #2 - links_to_use = await filtered_planner._fetch_links() - assert len(filtered_planner.pending_links) == 0 - assert filtered_planner.current_ability_index == 2 - assert len(links_to_use) == 2 - assert links_to_use[0].paw == 'paw1' - assert links_to_use[0].command == 'test command' - assert links_to_use[0].ability.ability_id == '456' - assert links_to_use[1].paw == 'paw4' - assert links_to_use[1].command == 'test command' - assert links_to_use[1].ability.ability_id == '456' - - # fourth pass - skip ability #3 and go to #4 - links_to_use = await filtered_planner._fetch_links() - assert len(filtered_planner.pending_links) == 0 - assert filtered_planner.current_ability_index == 4 - assert len(links_to_use) == 2 - assert links_to_use[0].paw == 'paw1' - assert links_to_use[0].command == 'test command' - assert links_to_use[0].ability.ability_id == '1011' - assert links_to_use[1].paw == 'paw3' - assert links_to_use[1].command == 'test command' - assert links_to_use[1].ability.ability_id == '1011' - - # fifth pass - end - links_to_use = await filtered_planner._fetch_links() - assert len(filtered_planner.pending_links) == 0 - assert filtered_planner.current_ability_index == 4 - assert len(links_to_use) == 0 + + # End + links = await filtered_planner._fetch_links() + assert len(links) == 0 + + async def test_empty_ordering(self, generate_planner, dummy_agents): + planner = generate_planner([], dummy_agents) + links = await planner._fetch_links() + assert links == [] + + async def test_single_ability(self, generate_planner, dummy_agents): + planner = generate_planner(['123'], dummy_agents) + links = await planner._fetch_links() + assert len(links) > 0 + # Drain remaining + while True: + more = await planner._fetch_links() + if not more: + break + + +# --------------------------------------------------------------------------- +# Tests — fetch_and_run_links +# --------------------------------------------------------------------------- + +class TestFetchAndRunLinks: + async def test_runs_and_waits(self, planner_without_filter): + await planner_without_filter.fetch_and_run_links() + assert planner_without_filter.next_bucket == BUCKET_NAME + + async def test_sets_none_when_empty(self, generate_planner, dummy_agents): + planner = generate_planner([], dummy_agents) + await planner.fetch_and_run_links() + assert planner.next_bucket is None + + async def test_exhausts_all_links(self, planner_without_filter): + """Run fetch_and_run_links until planner stops.""" + iterations = 0 + while planner_without_filter.next_bucket is not None: + await planner_without_filter.fetch_and_run_links() + iterations += 1 + if iterations > 20: + pytest.fail('Planner did not terminate') + assert planner_without_filter.next_bucket is None + + +# --------------------------------------------------------------------------- +# Tests — execute +# --------------------------------------------------------------------------- + +class TestExecute: + async def test_execute_delegates(self): + from unittest.mock import AsyncMock + mock_planning_svc = AsyncMock() + op = DummyOperation(DummyAdversary([]), []) + planner = LogicalPlanner(op, mock_planning_svc) + await planner.execute() + mock_planning_svc.execute_planner.assert_awaited_once_with(planner) diff --git a/tests/test_hook.py b/tests/test_hook.py new file mode 100644 index 0000000..7334211 --- /dev/null +++ b/tests/test_hook.py @@ -0,0 +1,116 @@ +"""Tests for hook.py — plugin enable function and module-level attributes.""" +import os +import pytest +from unittest.mock import patch, MagicMock, AsyncMock + +from app.utility.base_world import BaseWorld + + +class TestHookModuleAttributes: + """Test the module-level constants in hook.py.""" + + def test_name(self): + from plugins.emu import hook + assert hook.name == 'Emu' + + def test_description(self): + from plugins.emu import hook + assert isinstance(hook.description, str) + assert len(hook.description) > 0 + + def test_address(self): + from plugins.emu import hook + assert hook.address == '/plugin/emu/gui' + + def test_access(self): + from plugins.emu import hook + assert hook.access == BaseWorld.Access.RED + + def test_data_dir(self): + from plugins.emu import hook + expected = os.path.join('plugins', 'emu', 'data') + assert hook.data_dir == expected + + +class TestHookEnable: + """Test the enable() coroutine.""" + + async def test_enable_clones_when_repo_missing(self): + from plugins.emu import hook + from plugins.emu.app.emu_svc import EmuService + + mock_app_svc = MagicMock() + mock_app_svc.application = MagicMock() + mock_app_svc.application.router = MagicMock() + + services = { + 'app_svc': mock_app_svc, + 'auth_svc': MagicMock(), + 'data_svc': MagicMock(), + } + + with patch.object(BaseWorld, 'apply_config'): + with patch.object(BaseWorld, 'strip_yml', return_value=[{'evals_c2_host': '127.0.0.1', 'evals_c2_port': 8888}]): + with patch.object(BaseWorld, 'get_config', return_value=None): + with patch.object(os.path, 'isdir', return_value=False): + with patch.object(EmuService, 'clone_repo', new_callable=AsyncMock) as mock_clone: + with patch.object(EmuService, 'decrypt_payloads', new_callable=AsyncMock) as mock_decrypt: + with patch.object(EmuService, 'populate_data_directory', new_callable=AsyncMock) as mock_pop: + await hook.enable(services) + + mock_clone.assert_awaited_once() + mock_decrypt.assert_awaited_once() + mock_pop.assert_awaited_once() + + async def test_enable_skips_clone_when_repo_exists(self): + from plugins.emu import hook + from plugins.emu.app.emu_svc import EmuService + + mock_app_svc = MagicMock() + mock_app_svc.application = MagicMock() + mock_app_svc.application.router = MagicMock() + + services = { + 'app_svc': mock_app_svc, + 'auth_svc': MagicMock(), + 'data_svc': MagicMock(), + } + + with patch.object(BaseWorld, 'apply_config'): + with patch.object(BaseWorld, 'strip_yml', return_value=[{'evals_c2_host': '127.0.0.1', 'evals_c2_port': 8888}]): + with patch.object(BaseWorld, 'get_config', return_value=None): + with patch.object(os.path, 'isdir', return_value=True): + with patch.object(EmuService, 'clone_repo', new_callable=AsyncMock) as mock_clone: + with patch.object(EmuService, 'decrypt_payloads', new_callable=AsyncMock): + with patch.object(EmuService, 'populate_data_directory', new_callable=AsyncMock): + await hook.enable(services) + + mock_clone.assert_not_awaited() + + async def test_enable_adds_gui_route(self): + from plugins.emu import hook + + mock_app_svc = MagicMock() + mock_app_svc.application = MagicMock() + mock_router = MagicMock() + mock_app_svc.application.router = mock_router + + services = { + 'app_svc': mock_app_svc, + 'auth_svc': MagicMock(), + 'data_svc': MagicMock(), + } + + with patch.object(BaseWorld, 'apply_config'): + with patch.object(BaseWorld, 'strip_yml', return_value=[{'evals_c2_host': '127.0.0.1', 'evals_c2_port': 8888}]): + with patch.object(BaseWorld, 'get_config', return_value=None): + with patch.object(os.path, 'isdir', return_value=True): + from plugins.emu.app.emu_svc import EmuService + with patch.object(EmuService, 'decrypt_payloads', new_callable=AsyncMock): + with patch.object(EmuService, 'populate_data_directory', new_callable=AsyncMock): + await hook.enable(services) + + # Should have registered GET /plugin/emu/gui route + route_calls = mock_router.add_route.call_args_list + get_routes = [c for c in route_calls if c[0][0] == 'GET' and c[0][1] == '/plugin/emu/gui'] + assert len(get_routes) >= 1 diff --git a/tests/test_requirements.py b/tests/test_requirements.py new file mode 100644 index 0000000..57a3734 --- /dev/null +++ b/tests/test_requirements.py @@ -0,0 +1,150 @@ +"""Tests for app/requirements/ — check_registered and check_lightneuron_registered.""" +import pytest +from unittest.mock import MagicMock, AsyncMock, patch + +from app.utility.base_service import BaseService +from tests.conftest import _StubFact, _StubLink + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +class DummyAgent: + def __init__(self, paw): + self.paw = paw + + +def _make_link(used_values): + """Create a stub link with used facts.""" + facts = [_StubFact(trait='agent_paw', value=v) for v in used_values] + return _StubLink(command='cmd', paw='test', used=facts) + + +class DummyOperation: + def __init__(self, active_paws): + self._active_paws = active_paws + + async def active_agents(self): + return [DummyAgent(paw) for paw in self._active_paws] + + +# --------------------------------------------------------------------------- +# Tests — check_registered (Requirement) +# --------------------------------------------------------------------------- + +class TestCheckRegistered: + def _get_requirement(self): + from plugins.emu.app.requirements.check_registered import Requirement + return Requirement() + + async def test_enforce_true_when_agent_active(self): + req = self._get_requirement() + op = DummyOperation(active_paws=['paw1', 'paw2']) + link = _make_link(['paw1']) + result = await req.enforce(link, op) + assert result is True + + async def test_enforce_false_when_agent_inactive(self): + req = self._get_requirement() + op = DummyOperation(active_paws=['paw2', 'paw3']) + link = _make_link(['paw1']) + result = await req.enforce(link, op) + assert result is False + + async def test_enforce_false_empty_used(self): + req = self._get_requirement() + op = DummyOperation(active_paws=['paw1']) + link = _make_link([]) + result = await req.enforce(link, op) + assert result is False + + async def test_enforce_false_no_active_agents(self): + req = self._get_requirement() + op = DummyOperation(active_paws=[]) + link = _make_link(['paw1']) + result = await req.enforce(link, op) + assert result is False + + async def test_enforce_multiple_used_first_matches(self): + req = self._get_requirement() + op = DummyOperation(active_paws=['paw1']) + link = _make_link(['paw1', 'paw2']) + result = await req.enforce(link, op) + assert result is True + + async def test_enforce_multiple_used_second_matches(self): + req = self._get_requirement() + op = DummyOperation(active_paws=['paw2']) + link = _make_link(['paw1', 'paw2']) + result = await req.enforce(link, op) + assert result is True + + +# --------------------------------------------------------------------------- +# Tests — check_lightneuron_registered (Requirement) +# --------------------------------------------------------------------------- + +class TestCheckLightneuronRegistered: + def _get_requirement(self): + from plugins.emu.app.requirements.check_lightneuron_registered import Requirement + return Requirement() + + def _setup_data_svc(self, agent_paws): + """Set up a mock data_svc with agents in its ram dict.""" + agents = [DummyAgent(paw) for paw in agent_paws] + data_svc = MagicMock() + data_svc.ram = {'agents': agents} + BaseService._services['data_svc'] = data_svc + return data_svc + + async def test_enforce_true_exact_match(self): + req = self._get_requirement() + self._setup_data_svc(['paw1', 'paw2']) + link = _make_link(['paw1']) + result = await req.enforce(link, MagicMock()) + assert result is True + + async def test_enforce_true_with_at_symbol(self): + """The lightneuron requirement strips @ from values.""" + req = self._get_requirement() + self._setup_data_svc(['paw1', 'paw2']) + link = _make_link(['p@aw1']) + result = await req.enforce(link, MagicMock()) + assert result is True + + async def test_enforce_false_no_match(self): + req = self._get_requirement() + self._setup_data_svc(['paw1', 'paw2']) + link = _make_link(['paw99']) + result = await req.enforce(link, MagicMock()) + assert result is False + + async def test_enforce_false_empty_used(self): + req = self._get_requirement() + self._setup_data_svc(['paw1']) + link = _make_link([]) + result = await req.enforce(link, MagicMock()) + assert result is False + + async def test_enforce_false_no_agents(self): + req = self._get_requirement() + self._setup_data_svc([]) + link = _make_link(['paw1']) + result = await req.enforce(link, MagicMock()) + assert result is False + + async def test_enforce_at_symbol_stripped_multiple(self): + """Multiple @ symbols should all be removed.""" + req = self._get_requirement() + self._setup_data_svc(['abc']) + link = _make_link(['a@b@c']) + result = await req.enforce(link, MagicMock()) + assert result is True + + async def test_enforce_multiple_used_facts(self): + req = self._get_requirement() + self._setup_data_svc(['paw2']) + link = _make_link(['paw1', 'paw2']) + result = await req.enforce(link, MagicMock()) + assert result is True diff --git a/tests/test_vssadmin_shadow_parser.py b/tests/test_vssadmin_shadow_parser.py index 2862fd3..1a6f422 100644 --- a/tests/test_vssadmin_shadow_parser.py +++ b/tests/test_vssadmin_shadow_parser.py @@ -1,6 +1,18 @@ +"""Exhaustive tests for app/parsers/vssadmin_shadow.py.""" import pytest from plugins.emu.app.parsers.vssadmin_shadow import Parser +from tests.conftest import _StubFact + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture +def parser(): + p = Parser() + return p @pytest.fixture @@ -15,10 +27,104 @@ def dummy_blob(): """ -class TestVssadminShadowParser: +@pytest.fixture +def multi_volume_blob(): + return """ + Shadow Copy Volume Name: \\\\?\\GLOBALROOT\\Device\\HarddiskVolumeShadowCopy1 + Shadow Copy Volume Name: \\\\?\\GLOBALROOT\\Device\\HarddiskVolumeShadowCopy2 +""" + + +# --------------------------------------------------------------------------- +# Tests — _get_volume_name +# --------------------------------------------------------------------------- + +class TestGetVolumeName: def test_parse_blob(self, dummy_blob): results = Parser._get_volume_name(dummy_blob) assert results == r'\\?\GLOBALROOT\Device\HarddiskVolumeShadowCopy7' def test_parse_blob_without_match(self): assert not Parser._get_volume_name('no volume name here') + + def test_empty_string(self): + assert not Parser._get_volume_name('') + + def test_partial_match(self): + blob = 'Shadow Copy Volume Name:' + assert not Parser._get_volume_name(blob) + + def test_returns_first_match_only(self, multi_volume_blob): + result = Parser._get_volume_name(multi_volume_blob) + assert result is not None + # Should return the first match + + def test_leading_whitespace(self): + blob = ' Shadow Copy Volume Name: \\\\?\\GLOBALROOT\\Device\\ShadowCopy99' + result = Parser._get_volume_name(blob) + assert result is not None + assert 'ShadowCopy99' in result + + def test_no_whitespace_prefix(self): + blob = 'Shadow Copy Volume Name: \\\\?\\GLOBALROOT\\Device\\ShadowCopy1' + result = Parser._get_volume_name(blob) + assert result is not None + + def test_multiline_output(self): + blob = """some output +other line + Shadow Copy Volume Name: \\\\?\\GLOBALROOT\\Device\\HarddiskVolumeShadowCopy5 +trailing output +""" + result = Parser._get_volume_name(blob) + assert 'HarddiskVolumeShadowCopy5' in result + + +# --------------------------------------------------------------------------- +# Tests — parse method +# --------------------------------------------------------------------------- + +class TestParse: + def test_parse_returns_empty_on_no_match(self, parser): + result = parser.parse('no matching output') + assert result == [] + + def test_parse_returns_empty_on_empty(self, parser): + result = parser.parse('') + assert result == [] + + def test_parse_with_mappers(self, parser, dummy_blob): + class MockMapper: + def __init__(self): + self.source = 'src_trait' + self.target = 'tgt_trait' + self.edge = 'has' + + parser.mappers = [MockMapper()] + result = parser.parse(dummy_blob) + assert len(result) == 1 + rel = result[0] + assert rel.source.trait == 'src_trait' + assert rel.target.trait == 'tgt_trait' + assert rel.edge == 'has' + + def test_parse_with_multiple_mappers(self, parser, dummy_blob): + class MockMapper: + def __init__(self, source, target, edge): + self.source = source + self.target = target + self.edge = edge + + parser.mappers = [ + MockMapper('s1', 't1', 'edge1'), + MockMapper('s2', 't2', 'edge2'), + ] + result = parser.parse(dummy_blob) + assert len(result) == 2 + assert result[0].edge == 'edge1' + assert result[1].edge == 'edge2' + + def test_parse_no_mappers(self, parser, dummy_blob): + parser.mappers = [] + result = parser.parse(dummy_blob) + assert result == []