diff --git a/app/api/v2/schemas/caldera_info_schemas.py b/app/api/v2/schemas/caldera_info_schemas.py index c7c1fa7ba..efb52bcd5 100644 --- a/app/api/v2/schemas/caldera_info_schemas.py +++ b/app/api/v2/schemas/caldera_info_schemas.py @@ -9,6 +9,3 @@ class CalderaInfoSchema(schema.Schema): version = fields.String() access = fields.String() plugins = fields.List(fields.Nested(Plugin.display_schema)) - - class Meta: - ordered = True diff --git a/app/api/v2/schemas/error_schemas.py b/app/api/v2/schemas/error_schemas.py index 7d15f2f2b..89d6ce061 100644 --- a/app/api/v2/schemas/error_schemas.py +++ b/app/api/v2/schemas/error_schemas.py @@ -6,9 +6,6 @@ class JsonHttpErrorSchema(schema.Schema): error = fields.String(required=True) details = fields.Dict() - class Meta: - ordered = True - @classmethod def make_dict(cls, error, details=None): obj = {'error': error} diff --git a/app/contacts/contact_ftp.py b/app/contacts/contact_ftp.py index 08a756475..b86a725bb 100644 --- a/app/contacts/contact_ftp.py +++ b/app/contacts/contact_ftp.py @@ -3,7 +3,6 @@ import re import asyncio import aioftp -import sys from app.utility.base_world import BaseWorld @@ -39,18 +38,14 @@ def __init__(self, services): self.user = self.get_config('app.contact.ftp.user') self.pword = self.get_config('app.contact.ftp.pword') self.server = None - self.task = None async def start(self): self.set_up_server() - if sys.version_info >= (3, 7): - self.task = asyncio.create_task(self.ftp_server_python_new()) - else: - self.task = asyncio.create_task(self.ftp_server_python_old()) - await self.task + await self.ftp_server() async def stop(self): - self.task.cancel() + if self.server: + await self.server.close() def set_up_server(self): user = self.setup_ftp_users() @@ -89,12 +84,9 @@ def setup_ftp_users(self): ), ) - async def ftp_server_python_old(self): + async def ftp_server(self): await self.server.start(host=self.host, port=self.port) - async def ftp_server_python_new(self): - await self.server.run(host=self.host, port=self.port) - def check_config(self): if not self.get_config(FTP_HOST_PROPERTY): self.set_config('main', FTP_HOST_PROPERTY, FTP_HOST_DEFAULT) diff --git a/app/contacts/contact_tcp.py b/app/contacts/contact_tcp.py index 3cb2300c0..76a350438 100644 --- a/app/contacts/contact_tcp.py +++ b/app/contacts/contact_tcp.py @@ -17,33 +17,72 @@ def __init__(self, services): self.log = self.create_logger('contact_tcp') self.contact_svc = services.get('contact_svc') self.tcp_handler = TcpSessionHandler(services, self.log) + self.server_task = None + self.op_loop_task = None + self.server = None async def start(self): loop = asyncio.get_event_loop() tcp = self.get_config('app.contact.tcp') - loop.create_task(asyncio.start_server(self.tcp_handler.accept, *tcp.split(':'))) - loop.create_task(self.operation_loop()) + self.server_task = loop.create_task(self.start_server(*tcp.split(':'))) + self.op_loop_task = loop.create_task(self.operation_loop()) + + async def stop(self): + tasks_to_stop = [t for t in (self.server_task, self.op_loop_task) if t is not None] + for t in tasks_to_stop: + if t: + t.cancel() + if tasks_to_stop: + _ = await asyncio.gather(*tasks_to_stop, return_exceptions=True) + + async def start_server(self, host, port): + try: + self.server = await asyncio.start_server(self.tcp_handler.accept, host, port) + async with self.server: + await self.server.serve_forever() + except asyncio.CancelledError: + self.log.debug('Canceling TCP contact server task.') + if self.server: + self.log.debug('Closing TCP contact server.') + self.server.close() + await self.server.wait_closed() + self.log.debug('Closed TCP contact server.') + raise async def operation_loop(self): - while True: - await self.tcp_handler.refresh() - for session in self.tcp_handler.sessions: - _, instructions = await self.contact_svc.handle_heartbeat(paw=session.paw) - for instruction in instructions: - try: - self.log.debug('TCP instruction: %s' % instruction.id) - status, _, response, agent_reported_time = await self.tcp_handler.send( - session.id, - self.decode_bytes(instruction.command), - timeout=instruction.timeout - ) - beacon = dict(paw=session.paw, - results=[dict(id=instruction.id, output=self.encode_string(response), status=status, agent_reported_time=agent_reported_time)]) - await self.contact_svc.handle_heartbeat(**beacon) - await asyncio.sleep(instruction.sleep) - except Exception as e: - self.log.debug('[-] operation exception: %s' % e) - await asyncio.sleep(20) + try: + while True: + await self.tcp_handler.refresh() + await self.handle_sessions() + await asyncio.sleep(20) + except asyncio.CancelledError: + self.log.debug('Canceling TCP contact operation loop task.') + for sess in self.tcp_handler.sessions: + self.log.debug(f'Closing session {sess.id}.') + sess.writer.close() + await sess.writer.wait_closed() + self.log.debug('Closed TCP contact sessions.') + raise + + async def handle_sessions(self): + for session in self.tcp_handler.sessions: + _, instructions = await self.contact_svc.handle_heartbeat(paw=session.paw) + for instruction in instructions: + try: + self.log.debug('TCP instruction: %s' % instruction.id) + status, _, response, agent_reported_time = await self.tcp_handler.send( + session.id, + self.decode_bytes(instruction.command), + timeout=instruction.timeout + ) + beacon = dict(paw=session.paw, + results=[dict(id=instruction.id, output=self.encode_string(response), status=status, agent_reported_time=agent_reported_time)]) + await self.contact_svc.handle_heartbeat(**beacon) + await asyncio.sleep(instruction.sleep) + except asyncio.CancelledError: + raise + except Exception as e: + self.log.debug('[-] operation exception: %s' % e) class TcpSessionHandler(BaseWorld): @@ -68,6 +107,7 @@ async def refresh(self): index += 1 async def accept(self, reader, writer): + self.log.debug('Accepting connection.') try: profile = await self._handshake(reader) except Exception as e: diff --git a/requirements.txt b/requirements.txt index 326019968..520edd8af 100644 --- a/requirements.txt +++ b/requirements.txt @@ -14,6 +14,7 @@ marshmallow==3.26.2 dirhash==0.2.1 marshmallow-enum==1.5.1 ldap3==2.9.1 +pyasn1~=0.5.1 reportlab==4.0.4 # debrief rich==13.7.0 lxml==6.0.2 # debrief diff --git a/tests/api/v2/test_knowledge.py b/tests/api/v2/test_knowledge.py index 70ccbf4cb..2ec53133a 100644 --- a/tests/api/v2/test_knowledge.py +++ b/tests/api/v2/test_knowledge.py @@ -42,7 +42,7 @@ def base_world(): @pytest.fixture -async def knowledge_webapp(event_loop, base_world, data_svc): +async def knowledge_webapp(base_world, data_svc): app_svc = AppService(web.Application()) app_svc.add_service('auth_svc', AuthService()) app_svc.add_service('knowledge_svc', KnowledgeService()) diff --git a/tests/conftest.py b/tests/conftest.py index 5ea34a25b..dd95c6ad4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -332,7 +332,7 @@ def agent_config(): @pytest.fixture -async def api_v2_client(event_loop, aiohttp_client, contact_svc): +async def api_v2_client(aiohttp_client, contact_svc): def make_app(svcs): warnings.filterwarnings( "ignore", diff --git a/tests/contacts/test_contact_ftp.py b/tests/contacts/test_contact_ftp.py index 6db9af6f9..8a771f529 100644 --- a/tests/contacts/test_contact_ftp.py +++ b/tests/contacts/test_contact_ftp.py @@ -1,5 +1,6 @@ import pytest import os +import shutil from app.contacts import contact_ftp from app.utility.base_world import BaseWorld @@ -26,7 +27,7 @@ def base_world(): BaseWorld.apply_config(name='main', config={'app.contact.ftp.host': '0.0.0.0', 'app.contact.ftp.port': '2222', 'app.contact.ftp.pword': 'caldera', - 'app.contact.ftp.server.dir': 'ftp_dir', + 'app.contact.ftp.server.dir': 'ftp_dir_testing', 'app.contact.ftp.user': 'caldera_user', 'plugins': ['sandcat', 'stockpile'], 'crypt_salt': 'BLAH', @@ -64,7 +65,7 @@ def test_server_setup(ftp_c2): 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' + assert ftp_c2.directory == 'ftp_dir_testing' assert ftp_c2.user == 'caldera_user' assert ftp_c2.pword == 'caldera' assert ftp_c2.server is None @@ -80,6 +81,6 @@ def test_my_server_setup(ftp_c2_my_server): 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') + 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) - os.rmdir(ftp_c2_my_server.ftp_server_dir) + shutil.rmtree(ftp_c2_my_server.ftp_server_dir) diff --git a/tests/services/test_data_svc.py b/tests/services/test_data_svc.py index bf6528bb7..279ddc8f1 100644 --- a/tests/services/test_data_svc.py +++ b/tests/services/test_data_svc.py @@ -1,4 +1,3 @@ -import asyncio import glob import json import yaml @@ -81,12 +80,6 @@ def strip_payload_yaml(path): return PAYLOAD_CONFIG_YAMLS.get(path, []) -def async_mock_return(to_return): - mock_future = asyncio.Future() - mock_future.set_result(to_return) - return mock_future - - class TestDataService: mock_payload_config = dict() @@ -174,8 +167,8 @@ def test_no_autogen_cleanup_cmds(self, event_loop, data_svc): assert not executor.cleanup @mock.patch.object(BaseWorld, 'strip_yml', wraps=strip_payload_yaml) - @mock.patch.object(DataService, '_apply_special_payload_hooks', return_value=async_mock_return(None)) - @mock.patch.object(DataService, '_apply_special_extension_hooks', return_value=async_mock_return(None)) + @mock.patch.object(DataService, '_apply_special_payload_hooks', return_value=None) + @mock.patch.object(DataService, '_apply_special_extension_hooks', return_value=None) def test_load_payloads(self, mock_ext_hooks, mock_payload_hooks, mock_strip_yml, event_loop, data_svc): def _mock_apply_payload_config(config=None, **_): TestDataService.mock_payload_config = config