From d1140f96279f735adf7b470d0c5f9b51f0ba46b1 Mon Sep 17 00:00:00 2001 From: bjeffries Date: Tue, 3 Dec 2024 13:07:45 -0500 Subject: [PATCH 1/7] generalized hook in planning_svc and added missing check in op_api_manager --- app/api/v2/managers/operation_api_manager.py | 67 ++------------------ app/service/planning_svc.py | 4 +- 2 files changed, 7 insertions(+), 64 deletions(-) diff --git a/app/api/v2/managers/operation_api_manager.py b/app/api/v2/managers/operation_api_manager.py index e7409a62f..8ad547b51 100644 --- a/app/api/v2/managers/operation_api_manager.py +++ b/app/api/v2/managers/operation_api_manager.py @@ -21,7 +21,6 @@ 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) @@ -88,8 +87,6 @@ async def update_operation_link(self, operation_id: str, link_id: str, link_data if not link.is_valid_status(link_status): raise JsonHttpBadRequest(f'Cannot update link {link_id} due to invalid link status.') link.status = link_status - if link.can_ignore(): - operation.add_ignored_link(link.id) return link.display async def create_potential_link(self, operation_id: str, data: dict, access: BaseWorld.Access): @@ -98,10 +95,13 @@ async def create_potential_link(self, operation_id: str, data: dict, access: Bas agent = await self.get_agent(operation, data) if data['executor']['name'] not in agent.executors: raise JsonHttpBadRequest(f'Agent {agent.paw} missing specified executor') - encoded_command = self._encode_string(agent.replace(self._encode_string(data['executor']['command']), - file_svc=self.services['file_svc'])) + executor = self.build_executor(data=data.pop('executor', {}), agent=agent) ability = self.build_ability(data=data.pop('ability', {}), executor=executor) + for hook, fcall in executor.HOOKS.items(): + await fcall(ability, executor) + encoded_command = self._encode_string(agent.replace(self._encode_string(executor.command), + file_svc=self.services['file_svc'])) link = Link.load(dict(command=encoded_command, plaintext_command=encoded_command, paw=agent.paw, ability=ability, executor=executor, status=operation.link_status(), score=data.get('score', 0), jitter=data.get('jitter', 0), cleanup=data.get('cleanup', 0), pin=data.get('pin', 0), @@ -207,63 +207,6 @@ async def get_agent(self, operation: Operation, data: dict): raise JsonHttpNotFound(f'Agent {data["paw"]} was not found.') return agent - def get_agents(self, operation: dict): - agents = {} - chain = operation.get('chain', []) - for link in chain: - paw = link.get('paw') - if paw and paw not in agents: - tmp_agent = self.find_object('agents', {'paw': paw}).display - tmp_agent['links'] = [] - agents[paw] = tmp_agent - agents[paw]['links'].append(link) - return agents - - async def get_hosts(self, operation: dict): - hosts = {} - chain = operation.get('chain', []) - for link in chain: - host = link.get('host') - if not host: - continue - if host not in hosts: - tmp_agent = self.find_object('agents', {'host': host}).display - tmp_host = { - 'host': tmp_agent.get('host'), - 'host_ip_addrs': tmp_agent.get('host_ip_addrs'), - 'platform': tmp_agent.get('platform'), - 'reachable_hosts': await self.get_reachable_hosts(agent=tmp_agent) - } - hosts[host] = tmp_host - return hosts - - async def get_reachable_hosts(self, agent: dict = None, operation: dict = None): - """ - NOTE: When agent is supplied, only hosts discovered by agent - are retrieved. - """ - trait_names = BaseWorld.get_config('reachable_host_traits') or [] - paws = () - - if agent is not None: - paws = paws + (agent.get('paw'),) - else: - for agent in operation.get('host_group', []): - paw = agent.get('paw') - if paw: - paws = paws + (paw,) - - hosts = [] - for trait in trait_names: - fqdns = await self.services['knowledge_svc'].get_facts({ - 'trait': trait, - 'collected_by': paws, - }) - for name in fqdns: - hosts.append(name.value) - - return hosts - def build_executor(self, data: dict, agent: Agent): if not data.get('timeout'): data['timeout'] = 60 diff --git a/app/service/planning_svc.py b/app/service/planning_svc.py index 335da59d9..5a71cfc07 100644 --- a/app/service/planning_svc.py +++ b/app/service/planning_svc.py @@ -352,8 +352,8 @@ async def _generate_new_links(self, operation, agent, abilities, link_status): if not executor: continue - if executor.HOOKS and executor.language and executor.language in executor.HOOKS: - await executor.HOOKS[executor.language](ability, executor) + for hook, fcall in executor.HOOKS.items(): + await fcall(ability, executor) if executor.command: link = Link.load(dict(command=self.encode_string(executor.test), paw=agent.paw, score=0, ability=ability, executor=executor, status=link_status, From 76b17e74c5f1105fda47a7d0614ccb5a4b767dcd Mon Sep 17 00:00:00 2001 From: bjeffries Date: Tue, 3 Dec 2024 13:44:02 -0500 Subject: [PATCH 2/7] was using old operation_api_manager.py, now leveraging latest --- app/api/v2/managers/operation_api_manager.py | 61 +++++++++++++++++++- 1 file changed, 60 insertions(+), 1 deletion(-) diff --git a/app/api/v2/managers/operation_api_manager.py b/app/api/v2/managers/operation_api_manager.py index 8ad547b51..73d8beda9 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) @@ -87,6 +88,8 @@ async def update_operation_link(self, operation_id: str, link_id: str, link_data if not link.is_valid_status(link_status): raise JsonHttpBadRequest(f'Cannot update link {link_id} due to invalid link status.') link.status = link_status + if link.can_ignore(): + operation.add_ignored_link(link.id) return link.display async def create_potential_link(self, operation_id: str, data: dict, access: BaseWorld.Access): @@ -95,7 +98,6 @@ async def create_potential_link(self, operation_id: str, data: dict, access: Bas agent = await self.get_agent(operation, data) if data['executor']['name'] not in agent.executors: raise JsonHttpBadRequest(f'Agent {agent.paw} missing specified executor') - executor = self.build_executor(data=data.pop('executor', {}), agent=agent) ability = self.build_ability(data=data.pop('ability', {}), executor=executor) for hook, fcall in executor.HOOKS.items(): @@ -207,6 +209,63 @@ async def get_agent(self, operation: Operation, data: dict): raise JsonHttpNotFound(f'Agent {data["paw"]} was not found.') return agent + def get_agents(self, operation: dict): + agents = {} + chain = operation.get('chain', []) + for link in chain: + paw = link.get('paw') + if paw and paw not in agents: + tmp_agent = self.find_object('agents', {'paw': paw}).display + tmp_agent['links'] = [] + agents[paw] = tmp_agent + agents[paw]['links'].append(link) + return agents + + async def get_hosts(self, operation: dict): + hosts = {} + chain = operation.get('chain', []) + for link in chain: + host = link.get('host') + if not host: + continue + if host not in hosts: + tmp_agent = self.find_object('agents', {'host': host}).display + tmp_host = { + 'host': tmp_agent.get('host'), + 'host_ip_addrs': tmp_agent.get('host_ip_addrs'), + 'platform': tmp_agent.get('platform'), + 'reachable_hosts': await self.get_reachable_hosts(agent=tmp_agent) + } + hosts[host] = tmp_host + return hosts + + async def get_reachable_hosts(self, agent: dict = None, operation: dict = None): + """ + NOTE: When agent is supplied, only hosts discovered by agent + are retrieved. + """ + trait_names = BaseWorld.get_config('reachable_host_traits') or [] + paws = () + + if agent is not None: + paws = paws + (agent.get('paw'),) + else: + for agent in operation.get('host_group', []): + paw = agent.get('paw') + if paw: + paws = paws + (paw,) + + hosts = [] + for trait in trait_names: + fqdns = await self.services['knowledge_svc'].get_facts({ + 'trait': trait, + 'collected_by': paws, + }) + for name in fqdns: + hosts.append(name.value) + + return hosts + def build_executor(self, data: dict, agent: Agent): if not data.get('timeout'): data['timeout'] = 60 From a06c97ea8b33889d8bf95deedca6049e84825ab2 Mon Sep 17 00:00:00 2001 From: bjeffries Date: Tue, 3 Dec 2024 13:47:00 -0500 Subject: [PATCH 3/7] fixed encoded command ref --- app/api/v2/managers/operation_api_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app/api/v2/managers/operation_api_manager.py b/app/api/v2/managers/operation_api_manager.py index 73d8beda9..e08ee9a8d 100644 --- a/app/api/v2/managers/operation_api_manager.py +++ b/app/api/v2/managers/operation_api_manager.py @@ -102,7 +102,7 @@ async def create_potential_link(self, operation_id: str, data: dict, access: Bas ability = self.build_ability(data=data.pop('ability', {}), executor=executor) for hook, fcall in executor.HOOKS.items(): await fcall(ability, executor) - encoded_command = self._encode_string(agent.replace(self._encode_string(executor.command), + encoded_command = self._encode_string(agent.replace(self._encode_string(data['executor']['command']), file_svc=self.services['file_svc'])) link = Link.load(dict(command=encoded_command, plaintext_command=encoded_command, paw=agent.paw, ability=ability, executor=executor, status=operation.link_status(), score=data.get('score', 0), jitter=data.get('jitter', 0), From 965408859db9f29415432da156a76def625d93b7 Mon Sep 17 00:00:00 2001 From: Blaine Jeffries Date: Wed, 18 Dec 2024 20:16:11 +0000 Subject: [PATCH 4/7] Added inline comment to hook purpose and fixed style issue with unused dictionary key. --- app/api/v2/managers/operation_api_manager.py | 5 +++-- app/service/planning_svc.py | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/app/api/v2/managers/operation_api_manager.py b/app/api/v2/managers/operation_api_manager.py index e08ee9a8d..f30c4b85e 100644 --- a/app/api/v2/managers/operation_api_manager.py +++ b/app/api/v2/managers/operation_api_manager.py @@ -100,10 +100,11 @@ async def create_potential_link(self, operation_id: str, data: dict, access: Bas raise JsonHttpBadRequest(f'Agent {agent.paw} missing specified executor') executor = self.build_executor(data=data.pop('executor', {}), agent=agent) ability = self.build_ability(data=data.pop('ability', {}), executor=executor) - for hook, fcall in executor.HOOKS.items(): + """This dictionary provides plugins a way to hook into abilities at runtime""" + for _hook, fcall in executor.HOOKS.items(): await fcall(ability, executor) encoded_command = self._encode_string(agent.replace(self._encode_string(data['executor']['command']), - file_svc=self.services['file_svc'])) + file_svc=self.services['file_svc'])) link = Link.load(dict(command=encoded_command, plaintext_command=encoded_command, paw=agent.paw, ability=ability, executor=executor, status=operation.link_status(), score=data.get('score', 0), jitter=data.get('jitter', 0), cleanup=data.get('cleanup', 0), pin=data.get('pin', 0), diff --git a/app/service/planning_svc.py b/app/service/planning_svc.py index 5a71cfc07..755ff2da5 100644 --- a/app/service/planning_svc.py +++ b/app/service/planning_svc.py @@ -351,8 +351,8 @@ async def _generate_new_links(self, operation, agent, abilities, link_status): executor = await agent.get_preferred_executor(ability) if not executor: continue - - for hook, fcall in executor.HOOKS.items(): + """This dictionary provides plugins a way to hook into abilities at runtime""" + for _hook, fcall in executor.HOOKS.items(): await fcall(ability, executor) if executor.command: link = Link.load(dict(command=self.encode_string(executor.test), paw=agent.paw, score=0, From 858164c5b667e209f703e9eff08eed2a7c04197b Mon Sep 17 00:00:00 2001 From: elegantmoose Date: Wed, 18 Dec 2024 19:19:38 -0500 Subject: [PATCH 5/7] pulling out function --- app/api/v2/managers/operation_api_manager.py | 9 ++++++--- app/service/planning_svc.py | 9 ++++++--- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/app/api/v2/managers/operation_api_manager.py b/app/api/v2/managers/operation_api_manager.py index f30c4b85e..57d8fdae6 100644 --- a/app/api/v2/managers/operation_api_manager.py +++ b/app/api/v2/managers/operation_api_manager.py @@ -100,9 +100,7 @@ async def create_potential_link(self, operation_id: str, data: dict, access: Bas raise JsonHttpBadRequest(f'Agent {agent.paw} missing specified executor') executor = self.build_executor(data=data.pop('executor', {}), agent=agent) ability = self.build_ability(data=data.pop('ability', {}), executor=executor) - """This dictionary provides plugins a way to hook into abilities at runtime""" - for _hook, fcall in executor.HOOKS.items(): - await fcall(ability, executor) + self._call_ability_plugin_hooks(ability, executor) encoded_command = self._encode_string(agent.replace(self._encode_string(data['executor']['command']), file_svc=self.services['file_svc'])) link = Link.load(dict(command=encoded_command, plaintext_command=encoded_command, paw=agent.paw, ability=ability, executor=executor, @@ -173,6 +171,11 @@ async def _construct_and_dump_source(self, source_id: str): if not source: source = (await self.services['data_svc'].locate('sources', match=dict(name='basic'))) return SourceSchema().dump(source[0]) + + async def _call_ability_plugin_hooks(self, ability, executor): + """Calls any plugin hooks (at runtime) that exist for the ability and executor.""" + for _hook, fcall in executor.HOOKS.items(): + await fcall(ability, executor) async def validate_operation_state(self, data: dict, existing: Operation = None): if not existing: diff --git a/app/service/planning_svc.py b/app/service/planning_svc.py index 755ff2da5..98a442e53 100644 --- a/app/service/planning_svc.py +++ b/app/service/planning_svc.py @@ -351,9 +351,7 @@ async def _generate_new_links(self, operation, agent, abilities, link_status): executor = await agent.get_preferred_executor(ability) if not executor: continue - """This dictionary provides plugins a way to hook into abilities at runtime""" - for _hook, fcall in executor.HOOKS.items(): - await fcall(ability, executor) + self._call_ability_plugin_hooks(ability, executor) if executor.command: link = Link.load(dict(command=self.encode_string(executor.test), paw=agent.paw, score=0, ability=ability, executor=executor, status=link_status, @@ -391,6 +389,11 @@ async def _generate_cleanup_links(self, operation, agent, link_status): lnk.apply_id(agent.host) links.append(lnk) return links + + async def _call_ability_plugin_hooks(self, ability, executor): + """Calls any plugin hooks (at runtime) that exist for the ability and executor.""" + for _hook, fcall in executor.HOOKS.items(): + await fcall(ability, executor) @staticmethod async def _apply_adjustments(operation, links): From 4bca515e8ee5683b2b0c32d95cea620658ff9d6c Mon Sep 17 00:00:00 2001 From: elegantmoose Date: Wed, 18 Dec 2024 19:38:39 -0500 Subject: [PATCH 6/7] flake 8 --- app/api/v2/managers/operation_api_manager.py | 2 +- app/service/planning_svc.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/app/api/v2/managers/operation_api_manager.py b/app/api/v2/managers/operation_api_manager.py index 57d8fdae6..5ed6adf9b 100644 --- a/app/api/v2/managers/operation_api_manager.py +++ b/app/api/v2/managers/operation_api_manager.py @@ -171,7 +171,7 @@ async def _construct_and_dump_source(self, source_id: str): if not source: source = (await self.services['data_svc'].locate('sources', match=dict(name='basic'))) return SourceSchema().dump(source[0]) - + async def _call_ability_plugin_hooks(self, ability, executor): """Calls any plugin hooks (at runtime) that exist for the ability and executor.""" for _hook, fcall in executor.HOOKS.items(): diff --git a/app/service/planning_svc.py b/app/service/planning_svc.py index 98a442e53..af3bf17c7 100644 --- a/app/service/planning_svc.py +++ b/app/service/planning_svc.py @@ -389,7 +389,7 @@ async def _generate_cleanup_links(self, operation, agent, link_status): lnk.apply_id(agent.host) links.append(lnk) return links - + async def _call_ability_plugin_hooks(self, ability, executor): """Calls any plugin hooks (at runtime) that exist for the ability and executor.""" for _hook, fcall in executor.HOOKS.items(): From 755f8f0c9390c2a0e74e24802babb3a68c417b14 Mon Sep 17 00:00:00 2001 From: elegantmoose Date: Wed, 18 Dec 2024 19:47:52 -0500 Subject: [PATCH 7/7] missing await --- app/api/v2/managers/operation_api_manager.py | 2 +- app/service/planning_svc.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/app/api/v2/managers/operation_api_manager.py b/app/api/v2/managers/operation_api_manager.py index 5ed6adf9b..0f3988581 100644 --- a/app/api/v2/managers/operation_api_manager.py +++ b/app/api/v2/managers/operation_api_manager.py @@ -100,7 +100,7 @@ async def create_potential_link(self, operation_id: str, data: dict, access: Bas raise JsonHttpBadRequest(f'Agent {agent.paw} missing specified executor') executor = self.build_executor(data=data.pop('executor', {}), agent=agent) ability = self.build_ability(data=data.pop('ability', {}), executor=executor) - self._call_ability_plugin_hooks(ability, executor) + await self._call_ability_plugin_hooks(ability, executor) encoded_command = self._encode_string(agent.replace(self._encode_string(data['executor']['command']), file_svc=self.services['file_svc'])) link = Link.load(dict(command=encoded_command, plaintext_command=encoded_command, paw=agent.paw, ability=ability, executor=executor, diff --git a/app/service/planning_svc.py b/app/service/planning_svc.py index af3bf17c7..da13e6d16 100644 --- a/app/service/planning_svc.py +++ b/app/service/planning_svc.py @@ -351,7 +351,7 @@ async def _generate_new_links(self, operation, agent, abilities, link_status): executor = await agent.get_preferred_executor(ability) if not executor: continue - self._call_ability_plugin_hooks(ability, executor) + await self._call_ability_plugin_hooks(ability, executor) if executor.command: link = Link.load(dict(command=self.encode_string(executor.test), paw=agent.paw, score=0, ability=ability, executor=executor, status=link_status,