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