async def test_custom_rejecting_login_handler(aiohttp_client): class RejectAllLoginHandler(LoginHandlerInterface): def __init__(self, services): super().__init__(services, 'Reject All Login Handler') async def handle_login(self, request, **kwargs): # Always reject login and return 401 raise web.HTTPUnauthorized(text='Automatic rejection') async def handle_login_redirect(self, request, **kwargs): # Always reject and return 401 raise web.HTTPUnauthorized(text='Automatic rejection') login_handler = RejectAllLoginHandler(BaseService.get_services()) await BaseService.get_service('auth_svc').set_login_handlers( BaseService.get_services(), login_handler) resp = await aiohttp_client.post('/enter', allow_redirects=False, data=dict(username='******', password='******')) assert resp.status == HTTPStatus.UNAUTHORIZED assert await resp.text() == 'Automatic rejection' resp = await aiohttp_client.get('/', allow_redirects=False) assert resp.status == HTTPStatus.UNAUTHORIZED assert await resp.text() == 'Automatic rejection'
async def test_add_fact_to_operation(knowledge_webapp, aiohttp_client, test_operation, setup_empty_operation): client = await aiohttp_client(knowledge_webapp) fact_data = { 'trait': 'demo', 'value': 'test', 'source': test_operation['id'] } resp = await client.post('/facts', json=fact_data, headers=headers) data = await resp.json() response = data['added'] assert len(response) == 1 assert response[0]['trait'] == 'demo' assert response[0]['value'] == 'test' assert response[0]['source'] == test_operation['id'] tmp = await client.get('/facts', json=fact_data, headers=headers) cur = await tmp.json() current = cur['found'] assert current == response data_svc = BaseService.get_service('data_svc') file_svc = BaseService.get_service('file_svc') matched_operations = await data_svc.locate('operations', {'id': test_operation['id']}) report = await matched_operations[0].report(file_svc, data_svc) assert response[0] in report['facts']
async def test_custom_accepting_login_handler(aiohttp_client): class AcceptAllLoginHandler(LoginHandlerInterface): def __init__(self, services): super().__init__(services, 'Accept All Login Handler') async def handle_login(self, request, **kwargs): # Always accept login data = await request.post() username = data.get('username', 'default username') auth_svc = self.services.get('auth_svc') if not auth_svc: raise Exception('Auth service not available.') await auth_svc.handle_successful_login(request, username) async def handle_login_redirect(self, request, **kwargs): await self.handle_login(request, **kwargs) login_handler = AcceptAllLoginHandler(BaseService.get_services()) await BaseService.get_service('auth_svc').set_login_handlers( BaseService.get_services(), login_handler) resp = await aiohttp_client.post('/enter', allow_redirects=False, data=dict(username='******', password='******')) assert resp.status == HTTPStatus.FOUND assert resp.headers.get('Location') == '/' assert 'API_SESSION' in resp.cookies
def sample_agent(loop, aiohttp_client): kwargs = dict(architecture='amd64', exe_name='sandcat.go', executors=['shellcode_amd64', 'sh'], group='red', host='testsystem.localdomain', location='./sandcat.go', pid=125266, platform='linux', ppid=124042, privilege='User', server='http://127.0.0.1:8888', username='******', paw=None, contact='http') agent = loop.run_until_complete( BaseService.get_service('data_svc').store(Agent(sleep_min=0, sleep_max=60, watchdog=0, **kwargs)) ) yield agent loop.run_until_complete( BaseService.get_service('data_svc').remove('agent', dict(paw=agent.paw)) )
def test_obfuscator(loop, api_v2_client): obfuscator = Obfuscator(name='test', description='a test obfuscator', module='testmodule') loop.run_until_complete( BaseService.get_service('data_svc').store(obfuscator)) return obfuscator
def test_agent(event_loop, 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 test_agent = Agent(paw='123', sleep_min=2, sleep_max=8, watchdog=0, executors=['sh'], platform='linux') event_loop.run_until_complete(BaseService.get_service('data_svc').store(test_agent)) return test_agent
def test_objective(loop, test_goal): objective = Objective(id='123', name='test objective', description='a test objective', goals=[test_goal]) loop.run_until_complete( BaseService.get_service('data_svc').store(objective)) return objective
async def which_plugin(self): file_svc = BaseService.get_service('file_svc') for plugin in os.listdir('plugins'): if await file_svc.walk_file_path( os.path.join('plugins', plugin, 'data', ''), '%s.yml' % self.planner_id): return plugin return None
async def all_relationships(self): knowledge_svc_handle = BaseService.get_service('knowledge_svc') seeded_relationships = [] if self.source: seeded_relationships = await knowledge_svc_handle.get_relationships( criteria=dict(origin=self.source.id)) learned_relationships = await knowledge_svc_handle.get_relationships( criteria=dict(origin=self.id)) return seeded_relationships + learned_relationships
def test_schedule(test_operation, event_loop): operation = OperationSchema().load(test_operation) schedule = ScheduleSchema().load( dict(id='123', schedule='03:00:00.000000', task=operation.schema.dump(operation))) event_loop.run_until_complete( BaseService.get_service('data_svc').store(schedule)) return schedule
def test_source(event_loop): test_fact = Fact(trait='remote.host.fqdn', value='dc') test_source = Source(id='123', name='test', facts=[test_fact], adjustments=[]) event_loop.run_until_complete( BaseService.get_service('data_svc').store(test_source)) return test_source
def test_agent(loop): agent = Agent(paw='123', sleep_min=2, sleep_max=8, watchdog=0, executors=['sh'], platform='linux') loop.run_until_complete(BaseService.get_service('data_svc').store(agent)) return agent
def setup_empty_operation(event_loop, test_operation): test_operation = OperationSchema().load(test_operation) test_operation.set_start_details() test_objective = Objective(id='123', name='test objective', description='test', goals=[]) test_operation.objective = test_objective event_loop.run_until_complete( BaseService.get_service('data_svc').store(test_operation))
def test_link_agent_reported_time_present_when_set_roundtrip( self, ability, executor): test_executor = executor(name='psh', platform='windows') test_ability = ability(ability_id='123') test_link = Link( command= 'sc.exe \\dc create sandsvc binpath= "s4ndc4t.exe -originLinkID 111111"', paw='123456', ability=test_ability, executor=test_executor, id=111111, agent_reported_time=BaseService.get_timestamp_from_string( '2021-02-23 11:50:16')) serialized_link = test_link.display loaded_link = Link.load(serialized_link) assert serialized_link['agent_reported_time'] == '2021-02-23 11:50:16' assert loaded_link.agent_reported_time == BaseService.get_timestamp_from_string( '2021-02-23 11:50:16')
def deploy_ability(test_executor, event_loop): ability = AbilitySchema().load(dict(ability_id='123', tactic='persistence', technique_id='auto-generated', technique_name='auto-generated', name='test deploy command', description='test ability', executors=[ExecutorSchema().dump(test_executor)])) event_loop.run_until_complete(BaseService.get_service('data_svc').store(ability)) return ability
async def all_facts(self): knowledge_svc_handle = BaseService.get_service('knowledge_svc') seeded_facts = [] if self.source: seeded_facts = await knowledge_svc_handle.get_facts(criteria=dict( source=self.source.id)) learned_facts = await knowledge_svc_handle.get_facts(criteria=dict( source=self.id)) learned_facts = [f for f in learned_facts if f.score > 0] return seeded_facts + learned_facts
def test_ability(test_executor, loop): ability = AbilitySchema().load( dict(ability_id='123', tactic='discovery', technique_id='auto-generated', technique_name='auto-generated', name='Manual Command', description='test ability', executors=[ExecutorSchema().dump(test_executor)])) loop.run_until_complete(BaseService.get_service('data_svc').store(ability)) return ability
def fake_event_svc(loop): class FakeEventService(BaseService, EventServiceInterface): def __init__(self): self.fired = {} def reset(self): self.fired = {} async def observe_event(self, callback, exchange=None, queue=None): pass async def fire_event(self, exchange=None, queue=None, timestamp=True, **callback_kwargs): self.fired[exchange, queue] = callback_kwargs service = FakeEventService() service.add_service('event_svc', service) yield service BaseService.remove_service('event_svc')
def test_adversary(loop): expected_adversary = {'name': 'test', 'description': 'an empty adversary profile', 'adversary_id': '123', 'objective': '495a9828-cab1-44dd-a0ca-66e58177d8cc', 'tags': [], 'atomic_ordering': [], 'plugin': ''} test_adversary = AdversarySchema().load(expected_adversary) loop.run_until_complete(BaseService.get_service('data_svc').store(test_adversary)) return test_adversary
async def _init_source(self): # seed knowledge_svc with source facts if self.source: knowledge_svc_handle = BaseService.get_service('knowledge_svc') for f in self.source.facts: f.origin_type = OriginType.SEEDED f.source = self.source.id await knowledge_svc_handle.add_fact(f) for r in self.source.relationships: r.origin = self.source.id await knowledge_svc_handle.add_relationship(r)
async def _save_fact(link, facts, fact, operation=None): fact.source_type = OriginType.LEARNED.name fact.source = operation.id if operation else link.id if all(fact.trait) and not any(fact == f for f in facts): fact.collected_by = link.paw fact.technique_id = link.ability.technique_id fact.links = [link] fact.relationships = [] knowledge_svc_handle = BaseService.get_service('knowledge_svc') await knowledge_svc_handle.add_fact(fact) link.facts.append(fact)
async def test_get_potential_links(self, api_v2_client, api_cookies, mocker, async_return): BaseService.get_service( 'rest_svc').build_potential_abilities = mocker.Mock() BaseService.get_service( 'rest_svc').build_potential_abilities.return_value = async_return( []) expected_link = Link(command='whoami', paw='123456', id='789') BaseService.get_service( 'rest_svc').build_potential_links = mocker.Mock() BaseService.get_service( 'rest_svc').build_potential_links.return_value = async_return( [expected_link]) resp = await api_v2_client.get( '/api/v2/operations/123/potential-links', cookies=api_cookies) result = await resp.json() assert len(result) == 1 assert result[0]['id'] == expected_link.id assert result[0]['paw'] == expected_link.paw assert result[0]['command'] == expected_link.command
async def _create_relationships(self, relationships, operation): for relationship in relationships: relationship.origin = operation.id if operation else self.id await self._save_fact(operation, relationship.source, relationship.score, relationship.shorthand) await self._save_fact(operation, relationship.target, relationship.score, relationship.shorthand) if all((relationship.source.trait, relationship.edge)): knowledge_svc_handle = BaseService.get_service('knowledge_svc') await knowledge_svc_handle.add_relationship(relationship) self.relationships.append(relationship)
def _emit_state_change_event(self, from_state, to_state): event_svc = BaseService.get_service('event_svc') task = asyncio.get_event_loop().create_task( event_svc.fire_event(exchange=Operation.EVENT_EXCHANGE, queue=Operation.EVENT_QUEUE_STATE_CHANGED, op=self.id, from_state=from_state, to_state=to_state)) return task
def _emit_status_change_event(self, from_status, to_status): event_svc = BaseService.get_service('event_svc') task = asyncio.get_event_loop().create_task( event_svc.fire_event(exchange=Link.EVENT_EXCHANGE, queue=Link.EVENT_QUEUE_STATUS_CHANGED, link=self.id, from_status=from_status, to_status=to_status)) return task
async def handle_link_completed(socket, path, services): data = json.loads(await socket.recv()) paw = data['agent']['paw'] data_svc = services.get('data_svc') await process_elasticsearch_result(data, services) agent = await data_svc.locate('agents', match=dict(paw=paw, access=data_svc.Access.RED)) if agent: pid = data['pid'] op_type = 'hidden' if BaseService.Access(data.get('access')) == BaseService.Access.HIDDEN else 'visible' return await services.get('response_svc').respond_to_pid(pid, agent[0], op_type)
def test_ability(loop, api_v2_client, executor): executor_linux = executor(name='sh', platform='linux') ability = Ability(ability_id='123', name='Test Ability', executors=[executor_linux], technique_name='collection', technique_id='1', description='', privilege='', tactic='discovery', plugin='testplugin') loop.run_until_complete(BaseService.get_service('data_svc').store(ability)) return ability
async def update_scores(operation, increment, used, facts): knowledge_svc_handle = BaseService.get_service('knowledge_svc') for uf in used: all_facts = await operation.all_facts() if operation else facts for found_fact in all_facts: if found_fact.unique == uf.unique: found_fact.score += increment await knowledge_svc_handle.update_fact( dict(trait=found_fact.trait, value=found_fact.value, source=found_fact.source), dict(score=found_fact.score)) break
def test_adversary(event_loop): expected_adversary = { 'name': 'ad-hoc', 'description': 'an empty adversary profile', 'adversary_id': 'ad-hoc', 'objective': '495a9828-cab1-44dd-a0ca-66e58177d8cc', 'tags': [], 'has_repeatable_abilities': False } test_adversary = AdversarySchema().load(expected_adversary) event_loop.run_until_complete( BaseService.get_service('data_svc').store(test_adversary)) return test_adversary
async def save_fact(self, operation, fact, score, relationship): knowledge_svc_handle = BaseService.get_service('knowledge_svc') all_facts = await operation.all_facts() if operation else self.facts source = operation.id if operation else self.id rl = [relationship] if relationship else [] if all([fact.trait, fact.value]): if operation and operation.source: if any([(fact.trait, fact.value) == (x.trait, x.value) for x in await knowledge_svc_handle.get_facts( criteria=dict(source=operation.source.id))]): source = operation.source.id fact.source = source # Manual addition to ensure the check works correctly if not await knowledge_svc_handle.check_fact_exists( fact, all_facts): f_gen = Fact(trait=fact.trait, value=fact.value, source=source, score=score, collected_by=[self.paw], technique_id=self.ability.technique_id, links=[self.id], relationships=rl, origin_type=OriginType.LEARNED) self.facts.append(f_gen) await knowledge_svc_handle.add_fact(f_gen) else: existing_fact = (await knowledge_svc_handle.get_facts( criteria=dict( trait=fact.trait, value=fact.value, source=fact.source) ))[0] if self.id not in existing_fact.links: existing_fact.links.append(self.id) if relationship not in existing_fact.relationships: existing_fact.relationships.append(relationship) if self.paw not in existing_fact.collected_by: existing_fact.collected_by.append(self.paw) await knowledge_svc_handle.update_fact( criteria=dict(trait=fact.trait, value=fact.value, source=fact.source), updates=dict(links=existing_fact.links, relationships=existing_fact.relationships, collected_by=existing_fact.collected_by)) existing_local_record = [ x for x in self.facts if x.trait == fact.trait and x.value == fact.value ] if existing_local_record: existing_local_record[0].links = existing_fact.links else: self.facts.append(existing_fact)