Merge branch 'master' into docker-compose

Author: clenk

app/api/rest_api.py

@@ -29,7 +29,9 @@ def __init__(self, services):
         asyncio.get_event_loop().create_task(AdvancedPack(services).enable())
 
     async def enable(self):
-        self.app_svc.application.router.add_static('/assets', 'plugins/magma/dist/assets/', append_version=True)
+        # check if plugin path is present
+        if os.path.exists("plugins/magma/dist/assets") and (len(os.listdir("plugins/magma/dist/assets")) > 0):
+            self.app_svc.application.router.add_static('/assets', 'plugins/magma/dist/assets/', append_version=True)
         # TODO: only serve static files in legacy plugin mode
         self.app_svc.application.router.add_static('/gui', 'static/', append_version=True)
         # unauthorized GUI endpoints

app/api/v2/handlers/agent_api.py

@@ -19,6 +19,7 @@ def add_routes(self, app: web.Application):
         router.add_get('/agents', self.get_agents)
         router.add_get('/agents/{paw}', self.get_agent_by_id)
         router.add_post('/agents', self.create_agent)
+        router.add_post('/agents/kill/{paw}', self.kill_agent)
         router.add_patch('/agents/{paw}', self.update_agent)
         router.add_put('/agents/{paw}', self.create_or_update_agent)
         router.add_delete('/agents/{paw}', self.delete_agent)
@@ -64,6 +65,14 @@ async def create_agent(self, request: web.Request):
         agent = await self.create_object(request)
         return web.json_response(agent.display)
 
