diff --git a/app/api/rest_api.py b/app/api/rest_api.py index 96e87d079..98216129e 100644 --- a/app/api/rest_api.py +++ b/app/api/rest_api.py @@ -137,9 +137,9 @@ async def download_file(self, request): ('FILENAME', display_name)]) return web.Response(body=content, headers=headers) except FileNotFoundError: - return web.HTTPNotFound(body='File not found') + raise web.HTTPNotFound(text='File not found') except Exception as e: - return web.HTTPNotFound(body=str(e)) + raise web.HTTPNotFound(text=str(e)) @check_authorization async def download_exfil_file(self, request): @@ -153,20 +153,17 @@ def is_in_exfil_dir(f): file = base64.b64decode(request.query.get('file')).decode('ascii') file = os.path.normpath(file) # normalize path to remove all directory traversal attempts then check for presence in exfil dir if not is_in_exfil_dir(file): - return web.HTTPNotFound(body="File not found in exfil dir") + raise web.HTTPNotFound(text="File not found in exfil dir") filename = file.split(os.sep)[-1] path = os.sep.join(file.split(os.sep)[:-1]) _, content = await self.file_svc.read_file(filename, location=path) headers = dict([('CONTENT-DISPOSITION', 'attachment; filename="%s"' % filename), ('FILENAME', filename)]) return web.Response(body=content, headers=headers) + except web.HTTPNotFound as e: + raise e except FileNotFoundError: - return web.HTTPNotFound(body='File not found') + raise web.HTTPNotFound(text='File not found') except Exception as e: - return web.HTTPNotFound(body=str(e)) - return web.HTTPBadRequest(body='A file needs to be specified for download') - - @staticmethod - def _request_errors(request): - errors = [] - return errors + raise web.HTTPNotFound(text=str(e)) + raise web.HTTPBadRequest(text='A file needs to be specified for download') diff --git a/app/api/v2/handlers/ability_api.py b/app/api/v2/handlers/ability_api.py index 2aaa14c37..14f0b487d 100644 --- a/app/api/v2/handlers/ability_api.py +++ b/app/api/v2/handlers/ability_api.py @@ -109,4 +109,4 @@ async def update_ability(self, request: web.Request): description='HTTP 204 Status Code (No Content)') async def delete_ability(self, request: web.Request): await self.delete_on_disk_object(request) - return web.HTTPNoContent() + raise web.HTTPNoContent() diff --git a/app/api/v2/handlers/payload_api.py b/app/api/v2/handlers/payload_api.py index 6c6350348..d5e0df6c8 100644 --- a/app/api/v2/handlers/payload_api.py +++ b/app/api/v2/handlers/payload_api.py @@ -136,7 +136,7 @@ async def delete_payloads(self, request: web.Request): response = web.HTTPNotFound() except PermissionError: response = web.HTTPForbidden(reason="Permission denied.") - return response + raise response @classmethod async def __generate_file_name_and_path(cls, sanitized_filename: str) -> [str, str]: diff --git a/app/ascii_banner.py b/app/ascii_banner.py index 0411aa3aa..b5a89bd43 100644 --- a/app/ascii_banner.py +++ b/app/ascii_banner.py @@ -36,7 +36,7 @@ " -BANNER_SECTION_3 = "\ +_BANNER_SECTION_3 = "\ ╚██████╗██║ ██║███████╗██████╔╝███████╗██║ ██║██║ ██║\n\ ╚═════╝╚═╝ ╚═╝╚══════╝╚═════╝ ╚══════╝╚═╝ ╚═╝╚═╝ ╚═╝\n\ " @@ -46,17 +46,18 @@ def no_color(): return int(os.environ.get("NO_COLOR", 0)) == 1 -if no_color(): - ASCII_BANNER = _BANNER -else: - ASCII_BANNER = f"{DARK_BLUE}{_BANNER_SECTION_1}{DARK_PURPLE}{_BANNER_SECTION_2}{DARK_RED}{BANNER_SECTION_3}{END}" +def get_ascii_banner(): + if no_color(): + return _BANNER + else: + return f"{DARK_BLUE}{_BANNER_SECTION_1}{DARK_PURPLE}{_BANNER_SECTION_2}{DARK_RED}{_BANNER_SECTION_3}{END}" def print_rich_banner(): """Print banner using Python Rich library""" if no_color(): - rich_print(f"{_BANNER_SECTION_1}{_BANNER_SECTION_2}{BANNER_SECTION_3}") + rich_print(f"{_BANNER_SECTION_1}{_BANNER_SECTION_2}{_BANNER_SECTION_3}") else: rich_print( - f"[blue]{_BANNER_SECTION_1}[/blue][purple]{_BANNER_SECTION_2}[/purple][red]{BANNER_SECTION_3}[/red]" + f"[blue]{_BANNER_SECTION_1}[/blue][purple]{_BANNER_SECTION_2}[/purple][red]{_BANNER_SECTION_3}[/red]" ) diff --git a/app/contacts/contact_ftp.py b/app/contacts/contact_ftp.py index b86a725bb..f0980b5cc 100644 --- a/app/contacts/contact_ftp.py +++ b/app/contacts/contact_ftp.py @@ -165,15 +165,12 @@ async def handle_agent_file(self, split_file_path, file_bytes): profile = json.loads(file_bytes.decode()) agent, instructions = await self.contact_caldera_server(profile) paw, contents = await self.create_beacon_response(agent, instructions) - self.write_file(paw, 'Response.txt', json.dumps(contents)) + self.write_file(paw, 'Response.txt', bytes(json.dumps(contents).encode('ascii'))) elif re.match(r'^Payload\.txt$', split_file_path[-1]): profile = json.loads(file_bytes.decode()) file_path, contents, display_name = await self.get_payload_file(profile) if file_path is not None: - self.write_file(profile.get('paw'), profile.get('file'), str(contents)) - elif re.match(r'^Results\.txt$', split_file_path[-1]): - profile = json.loads(file_bytes.decode()) - await self.contact_caldera_server(profile) + self.write_file(profile.get('paw'), profile.get('file'), contents) else: paw = split_file_path[-2] filename = split_file_path[-1] @@ -202,7 +199,7 @@ def write_file(self, paw, file_name, contents): if not os.path.exists(agent_dir_path): os.makedirs(agent_dir_path) file_path = os.path.join(agent_dir_path, file_name) - with open(file_path, 'w+') as f: + with open(file_path, 'w+b') as f: f.write(contents) self.logger.debug('File written to: %s' % agent_dir_path) except IOError: diff --git a/app/contacts/contact_gist.py b/app/contacts/contact_gist.py index 6dfb2de12..6078e9f0a 100644 --- a/app/contacts/contact_gist.py +++ b/app/contacts/contact_gist.py @@ -14,7 +14,7 @@ def api_access(func): async def process(*args, **kwargs): async with aiohttp.ClientSession(headers=dict(Authorization='token {}'.format(args[0].token)), - connector=aiohttp.TCPConnector(verify_ssl=False)) as session: + connector=aiohttp.TCPConnector(ssl=False)) as session: kwargs['session'] = session return await func(*args, **kwargs) return process diff --git a/app/contacts/contact_html.py b/app/contacts/contact_html.py index bdefa63d9..27de73967 100644 --- a/app/contacts/contact_html.py +++ b/app/contacts/contact_html.py @@ -16,8 +16,7 @@ def __init__(self, services): async def start(self): self.app_svc.application.router.add_route('*', self.get_config('app.contact.html'), self._accept_beacon) - @template('weather.html') - async def _accept_beacon(self, request): + async def _beacon_helper(self, request): try: profile = json.loads(self.decode_bytes(await request.text())) profile['paw'] = profile.get('paw') @@ -30,3 +29,7 @@ async def _accept_beacon(self, request): return dict(instructions=self.encode_string(json.dumps(response))) except Exception: return dict(instructions=[]) + + @template('weather.html') + async def _accept_beacon(self, request): + return await self._beacon_helper(request) diff --git a/app/contacts/contact_slack.py b/app/contacts/contact_slack.py index e194c03bc..4e5dbcd5c 100644 --- a/app/contacts/contact_slack.py +++ b/app/contacts/contact_slack.py @@ -14,7 +14,7 @@ def api_access(func): async def process(*args, **kwargs): async with aiohttp.ClientSession(headers=dict(Authorization='Bearer {}'.format(args[0].key)), - connector=aiohttp.TCPConnector(verify_ssl=False)) as session: + connector=aiohttp.TCPConnector(ssl=False)) as session: kwargs['session'] = session return await func(*args, **kwargs) return process diff --git a/app/objects/secondclass/c_variation.py b/app/objects/secondclass/c_variation.py index 0b85e6a0e..8d1be74b6 100644 --- a/app/objects/secondclass/c_variation.py +++ b/app/objects/secondclass/c_variation.py @@ -29,3 +29,8 @@ def __init__(self, description, command): super().__init__() self.description = description self._command = self.encode_string(command) + + def __eq__(self, other): + if isinstance(other, Variation): + return self.description == other.description and self._command == other._command + return False diff --git a/app/service/data_svc.py b/app/service/data_svc.py index 51d13ca49..d4b266a86 100644 --- a/app/service/data_svc.py +++ b/app/service/data_svc.py @@ -87,7 +87,7 @@ async def destroy(): if not os.path.exists(DATA_BACKUP_DIR): os.mkdir(DATA_BACKUP_DIR) - timestamp = datetime.datetime.utcnow().strftime('%Y%m%d%H%M%S') + timestamp = datetime.datetime.now(datetime.timezone.utc).strftime('%Y%m%d%H%M%S') tarball_path = os.path.join(DATA_BACKUP_DIR, f'backup-{timestamp}.tar.gz') with tarfile.open(tarball_path, 'w:gz') as tarball: @@ -178,8 +178,8 @@ async def load_ability_file(self, filename, access): 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.') + self.log.warning(f'Tactic for ability={ability_id} is not in the ability file path {filename}.') + self.log.warning('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, diff --git a/app/service/file_svc.py b/app/service/file_svc.py index c8d319696..6b199bee6 100644 --- a/app/service/file_svc.py +++ b/app/service/file_svc.py @@ -103,7 +103,7 @@ async def create_exfil_operation_directory(self, dir_name, agent_name): special_chars = {ord(c): '_' for c in r':<>"/\|?*'} agent_opid = [(x.name.translate(special_chars), '_', x.start.strftime("%Y-%m-%d_%H%M%SZ")) for x in op_list_filtered if agent_name in [y.paw for y in x.agents]] - path = os.path.join((dir_name), ''.join(agent_opid[0])) + path = os.path.join((dir_name), ''.join(agent_opid[0])) if agent_opid else dir_name if not os.path.exists(path): os.makedirs(path) return path diff --git a/server.py b/server.py index ad55e41ef..e7f723d48 100644 --- a/server.py +++ b/server.py @@ -16,7 +16,7 @@ import app.api.v2 from app import version -from app.ascii_banner import ASCII_BANNER, no_color, print_rich_banner +from app.ascii_banner import get_ascii_banner, no_color, print_rich_banner from app.api.rest_api import RestApi from app.api.v2.responses import apispec_request_validation_middleware from app.api.v2.security import pass_option_middleware @@ -199,7 +199,7 @@ def list_str(values): return values.split(",") parser = argparse.ArgumentParser( - description=ASCII_BANNER, + description=get_ascii_banner(), formatter_class=argparse.RawDescriptionHelpFormatter ) parser.add_argument( diff --git a/tests/api/v2/handlers/test_abilities_api.py b/tests/api/v2/handlers/test_abilities_api.py index 2d223fc29..8ded67084 100644 --- a/tests/api/v2/handlers/test_abilities_api.py +++ b/tests/api/v2/handlers/test_abilities_api.py @@ -100,13 +100,23 @@ async def test_get_nonexistent_ability_by_id(self, api_v2_client, api_cookies): resp = await api_v2_client.get('/api/v2/abilities/999', cookies=api_cookies) assert resp.status == HTTPStatus.NOT_FOUND - async def test_create_ability(self, api_v2_client, api_cookies, mocker, async_return, new_ability_payload): + async def test_create_and_delete_ability(self, api_v2_client, api_cookies, mocker, async_return, new_ability_payload): + # Test creation + target_ability_id = new_ability_payload.get('ability_id') resp = await api_v2_client.post('/api/v2/abilities', cookies=api_cookies, json=new_ability_payload) assert resp.status == HTTPStatus.OK ability_data = await resp.json() assert ability_data == new_ability_payload ability_exists = await BaseService.get_service('data_svc').locate('abilities', {'ability_id': '456'}) assert ability_exists + assert os.path.exists(f'data/abilities/collection/{target_ability_id}.yml') + + # Test deletion + resp = await api_v2_client.delete(f'/api/v2/abilities/{target_ability_id}', cookies=api_cookies) + assert resp.status in (HTTPStatus.OK, HTTPStatus.NO_CONTENT) + ability_exists = await BaseService.get_service('data_svc').locate('abilities', {'ability_id': '456'}) + assert not ability_exists + assert not os.path.exists(f'data/abilities/collection/{target_ability_id}.yml') async def test_unauthorized_create_ability(self, api_v2_client, new_ability_payload): resp = await api_v2_client.post('/api/v2/abilities', json=new_ability_payload) diff --git a/tests/api/v2/handlers/test_config_api.py b/tests/api/v2/handlers/test_config_api.py new file mode 100644 index 000000000..034a487a3 --- /dev/null +++ b/tests/api/v2/handlers/test_config_api.py @@ -0,0 +1,83 @@ +import pytest + +from unittest import mock + +from app.api.v2.managers.config_api_manager import ConfigApiManager, SENSITIVE_CONFIG_PROPS +from app.utility.base_world import BaseWorld + + +@pytest.fixture +def base_world(app_config, agent_config): + BaseWorld.clear_config() + BaseWorld.apply_config('main', app_config) + BaseWorld.apply_config('agents', agent_config) + + yield BaseWorld + + BaseWorld.clear_config() + + +@pytest.fixture +def filter_config(): + def _filter_config(config): + to_filter = config + for sensitive_prop in SENSITIVE_CONFIG_PROPS: + to_filter.pop(sensitive_prop, None) + return to_filter + return _filter_config + + +class TestConfigApi: + async def test_get_config_with_name(self, api_v2_client, api_cookies, base_world, filter_config, app_config, agent_config): + resp = await api_v2_client.get('/api/v2/config/main', cookies=api_cookies) + config_json = await resp.json() + want = filter_config(app_config) + assert config_json == want + + resp = await api_v2_client.get('/api/v2/config/agents', cookies=api_cookies) + config_json = await resp.json() + assert config_json == agent_config + + # Test nonexistent config + resp = await api_v2_client.get('/api/v2/config/doesnotexist', cookies=api_cookies) + resp_dict = await resp.json() + want = dict(error='Config not found: doesnotexist') + assert resp_dict == want + + async def test_get_update_main_config(self, api_v2_client, api_cookies, base_world, filter_config, app_config): + data = dict(prop='app.contact.html', value='/newhtmlcontact') + resp = await api_v2_client.patch('/api/v2/config/main', json=data, cookies=api_cookies) + config_json = await resp.json() + want = filter_config(app_config) + want['app.contact.html'] = '/newhtmlcontact' + assert config_json == want + + # Test sensitive field + data = dict(prop='host', value='127.0.0.3') + resp = await api_v2_client.patch('/api/v2/config/main', json=data, cookies=api_cookies) + resp_dict = await resp.json() + want = dict( + error='Update not allowed', + details=dict(property='host') + ) + assert resp_dict == want + + async def test_get_update_agents_config(self, api_v2_client, api_cookies, base_world, agent_config): + data = dict(watchdog='1', implant_name='newname', bootstrap_abilities=['abil1', '', 'DNE', 'abil2']) + with mock.patch.object(ConfigApiManager, '_get_loaded_ability_ids', return_value={'abil1', 'abil2'}): + resp = await api_v2_client.patch('/api/v2/config/agents', json=data, cookies=api_cookies) + config_json = await resp.json() + want = { + 'sleep_min': 30, + 'sleep_max': 60, + 'untrusted_timer': 90, + 'watchdog': 1, + 'implant_name': 'newname', + 'deadman_abilities': [ + 'this-is-a-fake-ability' + ], + 'bootstrap_abilities': [ + 'abil1', 'abil2' + ] + } + assert config_json == want diff --git a/tests/api/v2/handlers/test_operations_api.py b/tests/api/v2/handlers/test_operations_api.py index e936983b7..4b1ab8a37 100644 --- a/tests/api/v2/handlers/test_operations_api.py +++ b/tests/api/v2/handlers/test_operations_api.py @@ -42,18 +42,24 @@ async def test_nonexistent_operation_get_operation_by_id(self, api_v2_client, ap resp = await api_v2_client.get('/api/v2/operations/999', cookies=api_cookies) assert resp.status == HTTPStatus.NOT_FOUND - async def test_get_operations_summary(self, api_v2_client, api_cookies, test_operation): - resp = await api_v2_client.get('/api/v2/operations/summary', cookies=api_cookies) - assert resp.status == HTTPStatus.OK - operations_list = await resp.json() - assert len(operations_list) == 1 - operation_dict = operations_list[0] - assert operation_dict['name'] == test_operation['name'] - assert operation_dict['id'] == test_operation['id'] - assert 'agents' in operation_dict - assert 'hosts' in operation_dict - assert 'chain' not in operation_dict - assert 'host_group' not in operation_dict + async def test_get_operations_summary(self, api_v2_client, api_cookies, mocker, async_return, + test_operation): + with mocker.patch('app.objects.c_operation.Operation.all_facts') as mock_all_facts: + mock_all_facts.return_value = async_return([]) + resp = await api_v2_client.get('/api/v2/operations/summary', cookies=api_cookies) + assert resp.status == HTTPStatus.OK + operations_list = await resp.json() + assert len(operations_list) == 1 + operation_dict = operations_list[0] + assert operation_dict['name'] == test_operation['name'] + assert operation_dict['id'] == test_operation['id'] + assert operation_dict['jitter'] == test_operation['jitter'] + assert operation_dict['planner']['name'] == test_operation['planner']['name'] + assert operation_dict['adversary']['name'] == test_operation['adversary']['name'] + assert 'agents' in operation_dict + assert 'hosts' in operation_dict + assert 'chain' not in operation_dict + assert 'host_group' not in operation_dict async def test_get_operations_summary_links_with_no_paw_skipped( self, api_v2_client, api_cookies, test_operation): diff --git a/tests/api/v2/handlers/test_payloads_api.py b/tests/api/v2/handlers/test_payloads_api.py index 4ffa34eb7..92f741c82 100644 --- a/tests/api/v2/handlers/test_payloads_api.py +++ b/tests/api/v2/handlers/test_payloads_api.py @@ -1,7 +1,9 @@ import os import pathlib import tempfile +from aiohttp import FormData from http import HTTPStatus +from unittest import mock import pytest @@ -82,3 +84,48 @@ async def test_get_payloads_name_filter_with_sort_and_add_path( async def test_unauthorized_get_payloads(self, api_v2_client): resp = await api_v2_client.get('/api/v2/payloads') assert resp.status == HTTPStatus.UNAUTHORIZED + + @mock.patch.object(pathlib.Path, 'rename') + async def test_post_payloads(self, mock_rename, api_v2_client, api_cookies): + file_data = bytes([0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]) + with tempfile.TemporaryFile(mode='w+b') as tmp_file: + tmp_file.write(file_data) + tmp_file.flush() + tmp_file.seek(0) + + m = mock.mock_open() + with mock.patch('builtins.open', m): + upload_data = FormData() + upload_data.add_field('file', tmp_file, filename='testpostpayload') + resp = await api_v2_client.post('/api/v2/payloads', + data=upload_data) + assert resp.status == HTTPStatus.OK + assert await resp.json() == dict(payloads=['testpostpayload']) + mock_rename.assert_called_with('data/payloads/testpostpayload') + m.assert_called_with('data/payloads/temp_testpostpayload', 'wb') + m().write.assert_called_once_with(file_data) + + async def test_delete_payloads(self, api_v2_client, api_cookies): + want_path = pathlib.Path('data/payloads/testtodelete').resolve() + with mock.patch.object(os, 'remove') as mock_remove: + resp = await api_v2_client.delete('/api/v2/payloads/testtodelete') + mock_remove.assert_called_once_with(want_path) + assert resp.status == 204 + + # Test ValueError + with mock.patch.object(os, 'remove', side_effect=ValueError('testvalueerror')) as mock_remove: + resp = await api_v2_client.delete('/api/v2/payloads/testtodelete') + assert resp.status == 404 + assert resp.reason == 'testvalueerror' + + # Test FileNotFoundError + with mock.patch.object(os, 'remove', side_effect=FileNotFoundError()) as mock_remove: + resp = await api_v2_client.delete('/api/v2/payloads/testtodelete') + assert resp.status == 404 + assert resp.reason == 'Not Found' + + # Test PermissionError + with mock.patch.object(os, 'remove', side_effect=PermissionError()) as mock_remove: + resp = await api_v2_client.delete('/api/v2/payloads/testtodelete') + assert resp.status == 403 + assert resp.reason == 'Permission denied.' diff --git a/tests/conftest.py b/tests/conftest.py index ac0ac60de..8a3b31ba8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -18,20 +18,7 @@ from aiohttp import web import aiohttp_jinja2 from pathlib import Path -from app.api.v2.handlers.agent_api import AgentApi -from app.api.v2.handlers.ability_api import AbilityApi -from app.api.v2.handlers.objective_api import ObjectiveApi -from app.api.v2.handlers.adversary_api import AdversaryApi -from app.api.v2.handlers.operation_api import OperationApi -from app.api.v2.handlers.contact_api import ContactApi -from app.api.v2.handlers.obfuscator_api import ObfuscatorApi -from app.api.v2.handlers.plugins_api import PluginApi -from app.api.v2.handlers.fact_source_api import FactSourceApi -from app.api.v2.handlers.fact_api import FactApi -from app.api.v2.handlers.planner_api import PlannerApi -from app.api.v2.handlers.health_api import HealthApi -from app.api.v2.handlers.schedule_api import ScheduleApi -from app.api.v2.handlers.payload_api import PayloadApi +from app.api.v2 import make_app from app.objects.c_obfuscator import Obfuscator from app.objects.c_objective import Objective from app.objects.c_planner import PlannerSchema @@ -59,8 +46,6 @@ from app.utility.base_object import BaseObject from app.utility.base_service import BaseService from app.utility.base_world import BaseWorld -from app.api.v2.responses import json_request_validation_middleware -from app.api.v2.security import authentication_required_middleware_factory from app.api.v2.responses import apispec_request_validation_middleware from app.api.rest_api import RestApi @@ -333,33 +318,10 @@ def agent_config(): @pytest.fixture async def api_v2_client(aiohttp_client, contact_svc): - def make_app(svcs): - warnings.filterwarnings( - "ignore", - message="Multiple schemas resolved to the name" - ) - - app = web.Application( - middlewares=[ - authentication_required_middleware_factory(svcs['auth_svc']), - json_request_validation_middleware - ] - ) - AgentApi(svcs).add_routes(app) - AbilityApi(svcs).add_routes(app) - OperationApi(svcs).add_routes(app) - AdversaryApi(svcs).add_routes(app) - ContactApi(svcs).add_routes(app) - ObjectiveApi(svcs).add_routes(app) - ObfuscatorApi(svcs).add_routes(app) - PluginApi(svcs).add_routes(app) - FactApi(svcs).add_routes(app) - FactSourceApi(svcs).add_routes(app) - PlannerApi(svcs).add_routes(app) - HealthApi(svcs).add_routes(app) - ScheduleApi(svcs).add_routes(app) - PayloadApi(svcs).add_routes(app) - return app + warnings.filterwarnings( + "ignore", + message="Multiple schemas resolved to the name" + ) async def initialize(): with open(Path(__file__).parents[1] / 'conf' / 'default.yml', 'r') as fle: @@ -386,7 +348,7 @@ async def initialize(): await auth_svc.apply(app_svc.application, auth_svc.get_config('users')) await auth_svc.set_login_handlers(services) - app_svc.register_subapp('/api/v2', make_app(svcs=services)) + app_svc.register_subapp('/api/v2', make_app(services)) aiohttp_apispec.setup_aiohttp_apispec( app=app_svc.application, title='Caldera', diff --git a/tests/contacts/test_contact_ftp.py b/tests/contacts/test_contact_ftp.py index 8a771f529..4361f444d 100644 --- a/tests/contacts/test_contact_ftp.py +++ b/tests/contacts/test_contact_ftp.py @@ -1,8 +1,15 @@ import pytest +import aioftp +import json import os import shutil +from unittest import mock + from app.contacts import contact_ftp +from app.objects.c_agent import Agent +from app.service.contact_svc import ContactService +from app.service.file_svc import FileSvc from app.utility.base_world import BaseWorld beacon_profile = {'architecture': 'amd64', @@ -21,18 +28,23 @@ } +DUMMY_EXFIL_DIR = '/tmp/testexfildir' + + @pytest.fixture(scope='session') def base_world(): BaseWorld.clear_config() - BaseWorld.apply_config(name='main', config={'app.contact.ftp.host': '0.0.0.0', - 'app.contact.ftp.port': '2222', + BaseWorld.apply_config(name='main', config={'app.contact.ftp.host': '127.0.0.1', + 'app.contact.ftp.port': '62221', 'app.contact.ftp.pword': 'caldera', - 'app.contact.ftp.server.dir': 'ftp_dir_testing', + 'app.contact.ftp.server.dir': 'test_dummy_ftp_dir', 'app.contact.ftp.user': 'caldera_user', 'plugins': ['sandcat', 'stockpile'], 'crypt_salt': 'BLAH', 'api_key': 'ADMIN123', - 'encryption_key': 'ADMIN123'}) + 'encryption_key': 'ADMIN123', + 'encrypt_files': False, + 'exfil_dir': DUMMY_EXFIL_DIR}) BaseWorld.apply_config(name='agents', config={'sleep_max': 5, 'sleep_min': 5, 'untrusted_timer': 90, @@ -43,6 +55,8 @@ def base_world(): ]}) yield BaseWorld BaseWorld.clear_config() + if os.path.exists(DUMMY_EXFIL_DIR): + shutil.rmtree(DUMMY_EXFIL_DIR) @pytest.fixture @@ -53,19 +67,63 @@ async def ftp_c2(app_svc, base_world, contact_svc, data_svc, file_svc, obfuscato @pytest.fixture -def ftp_c2_my_server(ftp_c2): +async def ftp_c2_handler_server(ftp_c2): ftp_c2.set_up_server() - return ftp_c2.server + await ftp_c2.server.start(host=ftp_c2.server.host, port=ftp_c2.server.port) + yield ftp_c2.server + await ftp_c2.server.close() + if os.path.exists(ftp_c2.server.ftp_server_dir): + shutil.rmtree(ftp_c2.server.ftp_server_dir) + + +@pytest.fixture +def ftp_dummy_agent(): + return Agent(paw=TestFtpServer.dummy_beacon_data.get('paw'), sleep_min=5, sleep_max=5, watchdog=0, executors=['sh', 'proc']) + + +@pytest.fixture +async def ftp_client(ftp_c2_handler_server, ftp_dummy_agent): + client = aioftp.Client() + await client.connect(ftp_c2_handler_server.host, port=int(ftp_c2_handler_server.port)) + await client.login(user=ftp_c2_handler_server.login, password=ftp_c2_handler_server.pword) + await client.make_directory(ftp_dummy_agent.paw) + await client.change_directory(ftp_dummy_agent.paw) + yield client + + await client.quit() class TestFtpServer: + dummy_beacon_data = { + 'architecture': 'arm64', + 'available_contacts': ['HTTP'], + 'contact': 'HTTP', + 'deadman_enabled': True, + 'exe_name': 'splunkd', + 'executors': ['proc', 'sh'], + 'group': 'red', + 'host': 'myhost', + 'host_ip_addrs': ['10.0.2.15'], + 'location': '/home/testuser/splunkd', + 'origin_link_id': '', + 'paw': 'testpaw', + 'pid': 63025, + 'platform': 'linux', + 'ppid': 4357, + 'privilege': 'User', + 'proxy_receivers': None, + 'server': 'http://0.0.0.0:8888', + 'upstream_dest': 'http://0.0.0.0:8888', + 'username': 'testuser' + } + @staticmethod def test_server_setup(ftp_c2): assert ftp_c2.name == 'ftp' assert ftp_c2.description == 'Accept agent beacons through ftp' - assert ftp_c2.host == '0.0.0.0' - assert ftp_c2.port == '2222' - assert ftp_c2.directory == 'ftp_dir_testing' + assert ftp_c2.host == '127.0.0.1' + assert ftp_c2.port == '62221' + assert ftp_c2.directory == 'test_dummy_ftp_dir' assert ftp_c2.user == 'caldera_user' assert ftp_c2.pword == 'caldera' assert ftp_c2.server is None @@ -76,11 +134,58 @@ def test_set_up_server(ftp_c2): assert ftp_c2.server is not None @staticmethod - def test_my_server_setup(ftp_c2_my_server): - assert ftp_c2_my_server.host == '0.0.0.0' - assert ftp_c2_my_server.port == '2222' - assert ftp_c2_my_server.login == 'caldera_user' - assert ftp_c2_my_server.pword == 'caldera' - assert ftp_c2_my_server.ftp_server_dir == os.path.join(os.getcwd(), 'ftp_dir_testing') - assert os.path.exists(ftp_c2_my_server.ftp_server_dir) - shutil.rmtree(ftp_c2_my_server.ftp_server_dir) + async def test_my_server_setup(ftp_c2_handler_server): + assert ftp_c2_handler_server.host == '127.0.0.1' + assert ftp_c2_handler_server.port == '62221' + assert ftp_c2_handler_server.login == 'caldera_user' + assert ftp_c2_handler_server.pword == 'caldera' + assert ftp_c2_handler_server.ftp_server_dir == os.path.join(os.getcwd(), 'test_dummy_ftp_dir') + assert os.path.exists(ftp_c2_handler_server.ftp_server_dir) + + async def test_beacon(self, ftp_c2_handler_server, ftp_dummy_agent, ftp_client): + beacon_file_data = bytes(json.dumps(self.dummy_beacon_data).encode('ascii')) + + ftp_dummy_agent.pending_contact = 'newcontact' + with mock.patch.object(ContactService, 'handle_heartbeat', return_value=(ftp_dummy_agent, [])): + async with ftp_client.upload_stream('Alive.txt') as upload_stream: + await upload_stream.write(beacon_file_data) + + resp_path = os.path.join(ftp_c2_handler_server.ftp_server_dir, 'testpaw', 'Response.txt') + assert os.path.exists(resp_path) + want_response_dict = dict( + paw='testpaw', + sleep=5, + watchdog=0, + instructions='[]', + new_contact='newcontact', + ) + with open(resp_path, 'rb') as resp_file: + response = resp_file.read() + resp_dict = json.loads(response) + assert want_response_dict == resp_dict + + async def test_download_payload(self, ftp_c2_handler_server, ftp_dummy_agent, ftp_client): + payload_req = dict(file='testdownload', platform='linux', paw='testpaw') + payload_req_data = bytes(json.dumps(payload_req).encode('ascii')) + dummy_payload_data = bytes([0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]) + + with mock.patch.object(FileSvc, 'get_file', return_value=('testplugin/payloads/testdownload', dummy_payload_data, 'testdownload')): + async with ftp_client.upload_stream('Payload.txt') as upload_stream: + await upload_stream.write(payload_req_data) + + resp_path = os.path.join(ftp_c2_handler_server.ftp_server_dir, 'testpaw', 'testdownload') + assert os.path.exists(resp_path) + with open(resp_path, 'rb') as resp_file: + assert dummy_payload_data == resp_file.read() + + async def test_upload_file(self, ftp_c2_handler_server, ftp_dummy_agent, ftp_client): + dummy_file_data = bytes([0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]) + + async with ftp_client.upload_stream('testupload') as upload_stream: + await upload_stream.write(dummy_file_data) + + assert os.path.exists(DUMMY_EXFIL_DIR) + upload_path = os.path.join(DUMMY_EXFIL_DIR, 'testpaw', 'testupload') + assert os.path.exists(upload_path) + with open(upload_path, 'rb') as upload_file: + assert dummy_file_data == upload_file.read() diff --git a/tests/contacts/test_contact_gist.py b/tests/contacts/test_contact_gist.py index 0a2ee16b5..7493dfb3c 100644 --- a/tests/contacts/test_contact_gist.py +++ b/tests/contacts/test_contact_gist.py @@ -1,64 +1,228 @@ import pytest +import aiohttp +import base64 +import contextlib +import json +import logging +import uuid +from unittest import mock -from unittest.mock import patch +from app.objects.c_agent import Agent +from app.objects.secondclass.c_instruction import Instruction +from app.contacts.contact_gist import Contact as GistContact +from app.service.contact_svc import ContactService +from app.service.file_svc import FileSvc -from app.contacts.contact_gist import Contact -from app.utility.base_world import BaseWorld +MOCK_GISTS = ''' +[ + { + "url": "https://api.github.com/gists/123", + "forks_url": "https://api.github.com/gists/123/forks", + "commits_url": "https://api.github.com/gists/123/commits", + "id": "123", + "node_id": "testnodeid", + "git_pull_url": "https://gist.github.com/123.git", + "git_push_url": "https://gist.github.com/123.git", + "html_url": "https://gist.github.com/testuser/123", + "files": { + "beacon-566631377468464325": { + "filename": "beacon-566631377468464325", + "type": "text/plain", + "language": null, + "raw_url": "https://gist.githubusercontent.com/testuser/123/raw/260cc70502fbb5ba3c64c8d943ac4cd2927d9a17/beacon-566631377468464325", + "size": 704 + } + }, + "public": false, + "description": "beacon-566631377468464325", + "comments_url": "https://api.github.com/gists/123/comments", + "truncated": false + }, + { + "url": "https://api.github.com/gists/456", + "forks_url": "https://api.github.com/gists/456/forks", + "commits_url": "https://api.github.com/gists/456/commits", + "id": "456", + "node_id": "testnodeid", + "git_pull_url": "https://gist.github.com/456.git", + "git_push_url": "https://gist.github.com/456.git", + "html_url": "https://gist.github.com/testuser/456", + "files": { + "results-566631377468464325": { + "filename": "results-566631377468464325", + "type": "text/plain", + "language": null, + "raw_url": "https://gist.githubusercontent.com/testuser/456/raw/260cc70502fbb5ba3c64c8d943ac4cd2927d9a17/results-566631377468464325", + "size": 704 + } + }, + "public": false, + "description": "results-566631377468464325", + "comments_url": "https://api.github.com/gists/456/comments", + "truncated": false + }, + { + "url": "https://api.github.com/gists/789", + "forks_url": "https://api.github.com/gists/789/forks", + "commits_url": "https://api.github.com/gists/789/commits", + "id": "789", + "node_id": "testnodeid", + "git_pull_url": "https://gist.github.com/789.git", + "git_push_url": "https://gist.github.com/789.git", + "html_url": "https://gist.github.com/testuser/789", + "files": { + "upload-566631377468464325": { + "filename": "upload-566631377468464325", + "type": "text/plain", + "language": null, + "raw_url": "https://gist.githubusercontent.com/testuser/789/raw/260cc70502fbb5ba3c64c8d943ac4cd2927d9a17/upload-566631377468464325", + "size": 704 + } + }, + "public": false, + "description": "upload:123:dGVzdHVwbG9hZA==:1:1", + "comments_url": "https://api.github.com/gists/789/comments", + "truncated": false + } +] +''' @pytest.fixture -def test_tokens_and_expected_results(): - token_old_len39 = '841fa61ea1cc4a2ea616f743bad140bf43bad14' - token_old_len40 = token_old_len39 + '0' - token_old_len255 = 'ad1401234567890' + (token_old_len40 * 6) - token_old_len256 = token_old_len255 + '9' - token_old_badchar = '841fa61ea1cc4a2ea616f74s2A82(*ad1403ri1x' - token_new_len39 = 'ghp_a0s9cu2hFJKjcjisIz3t92w80dulk3rfa35' - token_new_len40 = token_new_len39 + '2' - token_new_len255 = 'ghp_j' + ('a0s9cu2hfkjcjisiz3t92w8g0dulk3rfas35shkk310e0vkqjw' * 5) - token_new_len256 = token_new_len255 + '9' - token_new_badchar = 'ghp_a0s9cu2hfkjcjisiz3t92w8g0dulk3rf' - token_new_bad_suffix = 'gh0_a0s9cu2hfkjcjisiz3t92w8g0dulk3rf' - return [ - (token_old_len39, False), - (token_old_len40, True), - (token_old_len255, True), - (token_old_len256, False), - (token_old_badchar, False), - (token_new_len39, False), - (token_new_len40, True), - (token_new_len255, True), - (token_new_len256, False), - (token_new_badchar, False), - (token_new_bad_suffix, False), - ] +def gist_c2(app_svc, contact_svc, data_svc, file_svc, obfuscator): + services = app_svc.get_services() + return GistContact(services) -@pytest.fixture -def test_contact(app_svc): - return Contact(app_svc.get_services()) +def _mock_get(url, *, allow_redirects=True, **kwargs): + dummy_beacon_data = { + 'architecture': 'arm64', + 'available_contacts': ['HTTP'], + 'contact': 'HTTP', + 'deadman_enabled': True, + 'exe_name': 'splunkd', + 'executors': ['proc', 'sh'], + 'group': 'red', + 'host': 'myhost', + 'host_ip_addrs': ['10.0.2.15'], + 'location': '/home/testuser/splunkd', + 'origin_link_id': '', + 'paw': 'testpaw', + 'pid': 63025, + 'platform': 'linux', + 'ppid': 4357, + 'privilege': 'User', + 'proxy_receivers': None, + 'server': 'http://0.0.0.0:8888', + 'upstream_dest': 'http://0.0.0.0:8888', + 'username': 'testuser' + } + if url == 'https://api.github.com/gists': + mock_data = MOCK_GISTS + elif url.startswith('https://gist.githubusercontent.com'): + if url.endswith('/beacon-566631377468464325'): + mock_data = base64.b64encode(json.dumps(dummy_beacon_data).encode('ascii')) + elif url.endswith('/results-566631377468464325'): + dummy_beacon_data['results'] = 'mockresult' + mock_data = base64.b64encode(json.dumps(dummy_beacon_data).encode('ascii')) + elif url.endswith('/upload-566631377468464325'): + mock_data = base64.b64encode(b'testuploaddata') + else: + mock_data = '' + + mock_resp = mock.Mock( + spec=aiohttp.ClientResponse, + **{'text.return_value': mock_data} + ) + return contextlib.nullcontext(mock_resp) class TestContactGist: - def test_valid_config(self, test_contact, test_tokens_and_expected_results): - for token_and_expected_result in test_tokens_and_expected_results: - token = token_and_expected_result[0] - expected_result = token_and_expected_result[1] - assert test_contact.valid_config(token) == expected_result + @mock.patch.object(uuid, 'uuid4', return_value='mockuuid') + @mock.patch.object(FileSvc, 'read_file', return_value=('testpayload', bytes([0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]))) + @mock.patch.object(aiohttp.ClientSession, 'delete') + @mock.patch.object(aiohttp.ClientSession, 'post') + @mock.patch.object(aiohttp.ClientSession, 'get', side_effect=_mock_get) + async def test_handle_beacons(self, mock_get, mock_post, mock_delete, mock_read_file, mock_uuid, gist_c2): + dummy_agent = Agent(paw='testpaw', sleep_min=5, sleep_max=5, watchdog=0, executors=['sh', 'proc']) + dummy_agent.set_pending_executor_removal('sh') + dummy_agent.pending_contact = 'newcontact' + dummy_instruction = Instruction( + id='123', + sleep=5, + command='whoami', + executor='sh', + timeout=60, + payloads=['testpayload'], + uploads=[], + deadman=False, + delete_payload=True + ) + with mock.patch.object(ContactService, 'handle_heartbeat', return_value=(dummy_agent, [dummy_instruction])): + await gist_c2.handle_beacons(await gist_c2.get_beacons()) + mock_get.assert_any_call('https://api.github.com/gists') + mock_get.assert_any_call('https://gist.githubusercontent.com/testuser/123/raw/260cc70502fbb5ba3c64c8d943ac4cd2927d9a17/beacon-566631377468464325') + want_payload_post_body = dict( + description='payloads-testpaw-testpayload', + public=False, + files=dict(testpayload=dict(content='ASNFZ4mrze8=')) + ) + want_resp = dict( + paw='testpaw', + sleep=5, + watchdog=0, + instructions=json.dumps([json.dumps(dummy_instruction.display)]), + new_contact='newcontact', + executor_change=dict( + action='remove', + executor='sh', + ) + ) + want_resp_encoded = str(base64.b64encode(json.dumps(want_resp).encode('utf-8')), 'utf-8') + want_instruction_post_body = dict( + description='instructions-testpaw', + public=False, + files=dict(mockuuid=dict(content=want_resp_encoded)) + ) + mock_post.assert_any_call('https://api.github.com/gists', json=want_payload_post_body) + mock_post.assert_any_call('https://api.github.com/gists', json=want_instruction_post_body) + + @mock.patch.object(uuid, 'uuid4', return_value='mockuuid') + @mock.patch.object(FileSvc, 'read_file', return_value=('testpayload', bytes([0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]))) + @mock.patch.object(aiohttp.ClientSession, 'delete') + @mock.patch.object(aiohttp.ClientSession, 'post') + @mock.patch.object(aiohttp.ClientSession, 'get', side_effect=_mock_get) + async def test_handle_results(self, mock_get, mock_post, mock_delete, mock_read_file, mock_uuid, gist_c2): + dummy_agent = Agent(paw='testpaw', sleep_min=5, sleep_max=5, watchdog=0, executors=['sh', 'proc']) + with mock.patch.object(ContactService, 'handle_heartbeat', return_value=(dummy_agent, [])): + await gist_c2.handle_beacons(await gist_c2.get_results()) + mock_get.assert_any_call('https://api.github.com/gists') + mock_get.assert_any_call('https://gist.githubusercontent.com/testuser/456/raw/260cc70502fbb5ba3c64c8d943ac4cd2927d9a17/results-566631377468464325') + want_resp = dict( + paw='testpaw', + sleep=5, + watchdog=0, + instructions=json.dumps([]) + ) + want_resp_encoded = str(base64.b64encode(json.dumps(want_resp).encode('utf-8')), 'utf-8') + want_instruction_post_body = dict( + description='instructions-testpaw', + public=False, + files=dict(mockuuid=dict(content=want_resp_encoded)) + ) + mock_post.assert_called_once_with('https://api.github.com/gists', json=want_instruction_post_body) - async def test_retrieve_config(self, app_svc, test_tokens_and_expected_results): - expected_num_op_loops = 0 - with patch.object(Contact, '_start_operation_loop', return_value=None) as start_op_loop: - for token_and_expected_result in test_tokens_and_expected_results: - test_contact = Contact(app_svc.get_services()) - token = token_and_expected_result[0] - expected_result = token_and_expected_result[1] - with patch.object(BaseWorld, 'get_config', return_value=token) as get_config: - await test_contact.start() - assert get_config.call_count == 1 - if expected_result: - expected_num_op_loops += 1 - assert test_contact.retrieve_config() == token - else: - assert not test_contact.retrieve_config() - assert start_op_loop.call_count == expected_num_op_loops + @mock.patch.object(FileSvc, 'save_file') + @mock.patch.object(FileSvc, 'create_exfil_sub_directory', return_value='/tmp/mockcalderaexfil/566631377468464325') + @mock.patch.object(logging.Logger, 'debug') + @mock.patch.object(aiohttp.ClientSession, 'delete') + @mock.patch.object(aiohttp.ClientSession, 'get', side_effect=_mock_get) + async def test_handle_uploads(self, mock_get, mock_delete, mock_logger, mock_create_exfil_dir, mock_save_file, gist_c2): + await gist_c2.handle_uploads(await gist_c2.get_uploads()) + mock_logger.assert_any_call( + 'Received uploaded file chunk 1 out of 1 for paw 566631377468464325, upload ID 123, filename testupload ' + ) + mock_logger.assert_any_call('Upload 123 complete for paw 566631377468464325, filename testupload') + mock_logger.assert_any_call('Uploaded file /tmp/mockcalderaexfil/566631377468464325/testupload') + mock_create_exfil_dir.assert_called_once_with(dir_name='566631377468464325') + mock_save_file.assert_called_once_with('testupload-123', b'testuploaddata', '/tmp/mockcalderaexfil/566631377468464325') diff --git a/tests/contacts/test_contact_html.py b/tests/contacts/test_contact_html.py new file mode 100644 index 000000000..c17183922 --- /dev/null +++ b/tests/contacts/test_contact_html.py @@ -0,0 +1,91 @@ +import pytest +import base64 +import json +from unittest import mock + +from app.objects.c_agent import Agent +from app.contacts.contact_html import Contact as HTMLContact +from app.service.contact_svc import ContactService +from app.utility.base_world import BaseWorld + + +@pytest.fixture(scope='session') +def html_contact_base_world(): + BaseWorld.apply_config(name='main', config={'app.contact.http': 'http://0.0.0.0:8888', + 'host': '0.0.0.0', + 'port': '8888', + 'crypt_salt': 'BLAH', + 'api_key': 'ADMIN123', + 'encryption_key': 'ADMIN123', + 'exfil_dir': '/tmp'}) + BaseWorld.apply_config(name='agents', config={'sleep_max': 5, + 'sleep_min': 5, + 'untrusted_timer': 90, + 'watchdog': 0, + 'implant_name': 'splunkd', + 'bootstrap_abilities': [ + '43b3754c-def4-4699-a673-1d85648fda6a' + ]}) + + +@pytest.fixture +def html_c2(app_svc, contact_svc, data_svc, file_svc, obfuscator): + services = app_svc.get_services() + html_c2 = HTMLContact(services) + return html_c2 + + +@pytest.fixture +def encode_dict(): + def _encode_dict(input): + return bytes(base64.b64encode(json.dumps(input).encode('utf-8'))) + return _encode_dict + + +class _MockRequest(): + def __init__(self, data): + self._data = data + + async def text(self): + return self._data + + +@pytest.mark.usefixtures( + 'html_contact_base_world' +) +class TestContactHTML: + async def test_handler_beacon(self, aiohttp_client, html_c2, encode_dict): + dummy_beacon_data = { + 'architecture': 'arm64', + 'available_contacts': ['html'], + 'contact': 'html', + 'deadman_enabled': True, + 'exe_name': 'splunkd', + 'executors': ['proc', 'sh'], + 'group': 'red', + 'host': 'myhost', + 'host_ip_addrs': ['10.0.2.15'], + 'location': '/home/testuser/splunkd', + 'origin_link_id': '', + 'paw': 'testpaw', + 'pid': 63025, + 'platform': 'linux', + 'ppid': 4357, + 'privilege': 'User', + 'proxy_receivers': None, + 'server': 'http://0.0.0.0:8888', + 'upstream_dest': 'http://0.0.0.0:8888', + 'username': 'testuser' + } + encoded = encode_dict(dummy_beacon_data) + dummy_agent = Agent(paw=dummy_beacon_data.get('paw'), sleep_min=5, sleep_max=5, watchdog=0, executors=['sh', 'proc']) + mock_request = _MockRequest(encoded) + with mock.patch.object(ContactService, 'handle_heartbeat', return_value=(dummy_agent, [])): + inner_want_dict = dict( + paw='testpaw', + sleep=5, + watchdog=0, + instructions="[]", + ) + want_response_dict = dict(instructions=html_c2.encode_string(json.dumps(inner_want_dict))) + assert await html_c2._beacon_helper(mock_request) == want_response_dict diff --git a/tests/contacts/test_contact_http.py b/tests/contacts/test_contact_http.py new file mode 100644 index 000000000..a93ce1290 --- /dev/null +++ b/tests/contacts/test_contact_http.py @@ -0,0 +1,125 @@ +import pytest +import base64 +import json +from http import HTTPStatus +from unittest import mock + +from app.objects.c_agent import Agent +from app.objects.secondclass.c_instruction import Instruction +from app.contacts.contact_http import Contact as HTTPContact +from app.service.contact_svc import ContactService +from app.utility.base_world import BaseWorld + + +@pytest.fixture(scope='session') +def http_contact_base_world(): + BaseWorld.apply_config(name='main', config={'app.contact.http': 'http://0.0.0.0:8888', + 'host': '0.0.0.0', + 'port': '8888', + 'crypt_salt': 'BLAH', + 'api_key': 'ADMIN123', + 'encryption_key': 'ADMIN123', + 'exfil_dir': '/tmp'}) + BaseWorld.apply_config(name='agents', config={'sleep_max': 5, + 'sleep_min': 5, + 'untrusted_timer': 90, + 'watchdog': 0, + 'implant_name': 'splunkd', + 'bootstrap_abilities': [ + '43b3754c-def4-4699-a673-1d85648fda6a' + ]}) + + +@pytest.fixture +def http_c2(app_svc, contact_svc, data_svc, file_svc, obfuscator): + services = app_svc.get_services() + http_c2 = HTTPContact(services) + return http_c2 + + +@pytest.fixture +def encode_beacon_dict(): + def _encode_beacon_dict(input): + return bytes(base64.b64encode(json.dumps(input).encode('ascii'))) + return _encode_beacon_dict + + +class _MockRequest(): + def __init__(self, data): + self._data = data + + async def read(self): + return self._data + + +@pytest.mark.usefixtures( + 'http_contact_base_world' +) +class TestContactHTTP: + async def test_handler_beacon(self, http_c2, encode_beacon_dict): + dummy_beacon_data = { + 'architecture': 'arm64', + 'available_contacts': ['HTTP'], + 'contact': 'HTTP', + 'deadman_enabled': True, + 'exe_name': 'splunkd', + 'executors': ['proc', 'sh'], + 'group': 'red', + 'host': 'myhost', + 'host_ip_addrs': ['10.0.2.15'], + 'location': '/home/testuser/splunkd', + 'origin_link_id': '', + 'paw': 'testpaw', + 'pid': 63025, + 'platform': 'linux', + 'ppid': 4357, + 'privilege': 'User', + 'proxy_receivers': None, + 'server': 'http://0.0.0.0:8888', + 'upstream_dest': 'http://0.0.0.0:8888', + 'username': 'testuser' + } + encoded = encode_beacon_dict(dummy_beacon_data) + dummy_agent = Agent(paw=dummy_beacon_data.get('paw'), sleep_min=5, sleep_max=5, watchdog=0, executors=['sh', 'proc']) + dummy_agent.set_pending_executor_removal('sh') + dummy_agent.pending_contact = 'newcontact' + dummy_instruction = Instruction( + id='123', + sleep=5, + command='whoami', + executor='sh', + timeout=60, + payloads=[], + uploads=[], + deadman=False, + delete_payload=True + ) + mock_request = _MockRequest(encoded) + with mock.patch.object(ContactService, 'handle_heartbeat', return_value=(dummy_agent, [dummy_instruction])): + want_response_dict = dict( + paw='testpaw', + sleep=5, + watchdog=0, + instructions=[dict( + id='123', + sleep=5, + command='whoami', + executor='sh', + timeout=60, + payloads=[], + uploads=[], + deadman=False, + delete_payload=True + )], + new_contact='newcontact', + executor_change=dict( + action='remove', + executor='sh', + ) + ) + resp = await http_c2._beacon(mock_request) + assert resp.status == HTTPStatus.OK + decoded = base64.b64decode(resp.text) + resp_dict = json.loads(decoded) + resp_dict['instructions'] = [json.loads(a) for a in json.loads(resp_dict['instructions'])] + assert resp_dict == want_response_dict diff --git a/tests/contacts/test_contact_slack.py b/tests/contacts/test_contact_slack.py new file mode 100644 index 000000000..96371cdc7 --- /dev/null +++ b/tests/contacts/test_contact_slack.py @@ -0,0 +1,109 @@ +import pytest +import aiohttp +import base64 +import contextlib +import json +from unittest import mock + +from app.objects.c_agent import Agent +from app.objects.secondclass.c_instruction import Instruction +from app.contacts.contact_slack import Contact as SlackContact +from app.service.contact_svc import ContactService +from app.service.file_svc import FileSvc + +MOCK_SLACK = '''{ + "messages": [ + { + "text": "beacon | %s", + "ts": "dummytimestamp" + } + ] +} +''' + + +@pytest.fixture +def slack_c2(app_svc, contact_svc, data_svc, file_svc, obfuscator): + services = app_svc.get_services() + slack_contact = SlackContact(services) + slack_contact.channelid = 'mockchannelid' + yield slack_contact + + +def _mock_get(url, *, allow_redirects=True, **kwargs): + dummy_beacon_data = { + 'architecture': 'arm64', + 'available_contacts': ['HTTP'], + 'contact': 'HTTP', + 'deadman_enabled': True, + 'exe_name': 'splunkd', + 'executors': ['proc', 'sh'], + 'group': 'red', + 'host': 'myhost', + 'host_ip_addrs': ['10.0.2.15'], + 'location': '/home/testuser/splunkd', + 'origin_link_id': '', + 'paw': 'testpaw', + 'pid': 63025, + 'platform': 'linux', + 'ppid': 4357, + 'privilege': 'User', + 'proxy_receivers': None, + 'server': 'http://0.0.0.0:8888', + 'upstream_dest': 'http://0.0.0.0:8888', + 'username': 'testuser' + } + if url.startswith('https://slack.com/api/conversations.history?channel=mockchannelid&oldest='): + encoded_beacon = str(base64.b64encode(json.dumps(dummy_beacon_data).encode('utf-8')), 'utf-8') + mock_data = MOCK_SLACK % encoded_beacon + else: + mock_data = '' + + mock_resp = mock.Mock( + spec=aiohttp.ClientResponse, + **{'text.return_value': mock_data} + ) + return contextlib.nullcontext(mock_resp) + + +class TestContactSlack: + @mock.patch.object(FileSvc, 'read_file', return_value=('testpayload', bytes([0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]))) + @mock.patch.object(aiohttp.ClientSession, 'post') + @mock.patch.object(aiohttp.ClientSession, 'get', side_effect=_mock_get) + async def test_handle_beacons(self, mock_get, mock_post, mock_read_file, slack_c2): + dummy_agent = Agent(paw='testpaw', sleep_min=5, sleep_max=5, watchdog=0, executors=['sh', 'proc']) + dummy_agent.pending_contact = 'newcontact' + dummy_instruction = Instruction( + id='123', + sleep=5, + command='whoami', + executor='sh', + timeout=60, + payloads=['testpayload'], + uploads=[], + deadman=False, + delete_payload=True + ) + with mock.patch.object(ContactService, 'handle_heartbeat', return_value=(dummy_agent, [dummy_instruction])): + await slack_c2.handle_beacons(await slack_c2.get_beacons()) + want_payload_post_body = dict( + channels='mockchannelid', + initial_comment='payloads-testpaw-testpayload', + content='ASNFZ4mrze8=' + ) + want_resp = dict( + paw='testpaw', + sleep=5, + watchdog=0, + instructions=json.dumps([json.dumps(dummy_instruction.display)]), + new_contact='newcontact' + ) + want_resp_encoded = str(base64.b64encode(json.dumps(want_resp).encode('utf-8')), 'utf-8') + want_instruction_post_body = dict( + channel='mockchannelid', + text=f'instructions-testpaw | {want_resp_encoded}' + ) + want_delete_dict = {'channel': 'mockchannelid', 'ts': 'dummytimestamp'} + mock_post.assert_any_call('https://slack.com/api/chat.delete', data=want_delete_dict) + mock_post.assert_any_call('https://slack.com/api/files.upload', data=want_payload_post_body) + mock_post.assert_any_call('https://slack.com/api/chat.postMessage', json=want_instruction_post_body) diff --git a/tests/objects/secondclass/test_executor.py b/tests/objects/secondclass/test_executor.py new file mode 100644 index 000000000..c071236ab --- /dev/null +++ b/tests/objects/secondclass/test_executor.py @@ -0,0 +1,47 @@ +import pytest + +from app.objects.secondclass.c_executor import Executor, get_variations +from app.objects.secondclass.c_variation import Variation +from app.utility.base_world import BaseWorld + + +@pytest.fixture(scope='session') +def executors_base_world(): + BaseWorld.apply_config(name='main', config={'app.contact.http': 'http://0.0.0.0:8888', + 'host': '0.0.0.0', + 'port': '8888', + 'crypt_salt': 'BLAH', + 'api_key': 'ADMIN123', + 'encryption_key': 'ADMIN123', + 'exfil_dir': '/tmp'}) + + +@pytest.fixture +def test_executor(executor): + return executor(name='sh', platform='linux') + + +@pytest.mark.usefixtures( + 'executors_base_world' +) +class TestExecutor: + def test_is_global_variable(self): + assert not Executor.is_global_variable('notaglobalvariable') + assert Executor.is_global_variable('payload') + + def test_replace_cleanup(self, test_executor): + assert 'this has been replaced' == test_executor.replace_cleanup('this has been replaced', 'somepayload') + assert 'this has been replaced: somepayload' == test_executor.replace_cleanup('this has been replaced: #{payload}', 'somepayload') + + def test_get_variations(self): + assert [] == get_variations([]) + want_variation = Variation('test description', 'test command') + second_variation = Variation('test description', 'test command') + diff_variation = Variation('test description 2', 'test command 2') + assert want_variation == second_variation + assert not diff_variation == want_variation + var_dict = dict(description='test description', command='test command') + assert [want_variation] == get_variations([var_dict]) + result = get_variations([var_dict, want_variation]) + second_want = Variation('test description', Variation.encode_string('test command')) + assert [want_variation, second_want] == result diff --git a/tests/objects/test_agent.py b/tests/objects/test_agent.py index 220ffd52d..43bcf0408 100644 --- a/tests/objects/test_agent.py +++ b/tests/objects/test_agent.py @@ -1,4 +1,5 @@ from base64 import b64decode +from unittest import mock from datetime import timedelta from app.objects.c_ability import Ability @@ -128,6 +129,13 @@ def test_heartbeat_modification_during_pending_executor_removal(self, event_loop event_loop.run_until_complete(agent.heartbeat_modification(executors=original_executors)) assert agent.executors == ['cmd'] + def test_heartbeat_modification_executor(self, event_loop): + original_executors = ['cmd'] + new_executors = ['cmd', 'test'] + agent = Agent(paw='123', sleep_min=2, sleep_max=8, watchdog=0, executors=original_executors, platform='windows') + event_loop.run_until_complete(agent.heartbeat_modification(executors=new_executors)) + assert agent.executors == new_executors + def test_status_and_kill(self, event_loop, mocker, mock_time): BaseWorld.set_config(name='agents', prop='untrusted_timer', value=30) agent = Agent(paw='123', sleep_min=2, sleep_max=8, watchdog=0, executors=['cmd'], platform='windows') @@ -185,3 +193,58 @@ def test_store_existing_agent(self, data_svc): stored_agent = agent.store(data_svc.ram) assert len(data_svc.ram['agents']) == 1 assert stored_agent.schema.dump(stored_agent) == agent.schema.dump(agent) + + def test_is_global_variable(self): + assert Agent.is_global_variable('payload:thisshoudlreturntrue') + assert not Agent.is_global_variable('payload') + assert Agent.is_global_variable('server') + assert Agent.is_global_variable('group') + assert not Agent.is_global_variable('default') + + async def test_get_preferred_executor_empty(self, ability, executor): + agent = Agent(paw='123', sleep_min=5, group='red', sleep_max=5, watchdog=2, executors=['cmd'], + platform='windows', trusted=True) + + # No overlap between agent platform/executor and ability platform/executor + test_executor = executor(name='sh', platform='linux') + test_ability = ability(ability_id='123', executors=[test_executor]) + assert await agent.get_preferred_executor(test_ability) is None + + async def test_bootstrap(self, ability, executor, data_svc): + # Set empty executor list to exit agent.task function immmediately without mocking + agent = Agent(paw='123', sleep_min=5, group='red', sleep_max=5, watchdog=0, executors=[], + platform='windows', trusted=True) + test_executor = executor(name='sh', platform='linux') + test_ability = ability(ability_id='bootstrap-id', executors=[test_executor], name='Test name') + await data_svc.store(test_ability) + BaseWorld.apply_config(name='agents', + config={'sleep_max': 5, + 'sleep_min': 5, + 'untrusted_timer': 90, + 'watchdog': 0, + 'implant_name': 'splunkd', + 'bootstrap_abilities': ['bootstrap-id']}) + await agent.bootstrap(data_svc) + + # Test with mock + with mock.patch.object(Agent, 'task', return_value=[]) as mock_task: + await agent.bootstrap(data_svc) + mock_task.assert_called_once_with([test_ability], obfuscator='plain-text') + + async def test_deadman(self, ability, executor, data_svc): + agent = Agent(paw='123', sleep_min=5, group='red', sleep_max=5, watchdog=0, executors=[], + platform='windows', trusted=True) + test_executor = executor(name='sh', platform='linux') + test_ability = ability(ability_id='deadman-id', executors=[test_executor], name='Test name') + await data_svc.store(test_ability) + BaseWorld.apply_config(name='agents', + config={'sleep_max': 5, + 'sleep_min': 5, + 'untrusted_timer': 90, + 'watchdog': 0, + 'implant_name': 'splunkd', + 'deadman_abilities': ['deadman-id']}) + + with mock.patch.object(Agent, 'task', return_value=[]) as mock_task: + await agent.deadman(data_svc) + mock_task.assert_called_once_with([test_ability], obfuscator='plain-text', deadman=True) diff --git a/tests/services/test_app_svc.py b/tests/services/test_app_svc.py index 5cf2af066..bcc1d89e5 100644 --- a/tests/services/test_app_svc.py +++ b/tests/services/test_app_svc.py @@ -1,6 +1,17 @@ import pytest - +import asyncio +import logging +import subprocess +from datetime import datetime, timedelta, timezone +from unittest import mock + +from app.objects.c_agent import Agent +from app.service.app_svc import AppService +from app.service.contact_svc import ContactService +from app.service.data_svc import DataService +from app.service.file_svc import FileSvc from app.utility.base_service import BaseService +from app.utility.base_world import BaseWorld @pytest.fixture @@ -129,3 +140,77 @@ async def test_mark_agent_as_untrusted_cleanup_operation(self, setup_cleanup_ope assert operation.state == operation.States.CLEANUP.value assert test_agent in operation.agents assert not operation.untrusted_agents + + async def test_validate_requirements(self, app_svc): + reqs = dict( + go=dict( + command='go version', + type='installed_program', + version='1.24', + ), + python=dict( + attr='version', + module='sys', + type='python_module', + version='3.9.0' + ) + ) + BaseWorld.set_config('main', 'requirements', reqs) + + # Test success + with mock.patch.object(subprocess, 'check_output', return_value=b'go version go1.25.5 linux/arm64\n'): + await app_svc.validate_requirements() + for req, param in reqs.items(): + assert await app_svc.validate_requirement(req, param) + + # Test failure due to obsolete version + with mock.patch.object(subprocess, 'check_output', return_value=b'go version go1.19 linux/arm64\n'): + assert not await app_svc.validate_requirement('go', {'command': 'go version', 'type': 'installed_program', 'version': '1.24'}) + assert not await app_svc.validate_requirement('go', {'command': 'go version', 'type': 'installed_program', 'version': '1.24', 'optional': True}) + + # Test failure due to unknown version + with mock.patch.object(subprocess, 'check_output', return_value=b'go version X linux/arm64\n'): + assert not await app_svc.validate_requirement('go', {'command': 'go version', 'type': 'installed_program', 'version': '1.24'}) + + # Test FileNotFoundError due to bad command + BaseWorld.set_config('main', 'requirements', dict( + go=dict( + command='thiscommanddoesnotexist', + type='installed_program', + version='1.24', + ) + )) + assert not await app_svc.validate_requirement('go', {'command': 'thiscommanddoesnotexist', 'type': 'installed_program', 'version': '1.24'}) + + # Test Exception + with mock.patch.object(subprocess, 'check_output') as mock_check_output: + mock_check_output.side_effect = Exception('testexception') + assert not await app_svc.validate_requirement('go', {'command': 'go version', 'type': 'installed_program', 'version': '1.24'}) + + async def test_start_sniffer_untrusted_agents(self, app_svc): + trusted_agent = Agent(paw='test', trusted=True, sleep_max=1) + untrusted_agent = Agent(paw='test', trusted=True, sleep_max=1) + start_time = datetime.now(timezone.utc) + trusted_agent.last_trusted_seen = start_time + untrusted_agent.last_trusted_seen = start_time - timedelta(0, 30) + with mock.patch.object(asyncio, 'sleep'): + with mock.patch.object(AppService, 'get_config', return_value=10): + with mock.patch.object(AppService, 'update_operations_with_untrusted_agent') as mock_update_ops: + mock_update_ops.side_effect = Exception('terminate test_start_sniffer_untrusted_agents') + with mock.patch.object(DataService, 'locate', return_value=[trusted_agent, untrusted_agent]): + await app_svc.start_sniffer_untrusted_agents() + mock_update_ops.assert_called_once_with(untrusted_agent) + assert not untrusted_agent.trusted + assert trusted_agent.trusted + + @mock.patch.object(ContactService, 'build_filename', return_value='mock_filename') + @mock.patch.object(logging.Logger, 'debug') + @mock.patch.object(FileSvc, 'find_file_path', return_value=(None, 'mock_file')) + async def test_retrieve_compiled_file(self, mock_find_filepath, mock_debug, mock_build_filename, app_svc): + file_data = bytes([0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]) + m = mock.mock_open(read_data=file_data) + with mock.patch('builtins.open', m): + result = await app_svc.retrieve_compiled_file('test_file', 'windows') + assert result == ('test_file-windows', 'mock_filename') + mock_find_filepath.assert_called_once_with('test_file-windows', location='') + mock_debug.assert_called_once_with('test_file downloaded with hash=55c53f5d490297900cefa825d0c8e8e9532ee8a118abe7d8570762cd38be9818 and name=mock_filename') diff --git a/tests/services/test_data_svc.py b/tests/services/test_data_svc.py index f3b3fde16..af989fa8f 100644 --- a/tests/services/test_data_svc.py +++ b/tests/services/test_data_svc.py @@ -1,5 +1,11 @@ +import pytest + +import asyncio +import copy import glob import json +import os +import shutil import logging import yaml @@ -9,11 +15,14 @@ from app.objects.c_ability import Ability from app.objects.c_adversary import Adversary from app.objects.c_agent import Agent +from app.objects.c_objective import Objective from app.objects.c_operation import Operation from app.objects.c_planner import Planner from app.objects.c_plugin import Plugin +from app.objects.c_source import Source from app.objects.secondclass.c_executor import Executor from app.service.data_svc import DataService +from app.service.file_svc import FileSvc from app.utility.base_world import BaseWorld @@ -78,6 +87,42 @@ ABILITY_YAMLS = { + 'TESTV0ABILITYID': [yaml.safe_load(''' +- id: TESTV0ABILITYID + name: Account-type Admin Enumerator + description: Use PowerView to query the Active Directory server to determine remote admins + tactic: discovery + technique: + attack_id: T1069.002 + name: "Permission Groups Discovery: Domain Groups" + platforms: + windows: + psh: + command: | + testcommand + parsers: + plugins.stockpile.app.parsers.netlocalgroup: + - source: remote.host.fqdn + edge: has_admin + target: domain.user.name + payloads: + - powerview.ps1 + singleton: True + requirements: + - plugins.stockpile.app.requirements.not_exists: + - source: remote.host.fqdn + edge: has_admin + - plugins.stockpile.app.requirements.basic: + - source: backup.admin.ability + edge: first_failed + - plugins.stockpile.app.requirements.basic: + - source: domain.user.name + edge: has_password + target: domain.user.password + - plugins.stockpile.app.requirements.reachable: + - source: remote.host.fqdn + edge: isAccessibleFrom +''')], 'plugins/testing/data/discovery/764efa883dda1e11db47671c4a3bbd9e.yml': [yaml.safe_load(''' --- @@ -152,6 +197,25 @@ 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) + return mock_future + + +@pytest.fixture +def dir_to_delete(): + to_delete = '/tmp/caldera_test_delete_dir' + os.makedirs(to_delete) + yield to_delete + + if os.path.exists(to_delete): + try: + shutil.rmtree(to_delete) + except OSError: + pass + + class TestDataService: mock_payload_config = dict() @@ -307,7 +371,132 @@ 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_v0_ability_file(self, mock_strip_yml, data_svc): + await data_svc.load_ability_file('TESTV0ABILITYID', data_svc.Access.RED) + search_results = await data_svc.locate('abilities', match=dict(ability_id='TESTV0ABILITYID')) + assert len(search_results) == 1 + result = search_results[0] + assert result.name == 'Account-type Admin Enumerator' + assert result.description == 'Use PowerView to query the Active Directory server to determine remote admins' + assert result.tactic == 'discovery' + assert result.technique_name == 'Permission Groups Discovery: Domain Groups' + assert result.technique_id == 'T1069.002' + assert result.singleton + result_executor = result.find_executor('psh', 'windows') + assert result_executor.command == 'testcommand' + assert result_executor.payloads == ['powerview.ps1'] + assert len(result_executor.parsers) == 1 + result_parser = result_executor.parsers[0] + assert result_parser.unique == 'plugins.stockpile.app.parsers.netlocalgroup' + assert len(result_parser.parserconfigs) == 1 + result_parser_cfg = result_parser.parserconfigs[0] + assert result_parser_cfg.source == 'remote.host.fqdn' + assert result_parser_cfg.edge == 'has_admin' + assert result_parser_cfg.target == 'domain.user.name' + assert not result_parser_cfg.custom_parser_vals + assert len(result.requirements) == 4 + + def test_delete_file(self, dir_to_delete): + to_delete = os.path.join(dir_to_delete, 'caldera_test_delete_file') + open(to_delete, 'a').close() + assert os.path.exists(to_delete) + DataService._delete_file(to_delete) + assert not os.path.exists(to_delete) + DataService._delete_file(to_delete) # no-op delete + + assert os.path.exists(dir_to_delete) + DataService._delete_file(dir_to_delete) + assert not os.path.exists(dir_to_delete) + + async def test_destroy(self, data_svc, dir_to_delete): + def mock_glob(input): + if 'abilities' in input: + return ['/tmp/caldera_test_delete_dir/calderadummyfile.txt'] + else: + return [] + + with open('/tmp/caldera_test_delete_dir/calderadummyfile.txt', 'w') as input_file: + input_file.write('test') + + with mock.patch.object(glob, 'glob', side_effect=mock_glob): + with mock.patch.object(os.path, 'join', return_value='/tmp/caldera_test_delete_dir/test.tar.gz'): + with mock.patch.object(os, 'mkdir', return_value=True): + await data_svc.destroy() + assert not os.path.exists('/tmp/caldera_test_delete_dir/calderadummyfile.txt') + assert os.path.exists('/tmp/caldera_test_delete_dir/test.tar.gz') + + @mock.patch.object(os.path, 'exists', return_value=True) + async def test_save_and_restore_state(self, mock_exists, data_svc, dir_to_delete, ability): + test_abil = ability(ability_id='testsaverestoreabil', name='testsaverestoreabil') + await data_svc.store(test_abil) + with mock.patch.object(os.path, 'join', return_value='/tmp/caldera_test_delete_dir/object_store'): + with mock.patch.object(FileSvc, 'find_file_path', return_value=(None, '/tmp/caldera_test_delete_dir/object_store')): + await data_svc.save_state() + assert os.path.exists('/tmp/caldera_test_delete_dir/object_store') + assert await data_svc.locate('abilities', match=dict(ability_id='testsaverestoreabil')) + await data_svc.remove('abilities', match=dict(ability_id='testsaverestoreabil')) + assert not await data_svc.locate('abilities', match=dict(ability_id='testsaverestoreabil')) + await data_svc.restore_state() + assert await data_svc.locate('abilities', match=dict(ability_id='testsaverestoreabil')) + + async def test_search_object(self, data_svc, ability): + test_abil = ability(ability_id='tosearchfor', name='testsaverestoreabil', tags=['tag1', 'tag2']) + await data_svc.store(test_abil) + result = await data_svc.search('tag1', 'abilities') + assert len(result) == 1 + assert result[0] == test_abil + result = await data_svc.search('tag2', 'abilities') + assert len(result) == 1 + assert result[0] == test_abil + result = await data_svc.search('dne', 'abilities') + assert not result + + @mock.patch.object(DataService, '_update_payload_config') + @mock.patch.object(DataService, 'create_or_update_everything_adversary') + @mock.patch.object(FileSvc, 'add_special_payload') + @mock.patch.object(DataService, 'load_yaml_file') + @mock.patch.object(DataService, 'load_ability_file') + async def test_reload_data(self, mock_load_abil_file, mock_load_yaml_file, mock_add_special_payload, mock_everything_adversary, mock_update_payload_config, data_svc): + def _mock_iglob(pathname, *, root_dir=None, dir_fd=None, recursive=False, include_hidden=False): + if 'payloads' in pathname: + return iter([]) + elif 'abilities' in pathname: + return iter(['mockabilityforplugin.yml']) + elif 'objectives' in pathname: + return iter(['mockobjectiveforplugin.yml']) + elif 'adversaries' in pathname: + return iter(['mockadversaryforplugin.yml']) + elif 'planners' in pathname: + return iter(['mockplannerforplugin.yml']) + elif 'sources' in pathname: + return iter(['mocksourceforplugin.yml']) + elif 'packers' in pathname: + return iter([]) + elif 'data_encoders' in pathname: + return iter([]) + else: + return glob.iglob(pathname, root_dir=root_dir, dir_fd=dir_fd, recursive=recursive, include_hidden=include_hidden) + + BaseWorld.apply_config(name='payloads', + config=dict(extensions=dict(mockextension='mockextval'))) + + mock_plugin = Plugin(name='mockplugin', description='Mock plugin for unit tests', enabled=True, data_dir='mockplugin/data', + access=data_svc.Access.RED) + + data_svc.ram = copy.deepcopy(data_svc.schema) + + with mock.patch.object(glob, 'iglob', side_effect=_mock_iglob) as mock_iglob: + await data_svc.reload_data([mock_plugin]) + mock_iglob.assert_any_call('mockplugin/data/payloads/*.yml', recursive=False) + mock_load_abil_file.assert_any_call('mockabilityforplugin.yml', data_svc.Access.RED) + mock_load_yaml_file.assert_any_call(Objective, 'mockobjectiveforplugin.yml', data_svc.Access.RED) + mock_load_yaml_file.assert_any_call(Adversary, 'mockadversaryforplugin.yml', data_svc.Access.RED) + mock_load_yaml_file.assert_any_call(Planner, 'mockplannerforplugin.yml', data_svc.Access.RED) + mock_load_yaml_file.assert_any_call(Source, 'mocksourceforplugin.yml', data_svc.Access.RED) + mock_add_special_payload.assert_called_once_with('mockextension', 'mockextval') + + @mock.patch.object(logging.Logger, 'warning') @mock.patch.object(BaseWorld, 'strip_yml', wraps=strip_ability_yaml) async def test_load_ability_file(self, mock_strip_yml, mock_warn, data_svc): want_executors = [ diff --git a/tests/services/test_file_svc.py b/tests/services/test_file_svc.py index 97080d8e1..31c695d48 100644 --- a/tests/services/test_file_svc.py +++ b/tests/services/test_file_svc.py @@ -1,16 +1,20 @@ import base64 +import copy import json import os import pytest +import subprocess import yaml from aiohttp import web from base64 import b64encode from asyncio import Future +from unittest import mock from unittest.mock import AsyncMock, MagicMock from app.data_encoders.base64_basic import Base64Encoder from app.data_encoders.plain_text import PlainTextEncoder +from app.utility.base_world import BaseWorld from app.utility.file_decryptor import decrypt @@ -403,3 +407,31 @@ def _test_upload_file_with_encoding(event_loop, file_svc, data_svc, encoding, up assert written_data == decoded_content os.remove(uploaded_file_path) os.rmdir(upload_dir) + + @mock.patch.object(subprocess, 'check_output', return_value=b'mock output') + async def test_compile_go(self, mock_check_output, file_svc): + test_env = copy.copy(os.environ) + test_env['GOARCH'] = 'testarch' + test_env['GOOS'] = 'testplatform' + test_env['testcflag'] = 'testval' + want_args = ['go', 'build', 'testmode', '-ldflags', '-s -w', '-o', 'testoutput', 'testsrc.go'] + await file_svc.compile_go('testplatform', 'testoutput', 'testsrc.go', arch='testarch', + cflags='testcflag=testval', buildmode='testmode') + mock_check_output.assert_called_once_with(want_args, cwd='.', env=test_env) + + def test_get_payload_name_from_uuid(self, file_svc): + BaseWorld.set_config('payloads', 'standard_payloads', { + 'Akagi64.exe': { + 'description': 'UACME compiled binary', + 'id': 'testakagiid' + }, + 'wifi.ps1': { + 'description': 'Wifi manipulation script', + 'id': 'testwifiid', + 'obfuscation_name': ['obfuscated_payload.ps1'] + } + }) + BaseWorld.set_config('payloads', 'special_payloads', {}) + assert file_svc.get_payload_name_from_uuid('testakagiid') == ('Akagi64.exe', 'Akagi64.exe') + assert file_svc.get_payload_name_from_uuid('testwifiid') == ('wifi.ps1', 'obfuscated_payload.ps1') + assert file_svc.get_payload_name_from_uuid('DNE') == ('DNE', 'DNE') diff --git a/tests/services/test_learning_svc.py b/tests/services/test_learning_svc.py index b467c8526..403a96029 100644 --- a/tests/services/test_learning_svc.py +++ b/tests/services/test_learning_svc.py @@ -1,7 +1,10 @@ import pytest +from unittest import mock + from app.objects.c_adversary import Adversary from app.objects.secondclass.c_executor import Executor from app.objects.secondclass.c_fact import Fact +from app.service.data_svc import DataService from app.utility.base_world import BaseWorld @@ -71,3 +74,13 @@ async def test_build_relationships(self, setup_learning_service, learning_svc, k ] await learning_svc._store_results(link, facts) assert len(link.relationships) == 4 + + async def test_build_model(self, learning_svc, ability, executor, data_svc): + test_executor = executor(name='sh', platform='linux', command='test command #{var1} #{var2}') + test_ability = ability(ability_id='123123', executors=[test_executor]) + want = set() + want.add(frozenset(['var1', 'var2'])) + learning_svc.model = set() + with mock.patch.object(DataService, 'locate', return_value=[test_ability]): + await learning_svc.build_model() + assert want == learning_svc.model diff --git a/tests/services/test_rest_svc.py b/tests/services/test_rest_svc.py index 749e9f74b..91e72222a 100644 --- a/tests/services/test_rest_svc.py +++ b/tests/services/test_rest_svc.py @@ -1,3 +1,4 @@ +import aiohttp import base64 import pytest @@ -390,26 +391,27 @@ def make_request(file_path): req = MagicMock() req.query = {'file': encoded} # Simulate a real aiohttp Request so check_authorization passes the type check - from aiohttp.web_request import Request - req.__class__ = Request + req.__class__ = aiohttp.web_request.Request return req # Sibling directory /tmp/caldera2/evil must be rejected (HTTPNotFound) sibling_request = make_request('/tmp/caldera2/evil') - sibling_response = event_loop.run_until_complete( - rest_api.download_exfil_file(sibling_request) - ) - assert sibling_response.status == 404, ( + with pytest.raises(aiohttp.web_exceptions.HTTPNotFound) as exc_info: + sibling_response = event_loop.run_until_complete( + rest_api.download_exfil_file(sibling_request) + ) + assert exc_info.value.args[0] == "Not Found", ( "Sibling directory /tmp/caldera2/evil must be rejected with 404, got %s" % sibling_response.status ) # Path /tmp/calderaevil must also be rejected (starts-with prefix match must require sep) prefix_request = make_request('/tmp/calderaevil') - prefix_response = event_loop.run_until_complete( - rest_api.download_exfil_file(prefix_request) - ) - assert prefix_response.status == 404, ( - "Sibling path /tmp/calderaevil must be rejected with 404, got %s" % prefix_response.status + with pytest.raises(aiohttp.web_exceptions.HTTPNotFound) as exc_info: + prefix_response = event_loop.run_until_complete( + rest_api.download_exfil_file(prefix_request) + ) + assert exc_info.value.args[0] == "Not Found", ( + "Sibling directory /tmp/caldera2/evil must be rejected with 404, got %s" % prefix_response.status ) # A legitimate path inside exfil_dir must be accepted (file_svc.read_file is called) diff --git a/tests/test_banner.py b/tests/test_banner.py new file mode 100644 index 000000000..edb5b158c --- /dev/null +++ b/tests/test_banner.py @@ -0,0 +1,51 @@ +import os + +from unittest import mock +from unittest.mock import patch + +from app.ascii_banner import get_ascii_banner, print_rich_banner + +WANT_BANNER = """ + ██████╗ █████╗ ██╗ ██████╗ ███████╗██████╗ █████╗ +██╔════╝██╔══██╗██║ ██╔══██╗██╔════╝██╔══██╗██╔══██╗ +██║ ███████║██║ ██║ ██║█████╗ ██████╔╝███████║ +██║ ██╔══██║██║ ██║ ██║██╔══╝ ██╔══██╗██╔══██║ +╚██████╗██║ ██║███████╗██████╔╝███████╗██║ ██║██║ ██║ + ╚═════╝╚═╝ ╚═╝╚══════╝╚═════╝ ╚══════╝╚═╝ ╚═╝╚═╝ ╚═╝ +""" + + +WANT_BANNER_SECTION_1 = "\n\ + ██████╗ █████╗ ██╗ ██████╗ ███████╗██████╗ █████╗\n\ +██╔════╝██╔══██╗██║ ██╔══██╗██╔════╝██╔══██╗██╔══██╗\n\ +" + + +WANT_BANNER_SECTION_2 = "\ +██║ ███████║██║ ██║ ██║█████╗ ██████╔╝███████║\n\ +██║ ██╔══██║██║ ██║ ██║██╔══╝ ██╔══██╗██╔══██║\n\ +" + + +WANT_BANNER_SECTION_3 = "\ +╚██████╗██║ ██║███████╗██████╔╝███████╗██║ ██║██║ ██║\n\ + ╚═════╝╚═╝ ╚═╝╚══════╝╚═════╝ ╚══════╝╚═╝ ╚═╝╚═╝ ╚═╝\n\ +" + + +class TestBanner: + @mock.patch.dict(os.environ, {'NO_COLOR': '0'}, clear=True) + def test_banner_with_color(self): + with patch('app.ascii_banner.rich_print', return_value=None) as mock_rich_print: + print_rich_banner() + mock_rich_print.assert_called_once_with(f'[blue]{WANT_BANNER_SECTION_1}[/blue][purple]{WANT_BANNER_SECTION_2}[/purple][red]{WANT_BANNER_SECTION_3}[/red]') + + assert get_ascii_banner() == f'\x1b[38;5;20m{WANT_BANNER_SECTION_1}\x1b[38;5;92m{WANT_BANNER_SECTION_2}\x1b[38;5;1m{WANT_BANNER_SECTION_3}\033[0m' + + @mock.patch.dict(os.environ, {'NO_COLOR': '1'}, clear=True) + def test_banner_without_color(self): + with patch('app.ascii_banner.rich_print', return_value=None) as mock_rich_print: + print_rich_banner() + mock_rich_print.assert_called_once_with(f'{WANT_BANNER_SECTION_1}{WANT_BANNER_SECTION_2}{WANT_BANNER_SECTION_3}') + + assert get_ascii_banner() == WANT_BANNER diff --git a/tests/utility/test_base_obfuscator.py b/tests/utility/test_base_obfuscator.py new file mode 100644 index 000000000..b5d13b06a --- /dev/null +++ b/tests/utility/test_base_obfuscator.py @@ -0,0 +1,24 @@ +from app.objects.c_agent import Agent +from app.objects.secondclass.c_link import Link +from app.utility.base_obfuscator import BaseObfuscator + + +class MockObfuscator(BaseObfuscator): + def __init__(self, agent): + self.supported_platforms = dict(linux=['sh']) + super().__init__(agent) + + def sh(self, link, arg): + return f'obfuscated sh command {arg}' + + +class TestBaseObfuscator: + def test_run_obfuscator(self, ability, executor): + test_agent = Agent(paw='testpaw', sleep_min=5, sleep_max=5, watchdog=0, executors=['sh', 'proc'], platform='linux') + test_executor = executor(name='sh', platform='linux') + test_ability = ability(ability_id='123', executors=[test_executor]) + test_link = Link(command='original test command', paw='testpaw', ability=test_ability, id=111111, executor=test_executor) + want = 'obfuscated sh command val' + obf = MockObfuscator(test_agent) + assert want == obf.run(test_link, arg='val') + assert test_link.command_hash == 'eaa27692ee501903114ca1bde590e2ea80a14a615f965b5f4505aa4d73c0555e' diff --git a/tests/utility/test_base_parser.py b/tests/utility/test_base_parser.py new file mode 100644 index 000000000..27cfc4f36 --- /dev/null +++ b/tests/utility/test_base_parser.py @@ -0,0 +1,55 @@ +import pytest + +from app.objects.secondclass.c_fact import Fact +from app.utility.base_parser import BaseParser + + +@pytest.fixture +def mock_parser(): + return BaseParser(dict( + mappers=None, + used_facts=[ + Fact('trait1', value='value1'), + Fact('trait2', value='value2') + ], + source_facts=None + )) + + +class TestBaseParser: + def test_set_value(self, mock_parser): + assert len(mock_parser.used_facts) == 2 + assert BaseParser.set_value('', 'testmatch', mock_parser.used_facts) is None + assert BaseParser.set_value('testsearch', 'testmatch', mock_parser.used_facts) == 'testmatch' + assert BaseParser.set_value('trait2', 'testmatch', mock_parser.used_facts) == 'value2' + + def test_parsing(self): + # Email + assert BaseParser.email('notanemail') == [] + assert BaseParser.email('test@example.com') == ['test@example.com'] + assert BaseParser.email('blah test@example.com blah blah test2@example2.org blah') == ['test@example.com', 'test2@example2.org'] + + # Filename + assert BaseParser.filename('noextension') == [] + assert BaseParser.filename('test.xml') == ['test.xml'] + assert BaseParser.filename('blah blah test2.png blah test.xml blahblah') == ['test2.png', 'test.xml'] + + # Line splitting + assert BaseParser.line('singleline') == ['singleline'] + assert BaseParser.line('two\nlines') == ['two', 'lines'] + assert BaseParser.line('three\r\nline\r\nexample') == ['three', 'line', 'example'] + + # IP address + assert BaseParser.ip('notanip') == [] + assert BaseParser.ip('not.an.ip.address') == [] + assert BaseParser.ip('1.2.3') == [] + assert BaseParser.ip('1.2.3.4') == ['1.2.3.4'] + assert BaseParser.ip('blah 1.2.3.4 1.2. 1.2.3. 5.6.7.8 blah') == ['1.2.3.4', '5.6.7.8'] + + # Broadcast + assert BaseParser.broadcastip('notbroadcast') == [] + assert BaseParser.broadcastip('broadcast 1.2.3.4') == ['1.2.3.4'] + + # JSON + assert BaseParser.load_json('{"a":"b"}') == dict(a='b') + assert BaseParser.load_json('{"a":"b"') is None # malformed diff --git a/tests/utility/test_payload_encoder.py b/tests/utility/test_payload_encoder.py new file mode 100644 index 000000000..66bfd2c73 --- /dev/null +++ b/tests/utility/test_payload_encoder.py @@ -0,0 +1,39 @@ +from unittest import mock + +from app.utility.payload_encoder import xor_bytes, xor_file + + +input = bytes([0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]) + + +class TestPayloadEncoder: + def test_xor_bytes(self): + # Default key + want = bytes([0x33, 0x66, 0x77, 0xad, 0xbb, 0xee, 0xff, 0x25]) + assert xor_bytes(input) == want + assert xor_bytes(xor_bytes(input)) == input + + # Specific key + key = [0x65, 0x43, 0x21] + want = bytes([0x64, 0x60, 0x64, 0x02, 0xca, 0x8a, 0xa8, 0xac]) + assert xor_bytes(input, key) == want + assert xor_bytes(xor_bytes(input)) == input + + def test_xor_file(self): + # Without output file, use default key + want = bytes([0x33, 0x66, 0x77, 0xad, 0xbb, 0xee, 0xff, 0x25]) + m = mock.mock_open(read_data=input) + with mock.patch('builtins.open', m): + assert want == xor_file('test_file') + m.assert_called_once_with('test_file', 'rb') + m().write.assert_not_called() + + # With output file, specific key + key = [0x65, 0x43, 0x21] + want = bytes([0x64, 0x60, 0x64, 0x02, 0xca, 0x8a, 0xa8, 0xac]) + m = mock.mock_open(read_data=input) + with mock.patch('builtins.open', m): + assert want == xor_file('test_file', 'test_output', key) + m.assert_any_call('test_file', 'rb') + m.assert_called_with('test_output', 'wb') + m().write.assert_called_once_with(want) diff --git a/tests/web_server/test_core_endpoints.py b/tests/web_server/test_core_endpoints.py index b764c375f..eb38b1dbd 100644 --- a/tests/web_server/test_core_endpoints.py +++ b/tests/web_server/test_core_endpoints.py @@ -1,16 +1,20 @@ +import base64 import os from http import HTTPStatus from pathlib import Path +import tempfile +from unittest import mock import pytest import yaml -from aiohttp import web +from aiohttp import web, FormData from app.api.rest_api import RestApi from app.objects.c_agent import Agent from app.service.app_svc import AppService from app.service.auth_svc import AuthService from app.service.data_svc import DataService +from app.service.file_svc import FileSvc from app.service.rest_svc import RestService from app.service.interfaces.i_login_handler import LoginHandlerInterface from app.utility.base_service import BaseService @@ -30,6 +34,8 @@ async def initialize(): app_svc = AppService(web.Application()) _ = DataService() + file_svc = FileSvc() + file_svc.encrypt_output = False _ = RestService() auth_svc = AuthService() services = app_svc.get_services() @@ -110,6 +116,135 @@ async def test_invalid_request(aiohttp_client, authorized_cookies, sample_agent) assert messages == dict(sleep_min=['Not a valid integer.']) +async def test_command_overwrite_failure(aiohttp_client, authorized_cookies): + resp = await aiohttp_client.post('/api/rest', + cookies=authorized_cookies, + json=dict(index='configuration', + prop='requirements', + value=dict(go=dict(command='this should not get written', + type='installed program', + version='1.11',), + python=dict(attr='version', + module='sys', + type='python_module', + version='3.11.0')))) + + assert resp.status == HTTPStatus.BAD_REQUEST + + +async def test_upload_file(aiohttp_client): + file_data = bytes([0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]) + with tempfile.TemporaryFile(mode='w+b') as tmp_file: + tmp_file.write(file_data) + tmp_file.flush() + tmp_file.seek(0) + + m = mock.mock_open(read_data=file_data) + with mock.patch('builtins.open', m): + upload_data = FormData() + upload_data.add_field('file', tmp_file, filename='testupload') + resp = await aiohttp_client.post('/file/upload', + headers=dict(Directory='testdir'), + data=upload_data) + assert resp.status == HTTPStatus.OK + m.assert_called_with('data/payloads/testupload', 'wb') + m().write.assert_called_once_with(file_data) + + BaseWorld.set_config('main', 'exfil_dir', 'dummyexfildir') + with tempfile.TemporaryFile(mode='w+b') as tmp_file: + tmp_file.write(file_data) + tmp_file.flush() + tmp_file.seek(0) + + m = mock.mock_open(read_data=file_data) + with mock.patch.object(os, 'makedirs', return_value=None) as mock_makedirs: + with mock.patch('builtins.open', m): + upload_data = FormData() + upload_data.add_field('file', tmp_file, filename='testupload') + resp = await aiohttp_client.post('/file/upload', + headers={'X-Request-ID': 'testid'}, + data=upload_data) + assert resp.status == HTTPStatus.OK + mock_makedirs.assert_any_call('dummyexfildir/testid') + mock_makedirs.assert_called_with('dummyexfildir/testid') + m.assert_called_with('dummyexfildir/testid/testupload', 'wb') + m().write.assert_called_once_with(file_data) + + +async def test_download_file(aiohttp_client): + file_data = bytes([0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]) + m = mock.mock_open(read_data=file_data) + with mock.patch.object(FileSvc, 'find_file_path', return_value=('testplugin', 'testplugin/payloads/testdownload')) as mock_find_file_path: + with mock.patch('builtins.open', m): + resp = await aiohttp_client.post('/file/download', + headers=dict(file='testdownload')) + assert resp.status == HTTPStatus.OK + m.assert_called_with('testplugin/payloads/testdownload', 'rb') + m().read.assert_called_once_with() + mock_find_file_path.assert_called_with('testdownload', location='payloads') + assert await resp.content.read() == file_data + assert resp.headers.get('CONTENT-DISPOSITION') == 'attachment; filename="testdownload"' + assert resp.headers.get('FILENAME') == 'testdownload' + + # Test FileNotFound + with mock.patch.object(FileSvc, 'get_file') as mock_get_file: + mock_get_file.side_effect = FileNotFoundError() + resp = await aiohttp_client.post('/file/download', + headers=dict(file='testdownload')) + assert resp.status == HTTPStatus.NOT_FOUND + assert await resp.content.read() == b'File not found' + + # Test generic exception + with mock.patch.object(FileSvc, 'get_file') as mock_get_file: + mock_get_file.side_effect = Exception('Test exception') + resp = await aiohttp_client.post('/file/download', + headers=dict(file='testdownload')) + assert resp.status == HTTPStatus.NOT_FOUND + assert await resp.content.read() == b'Test exception' + + +async def test_download_exfil_file(aiohttp_client, authorized_cookies): + file_data = bytes([0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]) + encoded_file_path = base64.b64encode('dummyexfildir/path/testdownload'.encode('ascii')).decode('ascii') + bad_encoded_file_path = base64.b64encode('notexfildir/path/testdownload'.encode('ascii')).decode('ascii') + + BaseWorld.set_config('main', 'exfil_dir', 'dummyexfildir') + m = mock.mock_open(read_data=file_data) + with mock.patch('builtins.open', m): + with mock.patch.object(FileSvc, 'find_file_path', return_value=(None, 'dummyexfildir/path/testdownload')): + resp = await aiohttp_client.get(f'/file/download_exfil?file={encoded_file_path}', cookies=authorized_cookies) + assert resp.status == HTTPStatus.OK + m.assert_called_with('dummyexfildir/path/testdownload', 'rb') + m().read.assert_called_once_with() + assert await resp.content.read() == file_data + assert resp.headers.get('CONTENT-DISPOSITION') == 'attachment; filename="testdownload"' + assert resp.headers.get('FILENAME') == 'testdownload' + + # Test FileNotFound + with mock.patch.object(FileSvc, 'read_file') as mock_get_file: + mock_get_file.side_effect = FileNotFoundError() + resp = await aiohttp_client.get(f'/file/download_exfil?file={encoded_file_path}', cookies=authorized_cookies) + assert resp.status == HTTPStatus.NOT_FOUND + assert await resp.content.read() == b'File not found' + + # Test generic exception + with mock.patch.object(FileSvc, 'read_file') as mock_get_file: + mock_get_file.side_effect = Exception('Test exception') + resp = await aiohttp_client.get(f'/file/download_exfil?file={encoded_file_path}', cookies=authorized_cookies) + assert resp.status == HTTPStatus.NOT_FOUND + assert await resp.content.read() == b'Test exception' + + # Test payload not in exfil dir + resp = await aiohttp_client.get(f'/file/download_exfil?file={bad_encoded_file_path}', cookies=authorized_cookies) + assert resp.status == HTTPStatus.NOT_FOUND + assert await resp.content.read() == b'File not found in exfil dir' + + # Test missing URL parameter + resp = await aiohttp_client.get('/file/download_exfil', cookies=authorized_cookies) + assert resp.status == HTTPStatus.BAD_REQUEST + assert await resp.content.read() == b'A file needs to be specified for download' + + async def test_custom_rejecting_login_handler(aiohttp_client): class RejectAllLoginHandler(LoginHandlerInterface): def __init__(self, services):