From effd5d412e19489df28e2f940ecfd2b3802a3e29 Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Thu, 16 Oct 2025 20:00:57 +0000 Subject: [PATCH 01/16] fix asyncmock and atomic planner test --- tests/__init__.py | 6 ------ tests/conftest.py | 2 +- tests/planners/test_atomic.py | 19 +++++++++++-------- tests/services/test_file_svc.py | 2 +- tests/services/test_planning_svc.py | 3 +-- 5 files changed, 14 insertions(+), 18 deletions(-) diff --git a/tests/__init__.py b/tests/__init__.py index 41cf6b56e..e69de29bb 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1,6 +0,0 @@ -from unittest.mock import MagicMock - - -class AsyncMock(MagicMock): - async def __call__(self, *args, **kwargs): - return super(AsyncMock, self).__call__(*args, **kwargs) diff --git a/tests/conftest.py b/tests/conftest.py index c83aca648..5ea34a25b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -13,6 +13,7 @@ from datetime import datetime, timezone from base64 import b64encode from unittest import mock +from unittest.mock import AsyncMock from aiohttp_apispec import validation_middleware from aiohttp import web import aiohttp_jinja2 @@ -64,7 +65,6 @@ from app.api.rest_api import RestApi from app import version -from tests import AsyncMock DIR = os.path.dirname(os.path.abspath(__file__)) CONFIG_DIR = os.path.join(DIR, '..', 'conf') diff --git a/tests/planners/test_atomic.py b/tests/planners/test_atomic.py index bd481fa69..3a9adbdb1 100644 --- a/tests/planners/test_atomic.py +++ b/tests/planners/test_atomic.py @@ -1,7 +1,6 @@ -from tests import AsyncMock - import pytest +from unittest.mock import AsyncMock from app.planners.atomic import LogicalPlanner @@ -15,7 +14,10 @@ def __init__(self): self.adversary = AdversaryStub() self.agents = ['agent_1'] self.wait_for_links_completion = AsyncMock() - self.apply = AsyncMock() + self.apply = AsyncMock(side_effect=self._apply_side_effect) + + def _apply_side_effect(self, value): + return value.id class PlanningSvcStub(): @@ -34,9 +36,10 @@ def __init__(self, ability_id): class LinkStub(): def __init__(self, ability_id): self.ability = AbilityStub(ability_id) + self.id = 'link_' + ability_id def __eq__(self, other): - return self.ability.ability_id == other.ability.ability_id + return self.ability.ability_id == other.ability.ability_id and self.id == other.id @pytest.fixture @@ -64,8 +67,8 @@ def test_atomic_with_links_in_order(self, event_loop, atomic_planner): assert atomic_planner.operation.apply.call_count == 1 assert atomic_planner.operation.wait_for_links_completion.call_count == 1 - atomic_planner.operation.apply.assert_called_with(LinkStub('ability_b')) - atomic_planner.operation.wait_for_links_completion.assert_called_with([LinkStub('ability_b')]) + atomic_planner.operation.apply.assert_awaited_with(LinkStub('ability_b')) + atomic_planner.operation.wait_for_links_completion.assert_awaited_with(['link_ability_b']) def test_atomic_with_links_out_of_order(self, event_loop, atomic_planner): @@ -80,8 +83,8 @@ def test_atomic_with_links_out_of_order(self, event_loop, atomic_planner): assert atomic_planner.operation.apply.call_count == 1 assert atomic_planner.operation.wait_for_links_completion.call_count == 1 - atomic_planner.operation.apply.assert_called_with(LinkStub('ability_b')) - atomic_planner.operation.wait_for_links_completion.assert_called_with([LinkStub('ability_b')]) + atomic_planner.operation.apply.assert_awaited_with(LinkStub('ability_b')) + atomic_planner.operation.wait_for_links_completion.assert_awaited_with(['link_ability_b']) def test_atomic_no_links(self, event_loop, atomic_planner): diff --git a/tests/services/test_file_svc.py b/tests/services/test_file_svc.py index fe7c6c16b..5aecb7c84 100644 --- a/tests/services/test_file_svc.py +++ b/tests/services/test_file_svc.py @@ -5,8 +5,8 @@ import yaml from base64 import b64encode -from tests import AsyncMock from asyncio import Future +from unittest.mock import AsyncMock from app.data_encoders.base64_basic import Base64Encoder from app.data_encoders.plain_text import PlainTextEncoder diff --git a/tests/services/test_planning_svc.py b/tests/services/test_planning_svc.py index 655b3de16..1b2f7af4a 100644 --- a/tests/services/test_planning_svc.py +++ b/tests/services/test_planning_svc.py @@ -1,7 +1,7 @@ import pytest import asyncio import base64 -from unittest.mock import MagicMock +from unittest.mock import MagicMock, AsyncMock from app.objects.c_adversary import Adversary from app.objects.c_obfuscator import Obfuscator @@ -11,7 +11,6 @@ from app.objects.secondclass.c_fact import Fact from app.objects.secondclass.c_requirement import Requirement from app.utility.base_world import BaseWorld -from tests import AsyncMock stop_bucket_exhaustion_params = [ From 5a90efc8d95a7090d8bab1f56df81be85aa056a8 Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Fri, 17 Oct 2025 04:44:09 +0000 Subject: [PATCH 02/16] adjust TCP contact --- app/contacts/contact_tcp.py | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/app/contacts/contact_tcp.py b/app/contacts/contact_tcp.py index 3cb2300c0..cbdb78ffb 100644 --- a/app/contacts/contact_tcp.py +++ b/app/contacts/contact_tcp.py @@ -17,12 +17,34 @@ 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): + if self.op_loop_task: + self.op_loop_task.cancel() + if self.server_task: + self.server_task.cancel() + try: + await self.op_loop_task + except asyncio.CancelledError: + self.log.debug('Cancelled TCP contact operation loop task.') + try: + await self.server_task + except asyncio.CancelledError: + self.log.debug('Cancelled TCP contact server task.') + + async def start_server(self, host, port): + self.server = await asyncio.start_server(self.tcp_handler.accept, host, port) + async with self.server: + await self.server.serve_forever() async def operation_loop(self): while True: From 4adcde31a7e1301e14a7974a81660c152b825a19 Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Fri, 17 Oct 2025 04:44:23 +0000 Subject: [PATCH 03/16] address some warnings --- tests/api/v2/test_knowledge.py | 2 +- tests/conftest.py | 2 +- tests/services/test_data_svc.py | 10 ++-------- 3 files changed, 4 insertions(+), 10 deletions(-) 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/services/test_data_svc.py b/tests/services/test_data_svc.py index bf6528bb7..e354081e1 100644 --- a/tests/services/test_data_svc.py +++ b/tests/services/test_data_svc.py @@ -81,12 +81,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 +168,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 From 90dbb92156ff1e7f3792b0ec94c1db0c519b702c Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Fri, 17 Oct 2025 05:45:51 +0000 Subject: [PATCH 04/16] rework ftp contact --- app/contacts/contact_ftp.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/app/contacts/contact_ftp.py b/app/contacts/contact_ftp.py index 08a756475..9746a2171 100644 --- a/app/contacts/contact_ftp.py +++ b/app/contacts/contact_ftp.py @@ -39,18 +39,13 @@ 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() + await self.server.close() def set_up_server(self): user = self.setup_ftp_users() @@ -89,10 +84,7 @@ def setup_ftp_users(self): ), ) - async def ftp_server_python_old(self): - await self.server.start(host=self.host, port=self.port) - - async def ftp_server_python_new(self): + async def ftp_server(self): await self.server.run(host=self.host, port=self.port) def check_config(self): From 7150e47c69b608eb1480cc65e0a56127cc60764d Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Thu, 30 Oct 2025 13:14:51 +0000 Subject: [PATCH 05/16] remove unused imports --- app/contacts/contact_ftp.py | 1 - tests/services/test_data_svc.py | 1 - 2 files changed, 2 deletions(-) diff --git a/app/contacts/contact_ftp.py b/app/contacts/contact_ftp.py index 9746a2171..4f8847625 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 diff --git a/tests/services/test_data_svc.py b/tests/services/test_data_svc.py index e354081e1..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 From ef05c6ff55a41ef32c96ccfa85ee1519e47a06cb Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Thu, 30 Oct 2025 14:28:22 +0000 Subject: [PATCH 06/16] switch to start method to fix error in unit tests --- app/contacts/contact_ftp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app/contacts/contact_ftp.py b/app/contacts/contact_ftp.py index 4f8847625..b5c5e98ad 100644 --- a/app/contacts/contact_ftp.py +++ b/app/contacts/contact_ftp.py @@ -84,7 +84,7 @@ def setup_ftp_users(self): ) async def ftp_server(self): - await self.server.run(host=self.host, port=self.port) + await self.server.start(host=self.host, port=self.port) def check_config(self): if not self.get_config(FTP_HOST_PROPERTY): From 720d1bb265d18a9bfbc424d37aedd652faf7fd7c Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Thu, 30 Oct 2025 16:57:09 +0000 Subject: [PATCH 07/16] use non-breaking pyasn1 version to address unit test warning --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index c23045237..4192b7b6c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -14,6 +14,7 @@ marshmallow==3.20.1 dirhash==0.2.1 marshmallow-enum==1.5.1 ldap3==2.9.1 +pyasn1~=0.5.1 lxml~=4.9.1 # debrief reportlab==4.0.4 # debrief rich==13.7.0 From 791b5b0ee0f586b8f30163321630065587b4964c Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Mon, 8 Dec 2025 16:08:50 +0000 Subject: [PATCH 08/16] handle task cancelation --- app/contacts/contact_tcp.py | 57 ++++++++++++++++++++++--------------- 1 file changed, 34 insertions(+), 23 deletions(-) diff --git a/app/contacts/contact_tcp.py b/app/contacts/contact_tcp.py index cbdb78ffb..78f07733d 100644 --- a/app/contacts/contact_tcp.py +++ b/app/contacts/contact_tcp.py @@ -35,37 +35,48 @@ async def stop(self): try: await self.op_loop_task except asyncio.CancelledError: - self.log.debug('Cancelled TCP contact operation loop task.') + self.log.debug('Canceled TCP contact operation loop task.') try: await self.server_task except asyncio.CancelledError: - self.log.debug('Cancelled TCP contact server task.') + self.log.debug('Canceled TCP contact server task.') async def start_server(self, host, port): - self.server = await asyncio.start_server(self.tcp_handler.accept, host, port) - async with self.server: - await self.server.serve_forever() + 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.server.close() + await self.server.wait_closed() + 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: + 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) + except asyncio.CancelledError: + self.log.debug('Canceling TCP contact operation loop task.') + raise class TcpSessionHandler(BaseWorld): From 333854834a4b6e238360221c1c9f950e1e276b5c Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Mon, 8 Dec 2025 16:45:23 +0000 Subject: [PATCH 09/16] different method of suppressing exception --- app/contacts/contact_tcp.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/app/contacts/contact_tcp.py b/app/contacts/contact_tcp.py index 78f07733d..149ab2d0b 100644 --- a/app/contacts/contact_tcp.py +++ b/app/contacts/contact_tcp.py @@ -32,14 +32,7 @@ async def stop(self): self.op_loop_task.cancel() if self.server_task: self.server_task.cancel() - try: - await self.op_loop_task - except asyncio.CancelledError: - self.log.debug('Canceled TCP contact operation loop task.') - try: - await self.server_task - except asyncio.CancelledError: - self.log.debug('Canceled TCP contact server task.') + _ = await asyncio.gather(self.op_loop_task, self.server_task, return_exceptions=True) async def start_server(self, host, port): try: From 9e41d675db3a59a1e758c0476b82877faf1d81df Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Mon, 8 Dec 2025 18:59:58 +0000 Subject: [PATCH 10/16] more logging --- app/contacts/contact_tcp.py | 1 + 1 file changed, 1 insertion(+) diff --git a/app/contacts/contact_tcp.py b/app/contacts/contact_tcp.py index 149ab2d0b..752717032 100644 --- a/app/contacts/contact_tcp.py +++ b/app/contacts/contact_tcp.py @@ -44,6 +44,7 @@ async def start_server(self, host, port): if self.server: self.server.close() await self.server.wait_closed() + self.log.debug('Closed TCP contact server.') raise async def operation_loop(self): From 06e94f7b53eca023ed56e06c98b4caa5b8c02c95 Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Mon, 8 Dec 2025 21:45:24 +0000 Subject: [PATCH 11/16] use testing dir for ftp contact unit tests --- tests/contacts/test_contact_ftp.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) 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) From 373e238cbad218cdfe5da45fa569aca82bfbce1e Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Mon, 8 Dec 2025 22:29:05 +0000 Subject: [PATCH 12/16] close TCP sessions during contact close --- app/contacts/contact_tcp.py | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/app/contacts/contact_tcp.py b/app/contacts/contact_tcp.py index 752717032..38d336940 100644 --- a/app/contacts/contact_tcp.py +++ b/app/contacts/contact_tcp.py @@ -32,7 +32,7 @@ async def stop(self): self.op_loop_task.cancel() if self.server_task: self.server_task.cancel() - _ = await asyncio.gather(self.op_loop_task, self.server_task, return_exceptions=True) + _ = await asyncio.gather(self.server_task, self.op_loop_task, return_exceptions=True) async def start_server(self, host, port): try: @@ -42,14 +42,15 @@ async def start_server(self, host, port): 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: - try: + try: + while True: await self.tcp_handler.refresh() for session in self.tcp_handler.sessions: _, instructions = await self.contact_svc.handle_heartbeat(paw=session.paw) @@ -65,12 +66,19 @@ async def operation_loop(self): 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) await asyncio.sleep(20) - except asyncio.CancelledError: - self.log.debug('Canceling TCP contact operation loop task.') - raise + 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 session.writer.wait_closed() + self.log.debug('Closed TCP contact sessions.') + raise class TcpSessionHandler(BaseWorld): @@ -95,6 +103,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: @@ -139,3 +148,4 @@ async def _attempt_connection(self, session_id, reader, timeout): self.log.error("Timeout reached for session %s", session_id) return json.dumps(dict(status=1, pwd='~$ ', response=str(err))) return str(data, 'utf-8') + From fa43881a1ce840182a4381b8a73335f75a2be748 Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Tue, 9 Dec 2025 13:52:58 +0000 Subject: [PATCH 13/16] style fix --- app/contacts/contact_tcp.py | 1 - 1 file changed, 1 deletion(-) diff --git a/app/contacts/contact_tcp.py b/app/contacts/contact_tcp.py index 38d336940..5dfdb81ac 100644 --- a/app/contacts/contact_tcp.py +++ b/app/contacts/contact_tcp.py @@ -148,4 +148,3 @@ async def _attempt_connection(self, session_id, reader, timeout): self.log.error("Timeout reached for session %s", session_id) return json.dumps(dict(status=1, pwd='~$ ', response=str(err))) return str(data, 'utf-8') - From 678161ce1aceb9b968a61e6a7848c960cdce3ddd Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Tue, 9 Dec 2025 14:59:38 +0000 Subject: [PATCH 14/16] refactor --- app/contacts/contact_tcp.py | 41 ++++++++++++++++++++----------------- 1 file changed, 22 insertions(+), 19 deletions(-) diff --git a/app/contacts/contact_tcp.py b/app/contacts/contact_tcp.py index 5dfdb81ac..3cd6f344d 100644 --- a/app/contacts/contact_tcp.py +++ b/app/contacts/contact_tcp.py @@ -52,34 +52,37 @@ async def operation_loop(self): try: 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 asyncio.CancelledError: - raise - except Exception as e: - self.log.debug('[-] operation exception: %s' % e) + 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 session.writer.wait_closed() + 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): From bb0cd002b82c0392c2bf002189d94466be1e0036 Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Sat, 3 Jan 2026 02:21:10 +0000 Subject: [PATCH 15/16] handle possible null server/tasks --- app/contacts/contact_ftp.py | 3 ++- app/contacts/contact_tcp.py | 11 ++++++----- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/app/contacts/contact_ftp.py b/app/contacts/contact_ftp.py index b5c5e98ad..b86a725bb 100644 --- a/app/contacts/contact_ftp.py +++ b/app/contacts/contact_ftp.py @@ -44,7 +44,8 @@ async def start(self): await self.ftp_server() async def stop(self): - await self.server.close() + if self.server: + await self.server.close() def set_up_server(self): user = self.setup_ftp_users() diff --git a/app/contacts/contact_tcp.py b/app/contacts/contact_tcp.py index 3cd6f344d..76a350438 100644 --- a/app/contacts/contact_tcp.py +++ b/app/contacts/contact_tcp.py @@ -28,11 +28,12 @@ async def start(self): self.op_loop_task = loop.create_task(self.operation_loop()) async def stop(self): - if self.op_loop_task: - self.op_loop_task.cancel() - if self.server_task: - self.server_task.cancel() - _ = await asyncio.gather(self.server_task, self.op_loop_task, return_exceptions=True) + 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: From 78498fd848db6fb0278a0b4df3793b1e796c4ae4 Mon Sep 17 00:00:00 2001 From: uruwhy <58484522+uruwhy@users.noreply.github.com> Date: Mon, 5 Jan 2026 20:04:43 +0000 Subject: [PATCH 16/16] address marshmallow deprecated meta.ordered message --- app/api/v2/schemas/caldera_info_schemas.py | 3 --- app/api/v2/schemas/error_schemas.py | 3 --- 2 files changed, 6 deletions(-) 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}