async def update_on_disk_object(self, obj: Any, data: dict, ram_key: str,
id_property: str, obj_class: type):
obj_id = getattr(obj, id_property)
file_path = await self._get_existing_object_file_path(obj_id, ram_key)
existing_obj_data = AbilitySchema().dump(obj)
existing_obj_data.update(data)
self._validate_ability_data(create=False, data=existing_obj_data)
if existing_obj_data.get('tactic') not in file_path:
await self.remove_object_from_disk_by_id(obj_id, ram_key)
file_path = self._create_ability_filepath(data.get('tactic'),
obj_id)
await self._save_and_reload_object(file_path, existing_obj_data,
obj_class, obj.access)
return next(self.find_objects(ram_key, {id_property: obj_id}))
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
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
class LinkSchema(ma.Schema):
class Meta:
unknown = ma.EXCLUDE
id = ma.fields.String(missing='')
paw = ma.fields.String()
command = ma.fields.String()
status = ma.fields.Integer(missing=-3)
score = ma.fields.Integer(missing=0)
jitter = ma.fields.Integer(missing=0)
decide = ma.fields.DateTime(format='%Y-%m-%d %H:%M:%S')
pin = ma.fields.Integer(missing=0)
pid = ma.fields.String()
facts = ma.fields.List(ma.fields.Nested(FactSchema()))
relationships = ma.fields.List(ma.fields.Nested(RelationshipSchema()))
used = ma.fields.List(ma.fields.Nested(FactSchema()))
unique = ma.fields.String()
collect = ma.fields.DateTime(format='%Y-%m-%d %H:%M:%S', default='')
finish = ma.fields.String()
ability = ma.fields.Nested(AbilitySchema())
executor = ma.fields.Nested(ExecutorSchema())
cleanup = ma.fields.Integer(missing=0)
visibility = ma.fields.Nested(VisibilitySchema)
host = ma.fields.String(missing=None)
output = ma.fields.String()
deadman = ma.fields.Boolean()
agent_reported_time = ma.fields.DateTime(format='%Y-%m-%d %H:%M:%S',
missing=None)
@ma.pre_load()
def fix_ability(self, link, **_):
if 'ability' in link and isinstance(link['ability'], Ability):
ability = link.pop('ability')
link['ability'] = ability.schema.dump(ability)
return link
@ma.pre_load()
def fix_executor(self, link, **_):
if 'executor' in link and isinstance(link['executor'], Executor):
executor = link.pop('executor')
link['executor'] = executor.schema.dump(executor)
return link
@ma.post_load()
def build_link(self, data, **_):
return Link(**data)
@ma.post_dump()
def prepare_dump(self, data, **_):
if data.get('agent_reported_time', None) is None:
data.pop('agent_reported_time', None)
return data
def build_ability(self, data: dict, executor: Executor):
if not data.get('ability_id'):
data['ability_id'] = str(uuid.uuid4())
if not data.get('tactic'):
data['tactic'] = 'auto-generated'
if not data.get('technique_id'):
data['technique_id'] = 'auto-generated'
if not data.get('technique_name'):
data['technique_name'] = 'auto-generated'
if not data.get('name'):
data['name'] = 'Manual Command'
if not data.get('description'):
data['description'] = 'Manual command ability'
data['executors'] = [ExecutorSchema().dump(executor)]
ability = AbilitySchema().load(data)
return ability
class LinkSchema(ma.Schema):
class Meta:
unknown = ma.EXCLUDE
id = ma.fields.String(missing='')
paw = ma.fields.String()
command = ma.fields.String()
status = ma.fields.Integer(missing=-3)
score = ma.fields.Integer(missing=0)
jitter = ma.fields.Integer(missing=0)
decide = ma.fields.DateTime(format='%Y-%m-%d %H:%M:%S')
pin = ma.fields.Integer(missing=0)
pid = ma.fields.String()
facts = ma.fields.List(ma.fields.Nested(FactSchema()))
relationships = ma.fields.List(ma.fields.Nested(RelationshipSchema()))
used = ma.fields.List(ma.fields.Nested(FactSchema()))
unique = ma.fields.String()
collect = ma.fields.DateTime(format='%Y-%m-%d %H:%M:%S', default='')
finish = ma.fields.String()
ability = ma.fields.Nested(AbilitySchema())
cleanup = ma.fields.Integer(missing=0)
visibility = ma.fields.Nested(VisibilitySchema)
host = ma.fields.String(missing=None)
output = ma.fields.String()
deadman = ma.fields.Boolean()
@ma.pre_load()
def fix_ability(self, link, **_):
if 'ability' in link and isinstance(link['ability'], Ability):
link_input = link.pop('ability')
link['ability'] = link_input.schema.dump(link_input)
return link
@ma.post_load()
def build_link(self, data, **_):
return Link(**data)
@ma.pre_dump()
def prepare_link(self, data, **_):
# temp - can be simplified with AbilitySchema
data.executor = data.ability.executor if isinstance(
data.ability, Ability) else data.ability['executor']
return data
class AbilityApi(BaseObjectApi):
def __init__(self, services):
super().__init__(description='ability',
obj_class=Ability,
schema=AbilitySchema,
ram_key='abilities',
id_property='ability_id',
auth_svc=services['auth_svc'])
self._api_manager = AbilityApiManager(data_svc=services['data_svc'],
file_svc=services['file_svc'])
def add_routes(self, app: web.Application):
router = app.router
router.add_get('/abilities', self.get_abilities)
router.add_get('/abilities/{ability_id}', self.get_ability_by_id)
router.add_post('/abilities', self.create_ability)
router.add_put('/abilities/{ability_id}',
self.create_or_update_ability)
router.add_patch('/abilities/{ability_id}', self.update_ability)
router.add_delete('/abilities/{ability_id}', self.delete_ability)
@aiohttp_apispec.docs(
tags=['abilities'],
summary='Get all abilities.',
description='Provides a list of all available abilities.')
@aiohttp_apispec.querystring_schema(BaseGetAllQuerySchema)
@aiohttp_apispec.response_schema(
AbilitySchema(many=True, partial=True),
description='Returns a list of all abilities.')
async def get_abilities(self, request: web.Request):
abilities = await self.get_all_objects(request)
return web.json_response(abilities)
@aiohttp_apispec.docs(
tags=['abilities'],
summary='Get an ability.',
description='Provides one ability based on its ability id.',
parameters=[{
'in': 'path',
'name': 'ability_id',
'schema': {
'type': 'string'
},
'required': 'true',
'description': 'UUID of the Ability to be retrieved'
}])
@aiohttp_apispec.querystring_schema(BaseGetOneQuerySchema)
@aiohttp_apispec.response_schema(
AbilitySchema(partial=True),
description='JSON dictionary representation of the existing Ability.')
async def get_ability_by_id(self, request: web.Request):
ability = await self.get_object(request)
return web.json_response(ability)
@aiohttp_apispec.docs(
tags=['abilities'],
summary='Creates a new ability.',
description='Creates a new adversary based on the `AbilitySchema`. '
'"name", "tactic", and "executors" are all required fields.')
@aiohttp_apispec.request_schema(AbilitySchema)
@aiohttp_apispec.response_schema(
AbilitySchema,
description='JSON dictionary representation of the created Ability.')
async def create_ability(self, request: web.Request):
ability = await self.create_on_disk_object(request)
return web.json_response(ability.display)
@aiohttp_apispec.docs(
tags=['abilities'],
summary='Replaces an existing ability.',
description=
'Replaces an ability based on the `AbilitySchema` values provided '
'in the message body. "name", "tactic", and "executors" '
'are all required fields.',
parameters=[{
'in': 'path',
'name': 'ability_id',
'schema': {
'type': 'string'
},
'required': 'true',
'description': 'UUID of the Ability to be retrieved'
}])
@aiohttp_apispec.request_schema(AbilitySchema(partial=True))
@aiohttp_apispec.response_schema(
AbilitySchema,
description='JSON dictionary representation of the replaced Ability.')
async def create_or_update_ability(self, request: web.Request):
ability = await self.create_or_update_on_disk_object(request)
return web.json_response(ability.display)
@aiohttp_apispec.docs(
tags=['abilities'],
summary='Updates an existing ability.',
description=
'Updates an ability based on the `AbilitySchema` values provided '
'in the message body.',
parameters=[{
'in': 'path',
'name': 'ability_id',
'schema': {
'type': 'string'
},
'required': 'true',
'description': 'UUID of the Ability to be retrieved'
}])
@aiohttp_apispec.request_schema(
AbilitySchema(partial=True,
exclude=[
'ability_id', 'requirements', 'additional_info',
'access'
]))
@aiohttp_apispec.response_schema(
AbilitySchema,
description='JSON dictionary representation of the replaced Ability.')
async def update_ability(self, request: web.Request):
ability = await self.update_on_disk_object(request)
return web.json_response(ability.display)
@aiohttp_apispec.docs(tags=['abilities'],
summary='Deletes an ability.',
description='Deletes an existing ability.',
parameters=[{
'in':
'path',
'name':
'ability_id',
'schema': {
'type': 'string'
},
'required':
'true',
'description':
'UUID of the Ability to be retrieved'
}])
@aiohttp_apispec.response_schema(
AbilitySchema,
code=204,
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()
class AbilityApi(BaseObjectApi):
def __init__(self, services):
super().__init__(description='ability',
obj_class=Ability,
schema=AbilitySchema,
ram_key='abilities',
id_property='ability_id',
auth_svc=services['auth_svc'])
self._api_manager = AbilityApiManager(data_svc=services['data_svc'],
file_svc=services['file_svc'])
def add_routes(self, app: web.Application):
router = app.router
router.add_get('/abilities', self.get_abilities)
router.add_get('/abilities/{ability_id}', self.get_ability_by_id)
router.add_post('/abilities', self.create_ability)
router.add_put('/abilities/{ability_id}',
self.create_or_update_ability)
router.add_patch('/abilities/{ability_id}', self.update_ability)
router.add_delete('/abilities/{ability_id}', self.delete_ability)
@aiohttp_apispec.docs(tags=['abilities'])
@aiohttp_apispec.querystring_schema(BaseGetAllQuerySchema)
@aiohttp_apispec.response_schema(AbilitySchema(many=True, partial=True))
async def get_abilities(self, request: web.Request):
abilities = await self.get_all_objects(request)
return web.json_response(abilities)
@aiohttp_apispec.docs(tags=['abilities'])
@aiohttp_apispec.querystring_schema(BaseGetOneQuerySchema)
@aiohttp_apispec.response_schema(AbilitySchema(partial=True))
async def get_ability_by_id(self, request: web.Request):
ability = await self.get_object(request)
return web.json_response(ability)
@aiohttp_apispec.docs(tags=['abilities'])
@aiohttp_apispec.request_schema(AbilitySchema)
@aiohttp_apispec.response_schema(AbilitySchema)
async def create_ability(self, request: web.Request):
ability = await self.create_on_disk_object(request)
return web.json_response(ability.display)
@aiohttp_apispec.docs(tags=['abilities'])
@aiohttp_apispec.request_schema(AbilitySchema(partial=True))
@aiohttp_apispec.response_schema(AbilitySchema)
async def create_or_update_ability(self, request: web.Request):
ability = await self.create_or_update_on_disk_object(request)
return web.json_response(ability.display)
@aiohttp_apispec.docs(tags=['abilities'])
@aiohttp_apispec.request_schema(AbilitySchema(partial=True))
@aiohttp_apispec.response_schema(AbilitySchema)
async def update_ability(self, request: web.Request):
ability = await self.update_on_disk_object(request)
return web.json_response(ability.display)
@aiohttp_apispec.docs(tags=['abilities'])
@aiohttp_apispec.response_schema(AbilitySchema)
async def delete_ability(self, request: web.Request):
await self.delete_on_disk_object(request)
return web.HTTPNoContent()
class LinkSchema(ma.Schema):
class Meta:
unknown = ma.EXCLUDE
id = ma.fields.String(missing='')
paw = ma.fields.String()
command = ma.fields.String()
status = ma.fields.Integer(missing=-3)
score = ma.fields.Integer(missing=0)
jitter = ma.fields.Integer(missing=0)
decide = ma.fields.DateTime(format=BaseObject.TIME_FORMAT)
pin = ma.fields.Integer(missing=0)
pid = ma.fields.String()
facts = ma.fields.List(ma.fields.Nested(FactSchema()))
relationships = ma.fields.List(ma.fields.Nested(RelationshipSchema()))
used = ma.fields.List(ma.fields.Nested(FactSchema()))
unique = ma.fields.String()
collect = ma.fields.DateTime(format=BaseObject.TIME_FORMAT, default='')
finish = ma.fields.String()
ability = ma.fields.Nested(AbilitySchema())
executor = ma.fields.Nested(ExecutorSchema())
cleanup = ma.fields.Integer(missing=0)
visibility = ma.fields.Nested(VisibilitySchema())
host = ma.fields.String(missing=None)
output = ma.fields.String()
deadman = ma.fields.Boolean()
agent_reported_time = ma.fields.DateTime(format=BaseObject.TIME_FORMAT,
missing=None)
@ma.pre_load()
def fix_ability(self, link, **_):
if 'ability' in link and isinstance(link['ability'], Ability):
ability = link.pop('ability')
link['ability'] = ability.schema.dump(ability)
return link
@ma.pre_load()
def fix_executor(self, link, **_):
if 'executor' in link and isinstance(link['executor'], Executor):
executor = link.pop('executor')
link['executor'] = executor.schema.dump(executor)
return link
@ma.pre_load()
def remove_properties(self, data, **_):
data.pop('unique', None)
data.pop('decide', None)
data.pop('pid', None)
data.pop('facts', None)
data.pop('collect', None)
data.pop('finish', None)
data.pop('visibility', None)
data.pop('output', None)
data.pop('used.unique', None)
return data
@ma.post_load()
def build_link(self, data, **kwargs):
return None if kwargs.get('partial') is True else Link(**data)
@ma.post_dump()
def prepare_dump(self, data, **_):
if data.get('agent_reported_time', None) is None:
data.pop('agent_reported_time', None)
return data