+    @aiohttp_apispec.docs(tags=['agents'],
+                          summary="Kills an agent",
+                          description="Marks an agent to stop after the next beacon.")
+    async def kill_agent(self, request: web.Request):
+        target_paw = request.match_info.get('paw')
+        kill_resp = await self._api_manager.kill_agent(target_paw)
+        return web.json_response(kill_resp)
+
     @aiohttp_apispec.docs(tags=['agents'],
                           summary="Update an Agent",
                           description="Update the attributes of a specific Agent using its ID (paw). Use the paw "

app/api/v2/managers/agent_api_manager.py

@@ -1,4 +1,5 @@
 from app.api.v2.managers.base_api_manager import BaseApiManager
+from app.api.v2.responses import JsonHttpNotFound
 
 
 class AgentApiManager(BaseApiManager):
@@ -25,3 +26,11 @@ async def get_deploy_commands(self, ability_id: str = None):
         app_config.update({f'agents.{k}': v for k, v in self.get_config(name='agents').items()})
 
         return dict(abilities=raw_abilities, app_config=app_config)
+
+    async def kill_agent(self, target_paw: str):
+        agents = await self._data_svc.locate('agents', {'paw': target_paw})
+        if not agents:
+            raise JsonHttpNotFound(f'Agent {target_paw} not found.')
+        target = agents[0]
+        await target.kill()
+        return {'response': 'Ok'}

app/objects/c_agent.py

@@ -48,6 +48,8 @@ class AgentFieldsSchema(ma.Schema):
     links = ma.fields.List(ma.fields.Nested(LinkSchema), dump_only=True)
     pending_contact = ma.fields.String()
 
+    status = ma.fields.String(dump_only=True)
+
     @ma.pre_load
     def remove_nulls(self, in_data, **_):
         return {k: v for k, v in in_data.items() if v is not None}
@@ -58,6 +60,7 @@ def remove_properties(self, data, **_):
         data.pop('created', None)
         data.pop('last_seen', None)
         data.pop('links', None)
+        data.pop('status', None)
         return data
 
 
@@ -85,6 +88,20 @@ def unique(self):
     def display_name(self):
         return '{}${}'.format(self.host, self.username)
 
+    @property
+    def status(self):
+        now = datetime.now(timezone.utc)
+        untrusted_buffer = int(self.get_config(name='agents', prop='untrusted_timer'))
+        time_diff = (now - self.last_seen).total_seconds()
+        expired = time_diff > int(self.sleep_max) + untrusted_buffer
+        if self._marked_for_stop:
+            # If agent hasn't received the stop instruction yet in a beacon response, it's still pending stop
+            # Otherwise, if agent has received the stop instruction or takes too long to beacon back, mark as dead
+            return 'dead' if self._stop_delivered or expired else 'pending kill'
+        else:
+            # If agent hasn't beaconed in since max beacon time + untrusted timer, mark as dead
+            return 'dead' if expired else 'alive'
+
     @classmethod
     def is_global_variable(cls, variable):
         if variable.startswith('payload:'):
@@ -139,6 +156,8 @@ def __init__(self, sleep_min=30, sleep_max=60, watchdog=0, platform='unknown', s
             self.upstream_dest = self.server
         self._executor_change_to_assign = None
         self.log = self.create_logger('agent')
+        self._marked_for_stop = False
+        self._stop_delivered = False
 
     def store(self, ram):
         existing = self.retrieve(ram['agents'], self.unique)
@@ -213,15 +232,22 @@ async def heartbeat_modification(self, **kwargs):
             # Don't update executors if we're waiting to assign an executor change to the agent.
             self.update('executors', kwargs.get('executors'))
 
+        # Check if agent has been marked to stop
+        if self._marked_for_stop and not self._stop_delivered:
+            self._stop_delivered = True
+
     async def gui_modification(self, **kwargs):
         loaded = AgentFieldsSchema(only=('group', 'trusted', 'sleep_min', 'sleep_max', 'watchdog', 'pending_contact')).load(kwargs)
         for k, v in loaded.items():
             self.update(k, v)
 
     async def kill(self):
         self.update('watchdog', 1)
-        self.update('sleep_min', 60 * 2)
-        self.update('sleep_max', 60 * 2)
+        self.update('sleep_min', 3)
+        self.update('sleep_max', 3)
+
+        self._marked_for_stop = True
+        self._stop_delivered = False
 
     def replace(self, encoded_cmd, file_svc):
         decoded_cmd = b64decode(encoded_cmd).decode('utf-8', errors='ignore').replace('\n', '')

tests/api/v2/handlers/test_agents_api.py

@@ -163,6 +163,30 @@ async def test_unauthorized_create_agent(self, api_v2_client, new_agent_payload)
         resp = await api_v2_client.post('/api/v2/agents', json=new_agent_payload)
         assert resp.status == HTTPStatus.UNAUTHORIZED
 
+    async def test_kill_agent(self, api_v2_client, api_cookies, test_agent, mocker, mock_time):
+        with mocker.patch('app.objects.c_agent.datetime') as mock_datetime:
+            mock_datetime.return_value = mock_datetime
+            mock_datetime.now.return_value = mock_time
+
+            assert test_agent.watchdog == 0 and test_agent.sleep_min == 2 and test_agent.sleep_max == 8
+            assert not (test_agent._marked_for_stop or test_agent._stop_delivered)
+            assert test_agent.status == 'alive'
+
+            resp = await api_v2_client.post('/api/v2/agents/kill/123', cookies=api_cookies)
+            assert resp.status == HTTPStatus.OK
+            assert {'response': 'Ok'} == await resp.json()
+            assert test_agent.status == 'pending kill'
+            assert test_agent._marked_for_stop and not test_agent._stop_delivered
+            assert test_agent.watchdog == 1 and test_agent.sleep_min == 3 and test_agent.sleep_max == 3
+
+    async def test_unauthorized_kill_agent(self, api_v2_client):
+        resp = await api_v2_client.post('/api/v2/agents/kill/123')
+        assert resp.status == HTTPStatus.UNAUTHORIZED
+
+    async def test_nonexistent_kill_agent(self, api_v2_client, api_cookies):
+        resp = await api_v2_client.post('/api/v2/agents/kill/999', cookies=api_cookies)
+        assert resp.status == HTTPStatus.NOT_FOUND
+
     async def test_update_agent(self, api_v2_client, api_cookies, test_agent, updated_agent_fields_payload,
                                 expected_updated_agent_dump):
         resp = await api_v2_client.patch('/api/v2/agents/123', cookies=api_cookies, json=updated_agent_fields_payload)

tests/conftest.py

@@ -366,6 +366,8 @@ async def initialize():
             BaseWorld.apply_config('main', yaml.safe_load(fle))
         with open(Path(__file__).parents[1] / 'conf' / 'payloads.yml', 'r') as fle:
             BaseWorld.apply_config('payloads', yaml.safe_load(fle))
+        with open(Path(__file__).parents[1] / 'conf' / 'agents.yml', 'r') as fle:
+            BaseWorld.apply_config('agents', yaml.safe_load(fle))
 
         app_svc = AppService(web.Application(client_max_size=5120 ** 2))
         _ = DataService()

tests/objects/test_agent.py

@@ -1,9 +1,11 @@
 from base64 import b64decode
+from datetime import timedelta
 
 from app.objects.c_ability import Ability
 from app.objects.c_agent import Agent
 from app.objects.secondclass.c_executor import Executor
 from app.objects.secondclass.c_fact import Fact
+from app.utility.base_world import BaseWorld
 
 
 class TestAgent:
@@ -126,6 +128,42 @@ 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_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')
+        assert agent.status == 'alive'
+        event_loop.run_until_complete(agent.kill())
+        assert agent.status == 'pending kill'
+        assert agent.watchdog == 1 and agent.sleep_min == 3 and agent.sleep_max == 3
+        event_loop.run_until_complete(agent.heartbeat_modification())
+        assert agent.status == 'dead'
+
+        with mocker.patch('app.objects.c_agent.datetime') as mock_datetime:
+            mock_datetime.return_value = mock_datetime
+            mock_datetime.now.return_value = mock_time
+            second_agent = Agent(paw='123', sleep_min=2, sleep_max=8, watchdog=0, executors=['cmd'], platform='windows')
+            event_loop.run_until_complete(second_agent.kill())
+
+            mock_datetime.now.return_value = mock_time + timedelta(0, 10)
+            assert second_agent.status == 'pending kill'
+            mock_datetime.now.return_value = mock_time + timedelta(0, 32)
+            assert second_agent.status == 'pending kill'
+            mock_datetime.now.return_value = mock_time + timedelta(0, 34)
+            assert second_agent.status == 'dead'
+
+    def test_status_and_timeout(self, event_loop, mocker, mock_time):
+        BaseWorld.set_config(name='agents', prop='untrusted_timer', value=30)
+        with mocker.patch('app.objects.c_agent.datetime') as mock_datetime:
+            mock_datetime.return_value = mock_datetime
+            mock_datetime.now.return_value = mock_time
+            agent = Agent(paw='123', sleep_min=2, sleep_max=8, watchdog=0, executors=['cmd'], platform='windows')
+            assert agent.status == 'alive'
+
+            mock_datetime.now.return_value = mock_time + timedelta(0, 30)
+            assert agent.status == 'alive'
+            mock_datetime.now.return_value = mock_time + timedelta(0, 39)
+            assert agent.status == 'dead'
+
     def test_store_new_agent(self, data_svc):
         agent = Agent(paw='123', sleep_min=2, sleep_max=8, watchdog=0, executors=['cmd', 'test'], platform='windows')
         stored_agent = agent.store(data_svc.ram)

tests/services/test_rest_svc.py

@@ -88,7 +88,7 @@ def test_delete_operation(self, event_loop, rest_svc, data_svc):
                                               'privilege': 'User', 'proxy_receivers': {}, 'proxy_chain': [],
                                               'origin_link_id': '', 'deadman_enabled': False,
                                               'available_contacts': ['unknown'], 'pending_contact': 'unknown',
-                                              'host_ip_addrs': [], 'upstream_dest': '://None:None'}],
+                                              'host_ip_addrs': [], 'upstream_dest': '://None:None', 'status': 'alive'}],
                               'visibility': 50, 'autonomous': 1, 'chain': [], 'auto_close': False,
                               'obfuscator': 'plain-text', 'use_learning_parsers': False,
                               'group': '',
