diff --git a/app/api/v2/handlers/agent_api.py b/app/api/v2/handlers/agent_api.py index 8ef885f38..4b3d8d30d 100644 --- a/app/api/v2/handlers/agent_api.py +++ b/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 " diff --git a/app/api/v2/managers/agent_api_manager.py b/app/api/v2/managers/agent_api_manager.py index c5ba93afa..b8b145ba7 100644 --- a/app/api/v2/managers/agent_api_manager.py +++ b/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'} diff --git a/app/objects/c_agent.py b/app/objects/c_agent.py index 6c087bdd1..1fe8457bb 100644 --- a/app/objects/c_agent.py +++ b/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,6 +232,10 @@ 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(): @@ -220,8 +243,11 @@ async def gui_modification(self, **kwargs): 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', '') diff --git a/tests/api/v2/handlers/test_agents_api.py b/tests/api/v2/handlers/test_agents_api.py index d25683b78..5bb47d656 100644 --- a/tests/api/v2/handlers/test_agents_api.py +++ b/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) diff --git a/tests/conftest.py b/tests/conftest.py index 28fb32c66..3fe3e7554 100644 --- a/tests/conftest.py +++ b/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() diff --git a/tests/objects/test_agent.py b/tests/objects/test_agent.py index a414e4585..220ffd52d 100644 --- a/tests/objects/test_agent.py +++ b/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) diff --git a/tests/services/test_rest_svc.py b/tests/services/test_rest_svc.py index 22437cb17..db2b75739 100644 --- a/tests/services/test_rest_svc.py +++ b/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) diff --git a/tests/web_server/test_core_endpoints.py b/tests/web_server/test_core_endpoints.py index 6284d696f..8090d739e 100644 --- a/tests/web_server/test_core_endpoints.py +++ b/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()