class OperationSchema(ma.Schema):
id = ma.fields.String()
name = ma.fields.String(required=True)
host_group = ma.fields.List(ma.fields.Nested(AgentSchema()),
attribute='agents',
dump_only=True)
adversary = ma.fields.Nested(AdversarySchema())
jitter = ma.fields.String()
planner = ma.fields.Nested(PlannerSchema())
start = ma.fields.DateTime(format=BaseObject.TIME_FORMAT, dump_only=True)
state = ma.fields.String()
obfuscator = ma.fields.String()
autonomous = ma.fields.Integer()
chain = ma.fields.Function(lambda obj: [lnk.display for lnk in obj.chain])
auto_close = ma.fields.Boolean()
visibility = ma.fields.Integer()
objective = ma.fields.Nested(ObjectiveSchema())
use_learning_parsers = ma.fields.Boolean()
group = ma.fields.String(missing='')
source = ma.fields.Nested(SourceSchema())
@ma.pre_load()
def remove_properties(self, data, **_):
data.pop('host_group', None)
data.pop('start', None)
data.pop('chain', None)
data.pop('objective', None)
return data
@ma.post_load
def build_operation(self, data, **kwargs):
return None if kwargs.get('partial') is True else Operation(**data)
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 _construct_and_dump_adversary(self, adversary_id: str):
adv = await self.services['data_svc'].locate(
'adversaries', match=dict(adversary_id=adversary_id))
if not adv:
adv = Adversary.load(
dict(adversary_id='ad-hoc',
name='ad-hoc',
description='an empty adversary profile',
atomic_ordering=[]))
else:
adv = adv[0]
return AdversarySchema().dump(adv)
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
def test_operation(test_adversary, test_planner, test_source):
expected_operation = {
'name': '123',
'adversary': AdversarySchema().dump(test_adversary),
'state': 'paused',
'id': '123',
'group': 'red',
'autonomous': 0,
'planner': PlannerSchema().dump(test_planner),
'source': SourceSchema().dump(test_source),
'jitter': '2/8',
'visibility': 50,
'auto_close': False,
'obfuscator': 'plain-text',
'use_learning_parsers': False
}
return expected_operation
class OperationSchema(ma.Schema):
id = ma.fields.Integer()
name = ma.fields.String()
host_group = ma.fields.List(ma.fields.Nested(AgentSchema()), attribute='agents')
adversary = ma.fields.Nested(AdversarySchema())
jitter = ma.fields.String()
planner = ma.fields.Nested(PlannerSchema())
start = ma.fields.DateTime(format='%Y-%m-%d %H:%M:%S')
state = ma.fields.String()
obfuscator = ma.fields.String()
autonomous = ma.fields.Integer()
chain = ma.fields.Function(lambda obj: [lnk.display for lnk in obj.chain])
auto_close = ma.fields.Boolean()
visibility = ma.fields.Integer()
@ma.post_load
def build_planner(self, data, **_):
return Operation(**data)
class AdversaryApi(BaseObjectApi):
def __init__(self, services):
super().__init__(description='adversary',
obj_class=Adversary,
schema=AdversarySchema,
ram_key='adversaries',
id_property='adversary_id',
auth_svc=services['auth_svc'])
self._api_manager = AdversaryApiManager(data_svc=services['data_svc'],
file_svc=services['file_svc'])
def add_routes(self, app: web.Application):
router = app.router
adversaries_by_id_path = '/adversaries/{adversary_id}'
router.add_get('/adversaries', self.get_adversaries)
router.add_get(adversaries_by_id_path, self.get_adversary_by_id)
router.add_post('/adversaries', self.create_adversary)
router.add_patch(adversaries_by_id_path, self.update_adversary)
router.add_put(adversaries_by_id_path, self.create_or_update_adversary)
router.add_delete(adversaries_by_id_path, self.delete_adversary)
@aiohttp_apispec.docs(tags=['adversaries'])
@aiohttp_apispec.querystring_schema(BaseGetAllQuerySchema)
@aiohttp_apispec.response_schema(AdversarySchema(many=True, partial=True))
async def get_adversaries(self, request: web.Request):
adversaries = await self.get_all_objects(request)
return web.json_response(adversaries)
@aiohttp_apispec.docs(tags=['adversaries'])
@aiohttp_apispec.querystring_schema(BaseGetOneQuerySchema)
@aiohttp_apispec.response_schema(AdversarySchema(partial=True))
async def get_adversary_by_id(self, request: web.Request):
adversary = await self.get_object(request)
return web.json_response(adversary)
@aiohttp_apispec.docs(tags=['adversaries'])
@aiohttp_apispec.request_schema(AdversarySchema)
@aiohttp_apispec.response_schema(AdversarySchema)
async def create_adversary(self, request: web.Request):
adversary = await self.create_on_disk_object(request)
adversary = await self._api_manager.verify_adversary(adversary)
return web.json_response(adversary.display)
@aiohttp_apispec.docs(tags=['adversaries'])
@aiohttp_apispec.request_schema(
AdversarySchema(partial=True, exclude=['adversary_id']))
@aiohttp_apispec.response_schema(AdversarySchema)
async def update_adversary(self, request: web.Request):
adversary = await self.update_on_disk_object(request)
adversary = await self._api_manager.verify_adversary(adversary)
return web.json_response(adversary.display)
@aiohttp_apispec.docs(tags=['adversaries'])
@aiohttp_apispec.request_schema(AdversarySchema(partial=True))
@aiohttp_apispec.response_schema(AdversarySchema)
async def create_or_update_adversary(self, request: web.Request):
adversary = await self.create_or_update_on_disk_object(request)
adversary = await self._api_manager.verify_adversary(adversary)
return web.json_response(adversary.display)
@aiohttp_apispec.docs(tags=['adversaries'])
@aiohttp_apispec.response_schema(AdversarySchema)
async def delete_adversary(self, request: web.Request):
await self.delete_on_disk_object(request)
return web.HTTPNoContent()
class AdversaryApi(BaseObjectApi):
def __init__(self, services):
super().__init__(description='adversary', obj_class=Adversary, schema=AdversarySchema, ram_key='adversaries',
id_property='adversary_id', auth_svc=services['auth_svc'])
self._api_manager = AdversaryApiManager(data_svc=services['data_svc'], file_svc=services['file_svc'])
def add_routes(self, app: web.Application):
router = app.router
adversaries_by_id_path = '/adversaries/{adversary_id}'
router.add_get('/adversaries', self.get_adversaries)
router.add_get(adversaries_by_id_path, self.get_adversary_by_id)
router.add_post('/adversaries', self.create_adversary)
router.add_patch(adversaries_by_id_path, self.update_adversary)
router.add_put(adversaries_by_id_path, self.create_or_update_adversary)
router.add_delete(adversaries_by_id_path, self.delete_adversary)
@aiohttp_apispec.docs(tags=['adversaries'],
summary='Retrieve all adversaries',
description='Returns a list of all available adversaries in the system, including plugin, name, description, '
'and atomic ordering. Supply fields from the `AdversarySchema` to the include and exclude fields of the '
'`BaseGetAllQuerySchema` in the request body to filter retrieved adversaries.')
@aiohttp_apispec.querystring_schema(BaseGetAllQuerySchema)
@aiohttp_apispec.response_schema(AdversarySchema(many=True, partial=True),
description='Returns a list in `AdversarySchema` format of all available adversaries in the system.')
async def get_adversaries(self, request: web.Request):
adversaries = await self.get_all_objects(request)
return web.json_response(adversaries)
@aiohttp_apispec.docs(tags=['adversaries'],
summary='Retrieve adversary by ID',
description='Retrieve one adversary by ID. Use fields from the `AdversarySchema` in '
'the request body to filter retrieved adversary.',
parameters=[{
'in': 'path',
'name': 'adversary_id',
'schema': {'type': 'string'},
'required': 'true',
'description': 'UUID of the adversary to be retrieved'}])
@aiohttp_apispec.querystring_schema(BaseGetOneQuerySchema)
@aiohttp_apispec.response_schema(AdversarySchema(partial=True),
description='Returns single adversary in AdversarySchema format.')
async def get_adversary_by_id(self, request: web.Request):
adversary = await self.get_object(request)
return web.json_response(adversary)
@aiohttp_apispec.docs(tags=['adversaries'],
summary='Create a new adversary',
description='Create a new adversary using the format provided in the `AdversarySchema`.')
@aiohttp_apispec.request_schema(AdversarySchema)
@aiohttp_apispec.response_schema(AdversarySchema, description='A single adversary in AdversarySchema format.')
async def create_adversary(self, request: web.Request):
adversary = await self.create_on_disk_object(request)
adversary = await self._api_manager.verify_adversary(adversary)
return web.json_response(adversary.display)
@aiohttp_apispec.docs(tags=['adversaries'],
summary='Update an adversary',
description='Update an adversary using fields from the `AdversarySchema` in the request body.',
parameters=[{
'in': 'path',
'name': 'adversary_id',
'schema': {'type': 'string'},
'required': 'true',
'description': 'UUID of the adversary to be updated'
}])
@aiohttp_apispec.docs(tags=['adversaries'])
@aiohttp_apispec.request_schema(AdversarySchema(partial=True, exclude=['adversary_id']))
@aiohttp_apispec.response_schema(AdversarySchema,
description='The updated adversary in AdversarySchema format.')
async def update_adversary(self, request: web.Request):
adversary = await self.update_on_disk_object(request)
adversary = await self._api_manager.verify_adversary(adversary)
return web.json_response(adversary.display)
@aiohttp_apispec.docs(tags=['adversaries'],
summary='Create or update an adversary',
description='Attempt to update an adversaries using fields from the `AdversarySchema` '
'in the request body. If the adversary does not already exist, '
'then create a new one using the `AdversarySchema` format.',
parameters=[{
'in': 'path',
'name': 'adversary_id',
'schema': {'type': 'string'},
'required': 'true',
'description': 'UUID of the adversary to be created or updated'
}])
@aiohttp_apispec.request_schema(AdversarySchema(partial=True))
@aiohttp_apispec.response_schema(AdversarySchema,
description='A single adversary, either newly created or updated, in AdversarySchema format.')
async def create_or_update_adversary(self, request: web.Request):
adversary = await self.create_or_update_on_disk_object(request)
adversary = await self._api_manager.verify_adversary(adversary)
return web.json_response(adversary.display)
@aiohttp_apispec.docs(tags=['adversaries'], summary='Deletes an adversary.',
description='Deletes an existing adversary.',
parameters=[{
'in': 'path',
'name': 'adversary_id',
'schema': {'type': 'string'},
'required': 'true',
'description': 'UUID of the adversary to be retrieved'
}])
@aiohttp_apispec.response_schema(AdversarySchema(partial=True), code=204,
description='HTTP 204 Status Code (No Content)')
async def delete_adversary(self, request: web.Request):
await self.delete_on_disk_object(request)
return web.HTTPNoContent()
async def create_on_disk_object(self, request: web.Request):
data = await request.json()
data.pop('id', None)
await self._error_if_object_with_id_exists(data.get(self.id_property))
access = await self.get_request_permissions(request)
obj = await self._api_manager.create_on_disk_object(data, access, self.ram_key, self.id_property,
self.obj_class)
return obj
async def _parse_common_data_from_request(self, request) -> (dict, dict, str, dict, dict):
data = {}
raw_body = await request.read()
if raw_body:
data = json.loads(raw_body)
data.pop('id', None)
obj_id = request.match_info.get(self.id_property, '')
if obj_id:
data[self.id_property] = obj_id
access = await self.get_request_permissions(request)
query = {self.id_property: obj_id}
search = {**query, **access}
return data, access, obj_id, query, search