From 9d2999476dd9d166f461c7a86ab5dd0b3fe72e9b Mon Sep 17 00:00:00 2001 From: Daan Selen Date: Thu, 25 Sep 2025 09:49:19 +0200 Subject: [PATCH] refac: added some type checking and fixed a duplicate bug (apparently) --- meshbook.py | 15 +++++++++++---- modules/console.py | 1 + modules/executor.py | 3 ++- modules/utilities.py | 23 +++++++++++++++++------ 4 files changed, 31 insertions(+), 11 deletions(-) diff --git a/meshbook.py b/meshbook.py index 4da5eaa..f7ec5d4 100644 --- a/meshbook.py +++ b/meshbook.py @@ -53,8 +53,11 @@ async def gather_targets(args: argparse.Namespace, offline_list = [] target_os = meshbook.get("target_os") - ignore_categorisation = meshbook.get("ignore_categorisation", False) target_tag = meshbook.get("target_tag") + ignore_categorisation = meshbook.get("ignore_categorisation", False) + + assert target_os is not None, "target_os must be provided" + assert target_tag is not None, "target_tag must be provided" async def add_processed_devices(processed): """Helper to update target and offline lists.""" @@ -119,7 +122,7 @@ async def gather_targets(args: argparse.Namespace, elif isinstance(pseudo_target, str) and pseudo_target.lower() == "all": for group in group_list.values(): await process_group_helper(group) - elif isintance(pseudo_target, str): + elif isinstance(pseudo_target, str): console.nice_print( args, console.text_color.yellow + "The 'groups' key is being used, but only one string is given. Did you mean 'group'?", @@ -303,7 +306,7 @@ async def main(): group_list = await transform.compile_group_list(session) compiled_device_list = await gather_targets(args, meshbook, group_list, os_categories) - if len(compiled_device_list["target_list"]) == 0: + if "target_list" not in compiled_device_list or len(compiled_device_list["target_list"]) == 0: console.nice_print(args, console.text_color.red + "No targets found or targets unreachable, quitting.", True) @@ -327,6 +330,10 @@ async def main(): case {"devices": candidate_target_name}: target_name = str(candidate_target_name) + case _: + target_name = "" + + console.nice_print(args, console.text_color.yellow + "Executing meshbook on the target(s): " + console.text_color.green + target_name + console.text_color.yellow + ".") @@ -350,7 +357,7 @@ async def main(): except OSError as message: console.nice_print(args, - console.text_color.red + message, True) + console.text_color.red + f'{message}', True) if __name__ == "__main__": asyncio.run(main()) diff --git a/modules/console.py b/modules/console.py index be8407c..659cba1 100644 --- a/modules/console.py +++ b/modules/console.py @@ -14,6 +14,7 @@ class console: italic = "\x1B[3m" reset = "\x1B[0m" + @staticmethod def nice_print(args: argparse.Namespace, message: str, final: bool=False): ''' Helper function for terminal output, with a couple variables for the silent flag. Also clears terminal color each time. diff --git a/modules/executor.py b/modules/executor.py index 06f4cc1..c964895 100644 --- a/modules/executor.py +++ b/modules/executor.py @@ -11,6 +11,7 @@ from modules.utilities import transform intertask_delay = 1 class executor: + @staticmethod async def execute_meshbook(args: argparse.Namespace, session: meshctrl.Session, compiled_device_list: dict, meshbook: dict, group_list: dict) -> None: ''' Actual function that handles meshbook execution, also responsible for formatting the resulting JSON. @@ -28,7 +29,7 @@ class executor: if "powershell" in meshbook and meshbook["powershell"]: response = await session.run_command(nodeids=targets, command=task["command"],powershell=True,ignore_output=False,timeout=900) else: - response = await session.run_command(nodeids=targets, command=task["command"],ignore_output=False,timeout=900) + response = await session.run_command(nodeids=targets, command=task["command"],powershell=False,ignore_output=False,timeout=900) task_batch = [] for device in response: diff --git a/modules/utilities.py b/modules/utilities.py index 8437d1c..9744cc0 100644 --- a/modules/utilities.py +++ b/modules/utilities.py @@ -10,6 +10,7 @@ Creation and compilation of the MeshCentral nodes list (list of all nodes availa ''' class utilities: + @staticmethod async def load_config(args: argparse.Namespace, segment: str = 'meshcentral-account') -> dict: ''' @@ -32,17 +33,20 @@ class utilities: print(f'Segment "{segment}" not found in config file {conf_file}.') os._exit(1) - return config[segment] + return dict(config[segment]) - async def compile_book(meshbook_file: dict) -> dict: + @staticmethod + async def compile_book(meshbook_file: str) -> dict: ''' Simple function that opens the file and replaces placeholders through the next function. After that just return it. ''' - meshbook = open(meshbook_file, 'r') + with open(meshbook_file, 'r') as f: + meshbook = f.read() meshbook = await transform.replace_placeholders(yaml.safe_load(meshbook)) return meshbook + @staticmethod def get_os_variants(target_category: str, os_map: dict) -> set: ''' @@ -65,17 +69,19 @@ class utilities: return set() + @staticmethod async def filter_targets(devices: list[dict], os_categories: dict, - target_os: str = None, + target_os: str = "", ignore_categorisation: bool = False, - target_tag: str = None) -> dict: + target_tag: str = "") -> dict: ''' Filters devices based on reachability and optional OS criteria, supporting nested OS categories. ''' valid_devices = [] offline_devices = [] + allowed_os = set() # Identify correct OS filtering scope for key in os_categories: @@ -109,6 +115,7 @@ class utilities: "offline_devices": offline_devices } + @staticmethod async def process_device(device: str, group_list: dict, os_categories: dict, @@ -143,6 +150,7 @@ class utilities: import shlex class transform: + @staticmethod def process_shell_response(shlex_enable: bool, meshbook_result: dict) -> dict: for task_name, task_data in meshbook_result.items(): if task_name == "Offline": # Failsafe @@ -164,6 +172,7 @@ class transform: node_responses["result"] = clean_output return meshbook_result + @staticmethod async def translate_nodeid_to_name(target_id: str, group_list: dict) -> str: ''' Simple function that looks up nodeid to the human-readable name if existent - otherwise return None. @@ -173,8 +182,9 @@ class transform: for device in group_list[group]: if device["device_id"] == target_id: return device["device_name"] - return None + return "" + @staticmethod async def replace_placeholders(meshbook: dict) -> dict: ''' Replace the placeholders in both name and command fields of the tasks. According to the variables defined in the variables list. @@ -208,6 +218,7 @@ class transform: return meshbook + @staticmethod async def compile_group_list(session: meshctrl.Session) -> dict: ''' Function that retrieves the devices from MeshCentral and compiles it into a efficient list.