@@ -163,7 +163,7 @@ def test_create_operation(self, event_loop, rest_svc, data_svc):
                      'display_name': 'unknown$unknown', 'group': 'red', 'location': 'unknown', 'privilege': 'User',
                      'proxy_receivers': {}, 'proxy_chain': [], 'origin_link_id': '',
                      'deadman_enabled': False, 'available_contacts': ['unknown'], 'pending_contact': 'unknown',
-                     'host_ip_addrs': [], 'upstream_dest': '://None:None'}],
+                     'host_ip_addrs': [], 'upstream_dest': '://None:None', 'status': 'alive'}],
                 'visibility': 50, 'autonomous': 1, 'chain': [], 'auto_close': False, 'objective': '',
                 'obfuscator': 'plain-text', 'use_learning_parsers': False}
         internal_rest_svc = rest_svc(event_loop)

tests/web_server/test_core_endpoints.py

@@ -25,6 +25,8 @@ async def initialize():
             BaseWorld.apply_config('main', yaml.safe_load(fle))
         with open(Path(__file__).parents[2] / 'conf' / 'payloads.yml', 'r') as fle:
             BaseWorld.apply_config('payloads', yaml.safe_load(fle))
+        with open(Path(__file__).parents[2] / 'conf' / 'agents.yml', 'r') as fle:
+            BaseWorld.apply_config('agents', yaml.safe_load(fle))
 
         app_svc = AppService(web.Application())
         _ = DataService()
@@ -63,6 +65,8 @@ async def sample_agent(aiohttp_client):
 
 
 async def test_home(aiohttp_client):
+    if not (len(os.listdir("plugins/magma/dist/assets")) > 0):
+        pytest.xfail("Magma plugin not present, expecting failure")
     resp = await aiohttp_client.get('/')
     assert resp.status == HTTPStatus.OK
     assert resp.content_type == 'text/html'