Addressing some unit test warnings (#3222)

* fix asyncmock and atomic planner test

* adjust TCP contact

* address some warnings

* rework ftp contact

* remove unused imports

* switch to start method to fix error in unit tests

* use non-breaking pyasn1 version to address unit test warning

* handle task cancelation

* different method of suppressing exception

* more logging

* use testing dir for ftp contact unit tests

* close TCP sessions during contact close

* style fix

* refactor

* handle possible null server/tasks

* address marshmallow deprecated meta.ordered message

---------

Co-authored-by: Chris Lenk <clenk@users.noreply.github.com>

Author: uruwhy

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

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}

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)

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:

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

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())

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",

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)

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