diff --git a/app/api/v2/handlers/fact_source_api.py b/app/api/v2/handlers/fact_source_api.py index 4ff87e791..1112587e5 100644 --- a/app/api/v2/handlers/fact_source_api.py +++ b/app/api/v2/handlers/fact_source_api.py @@ -2,7 +2,7 @@ from aiohttp import web from app.api.v2.handlers.base_object_api import BaseObjectApi -from app.api.v2.managers.base_api_manager import BaseApiManager +from app.api.v2.managers.fact_source_manager import FactSourceApiManager from app.api.v2.schemas.base_schemas import BaseGetAllQuerySchema, BaseGetOneQuerySchema from app.objects.c_source import Source, SourceSchema @@ -11,7 +11,7 @@ class FactSourceApi(BaseObjectApi): def __init__(self, services): super().__init__(description='Fact Source', obj_class=Source, schema=SourceSchema, ram_key='sources', id_property='id', auth_svc=services['auth_svc']) - self._api_manager = BaseApiManager(data_svc=services['data_svc'], file_svc=services['file_svc']) + self._api_manager = FactSourceApiManager(data_svc=services['data_svc'], file_svc=services['file_svc'], knowledge_svc=services['knowledge_svc']) def add_routes(self, app: web.Application): router = app.router @@ -102,4 +102,7 @@ async def create_or_update_source(self, request: web.Request): @aiohttp_apispec.response_schema(SourceSchema, description='Returns DELETE status.') async def delete_source(self, request: web.Request): await self.delete_on_disk_object(request) + knowledge_svc_handle = self._api_manager.knowledge_svc + await knowledge_svc_handle.delete_fact(criteria=dict(source=request.match_info.get('id'))) + await knowledge_svc_handle.delete_relationship(criteria=dict(origin=request.match_info.get('id'))) return web.HTTPNoContent() diff --git a/app/api/v2/handlers/operation_api.py b/app/api/v2/handlers/operation_api.py index e8881d0f8..17bf13645 100644 --- a/app/api/v2/handlers/operation_api.py +++ b/app/api/v2/handlers/operation_api.py @@ -132,6 +132,9 @@ async def update_operation(self, request: web.Request): description='There is an empty response from a successful delete request.') async def delete_operation(self, request: web.Request): await self.delete_object(request) + knowledge_svc_handle = self._api_manager.knowledge_svc + await knowledge_svc_handle.delete_fact(criteria=dict(source=request.match_info.get('id'))) + await knowledge_svc_handle.delete_relationship(criteria=dict(origin=request.match_info.get('id'))) return web.HTTPNoContent() @aiohttp_apispec.docs(tags=['operations'], diff --git a/app/api/v2/managers/fact_source_manager.py b/app/api/v2/managers/fact_source_manager.py new file mode 100644 index 000000000..2a0afbb97 --- /dev/null +++ b/app/api/v2/managers/fact_source_manager.py @@ -0,0 +1,6 @@ +from app.api.v2.managers.base_api_manager import BaseApiManager + +class FactSourceApiManager(BaseApiManager): + def __init__(self, data_svc, file_svc, knowledge_svc): + super().__init__(data_svc=data_svc, file_svc=file_svc) + self.knowledge_svc = knowledge_svc diff --git a/app/api/v2/managers/operation_api_manager.py b/app/api/v2/managers/operation_api_manager.py index ada8be28f..e7409a62f 100644 --- a/app/api/v2/managers/operation_api_manager.py +++ b/app/api/v2/managers/operation_api_manager.py @@ -21,6 +21,7 @@ class OperationApiManager(BaseApiManager): def __init__(self, services): super().__init__(data_svc=services['data_svc'], file_svc=services['file_svc']) self.services = services + self.knowledge_svc = services['knowledge_svc'] async def get_operation_report(self, operation_id: str, access: dict, output: bool): operation = await self.get_operation_object(operation_id, access) diff --git a/app/objects/c_operation.py b/app/objects/c_operation.py index d7b553773..7f961d8f1 100644 --- a/app/objects/c_operation.py +++ b/app/objects/c_operation.py @@ -187,9 +187,10 @@ def update_untrusted_agents(self, agent): async def all_facts(self): knowledge_svc_handle = BaseService.get_service('knowledge_svc') + data_svc_handle = BaseService.get_service('data_svc') seeded_facts = [] if self.source: - seeded_facts = await knowledge_svc_handle.get_facts(criteria=dict(source=self.source.id)) + seeded_facts = await data_svc_handle.get_facts_from_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 diff --git a/app/service/data_svc.py b/app/service/data_svc.py index 5f23a0d58..f768823bd 100644 --- a/app/service/data_svc.py +++ b/app/service/data_svc.py @@ -481,3 +481,12 @@ async def _verify_adversary_profiles(self): def _get_plugin_name(self, filename): plugin_path = pathlib.PurePath(filename).parts return plugin_path[1] if 'plugins' in plugin_path else '' + + async def get_facts_from_source(self, fact_source_id): + fact_sources = await self.locate('sources', match=dict(id=fact_source_id)) + if len(fact_sources) == 0: + return [] + elif len(fact_sources) > 1: + self.log.error('Found multiple fact sources with the same id', fact_source_id) + return [] + return fact_sources[0].facts