improving unit test coverage (WIP)

app/api/rest_api.py

@@ -137,9 +137,9 @@
                             ('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 @@
                 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')

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

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]:

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]"
         )

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:

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

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)

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

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

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,

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

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(

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)

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

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