class TestPlanningService(TestBase):
def setUp(self):
self.initialize()
self.ability = Ability(ability_id='123',
executor='sh',
test=BaseWorld.encode_string('mkdir test'),
cleanup=BaseWorld.encode_string('rm -rf test'))
self.agent = Agent(sleep_min=1, sleep_max=2, watchdog=0)
self.operation = Operation(name='test1',
agents=self.agent,
adversary='hunter')
self.run_async(self.data_svc.store(self.ability))
def test_get_cleanup_links(self):
self.operation.add_link(
Link(operation=self.operation,
command='',
paw=self.agent.paw,
ability=self.ability,
status=0))
links = self.run_async(
self.planning_svc.get_cleanup_links(operation=self.operation,
agent=self.agent))
link_list = list(links)
self.assertEqual(len(link_list), 1)
self.assertEqual(link_list[0].command, self.ability.cleanup)
class TestLearningSvc(TestBase):
def setUp(self):
self.initialize()
self.ability = self.run_async(
self.data_svc.store(
Ability(ability_id='123',
tactic='discovery',
technique_id='T1033',
technique='Find',
name='test',
test='d2hvYW1pCg==',
description='find active user',
cleanup='',
executor='sh',
platform='darwin',
payload='wifi.sh',
parsers=[],
requirements=[],
privilege=None)))
self.operation = Operation(name='sample', agents=None, adversary=None)
self.run_async(self.data_svc.store(self.operation))
def test_learn(self):
link = Link(operation=self.operation.id,
ability=self.ability,
command=None,
paw=None)
self.operation.add_link(link)
self.run_async(
self.learning_svc.learn(
link=link,
blob=BaseWorld.encode_string(
'i contain 1 ip address 192.168.0.1 and one file /etc/host.txt. that is all.'
)))
self.assertEqual(2, len(link.facts))
def test_build_relationships(self):
self.learning_svc.model.add(
frozenset({'host.user.name', 'target.org.name'}))
self.learning_svc.model.add(
frozenset(
{'host.file.extension', 'host.user.name', 'domain.user.name'}))
facts = [
Fact(trait='target.org.name', value='something'),
Fact(trait='host.user.name', value='admin'),
Fact(trait='host.user.name', value='root'),
Fact(trait='domain.user.name', value='user'),
Fact(trait='not.really.here', value='should never be found')
]
link = Link(operation=self.operation.id,
ability=self.ability,
command=None,
paw=None)
self.run_async(self.learning_svc._build_relationships(link, facts))
self.assertEqual(4, len(link.relationships))
class ResponseService(BaseService):
def __init__(self, services):
self.log = self.add_service('response_svc', self)
self.data_svc = services.get('data_svc')
self.rest_svc = services.get('rest_svc')
self.agents = []
self.adversary = None
self.abilities = []
self.op = None
@template('response.html')
async def splash(self, request):
abilities = [
a for a in await self.data_svc.locate('abilities')
if await a.which_plugin() == 'response'
]
adversaries = [
a for a in await self.data_svc.locate('adversaries')
if await a.which_plugin() == 'response'
]
return dict(abilities=abilities, adversaries=adversaries)
@staticmethod
async def register_handler(event_svc):
await event_svc.observe_event('link/completed', handle_link_completed)
async def respond_to_pid(self, pid, agent):
available_agents = await self.get_available_agents(agent)
if not available_agents:
return
total_facts = []
total_links = []
for blue_agent in available_agents:
agent_facts, agent_links = await self.run_abilities_on_agent(
blue_agent, pid)
total_facts.extend(agent_facts)
total_links.extend(agent_links)
await self.save_to_operation(total_facts, total_links)
async def get_available_agents(self, agent_to_match):
await self.refresh_blue_agents_abilities()
available_agents = [
a for a in self.agents if a.host == agent_to_match.host
]
if not available_agents:
self.log.debug('No available blue agents to respond to red action')
return []
return available_agents
async def refresh_blue_agents_abilities(self):
self.agents = await self.data_svc.locate(
'agents', match=dict(access=self.Access.BLUE))
self.adversary = (await self.data_svc.locate(
'adversaries', match=dict(adversary_id=BLUE_ADVERSARY)))[0]
self.abilities = []
for a in self.adversary.atomic_ordering:
if a not in self.abilities:
self.abilities.append(a)
async def run_abilities_on_agent(self, blue_agent, original_pid):
facts = [Fact(trait='host.process.id', value=original_pid)]
links = []
for ability_id in self.abilities:
ability_facts, ability_links = await self.run_ability_on_agent(
blue_agent, ability_id, facts, original_pid)
links.extend(ability_links)
facts.extend(ability_facts)
return facts, links
async def run_ability_on_agent(self, blue_agent, ability_id, agent_facts,
original_pid):
links = await self.rest_svc.task_agent_with_ability(
blue_agent.paw, ability_id, agent_facts)
await self.wait_for_link_completion(links, blue_agent)
ability_facts = []
for link in links:
link.pin = int(original_pid)
unique_facts = link.facts[1:]
ability_facts.extend(unique_facts)
return ability_facts, links
@staticmethod
async def wait_for_link_completion(links, agent):
for link in links:
while not link.finish or link.can_ignore():
await asyncio.sleep(3)
if not agent.trusted:
break
async def create_fact_source(self, facts):
source_id = str(uuid.uuid4())
source_name = 'blue-pid-{}'.format(source_id)
return Source(id=source_id, name=source_name, facts=facts)
async def save_to_operation(self, facts, links):
if not self.op or await self.op.is_finished():
source = await self.create_fact_source(facts)
await self.create_operation(links=links, source=source)
else:
await self.update_operation(links)
await self.get_service('data_svc').store(self.op)
async def create_operation(self, links, source):
planner = (await self.get_service('data_svc').locate(
'planners', match=dict(name='sequential')))[0]
await self.get_service('data_svc').store(source)
self.op = Operation(name=BLUE_OP_NAME,
agents=self.agents,
adversary=self.adversary,
source=source,
access=self.Access.BLUE,
planner=planner,
state='running',
auto_close=False,
jitter='1/4')
self.op.set_start_details()
await self.update_operation(links)
async def update_operation(self, links):
for link in links:
link.operation = self.op.id
self.op.add_link(link)
