From 7085eaba27046d39ed4912992c7e42e581ec3a05 Mon Sep 17 00:00:00 2001 From: Daan Date: Wed, 27 Nov 2024 22:10:27 +0100 Subject: [PATCH] Added node translation from data in global_list --- meshbook/examples/apt_upgrade.yaml | 16 ++++++ meshbook/meshbook.py | 83 +++++++++++++++++++++--------- 2 files changed, 74 insertions(+), 25 deletions(-) create mode 100644 meshbook/examples/apt_upgrade.yaml diff --git a/meshbook/examples/apt_upgrade.yaml b/meshbook/examples/apt_upgrade.yaml new file mode 100644 index 0000000..f77ea5a --- /dev/null +++ b/meshbook/examples/apt_upgrade.yaml @@ -0,0 +1,16 @@ +--- +name: Refresh the apt cache +company: Temp-Agents +#device: MSI +tasks: + - name: refresh the cache + command: "apt update" + + - name: display available upgrades + command: "apt list --upgradable" + + - name: apply upgrades + command: "apt upgrade -y" + + - name: cleanup remaining packages + command: "apt autoremove -y" \ No newline at end of file diff --git a/meshbook/meshbook.py b/meshbook/meshbook.py index 06b2024..24c827b 100644 --- a/meshbook/meshbook.py +++ b/meshbook/meshbook.py @@ -16,14 +16,13 @@ expected_responses = 0 basic_ready_state = asyncio.Event() ready_for_next = asyncio.Event() global_list = [] - +responses_dict = {} class ScriptEndTrigger(Exception): """Custom Exception to handle script termination events.""" pass - -class MeshcallerUtilities: +class MeshbookUtilities: """Helper utility functions for the Meshcaller application.""" @staticmethod @@ -31,12 +30,6 @@ class MeshcallerUtilities: """Encode a string in Base64 format.""" return b64encode(string.encode('utf-8')).decode() - @staticmethod - def read_yaml(file_path: str) -> dict: - """Read a YAML file and return its content as a dictionary.""" - with open(file_path, 'r') as file: - return yaml.safe_load(file) - @staticmethod def get_target_ids(company: str = None, device: str = None) -> list: """Retrieve target IDs based on company or device.""" @@ -74,8 +67,35 @@ class MeshcallerUtilities: return my_config[segment] + @staticmethod + def read_yaml(file_path: str) -> dict: + """Read a YAML file and return its content as a dictionary.""" + with open(file_path, 'r') as file: + return yaml.safe_load(file) -class MeshcallerWebSocket: + @staticmethod + def translate_nodeids(batches_dict, global_list) -> dict: + for batch_name, items in batches_dict.items(): + + for item in items: # Process each item in the batch + node_id = item["nodeid"] # Get the nodeid field + + real_name = None + for company in global_list: + for machine in company["nodes"]: + if machine["node_id"] == node_id: + real_name = machine["node_name"] + break + if real_name: + break + + # If found, replace the nodeid with the real node name, otherwise mark it as "Unknown Node" + if real_name: + item["nodeid"] = real_name + + return batches_dict + +class MeshbookWebsocket: """Handles WebSocket connections and interactions.""" def __init__(self): @@ -116,7 +136,7 @@ class MeshcallerWebSocket: async def ws_handler(self, uri: str, username: str, password: str): """Main WebSocket connection handler.""" - login_string = f'{MeshcallerUtilities.base64_encode(username)},{MeshcallerUtilities.base64_encode(password)}' + login_string = f'{MeshbookUtilities.base64_encode(username)},{MeshbookUtilities.base64_encode(password)}' ws_headers = { 'User-Agent': 'MeshCentral API client', 'x-meshauth': login_string @@ -143,7 +163,7 @@ class MeshcallerWebSocket: print(f"An error occurred: {e}") -class MeshcallerProcessor: +class MeshbookProcessor: """Processes data received from the WebSocket.""" def __init__(self): @@ -196,23 +216,29 @@ class MeshcallerProcessor: basic_ready_state.set() ready_for_next.set() - async def receive_processor(self, python_client: MeshcallerWebSocket): + async def receive_processor(self, python_client: MeshbookWebsocket): """Processes messages received from the WebSocket.""" global response_counter + temp_responses_list = [] while True: message = await python_client.received_response_queue.get() action_type = message.get('action') if action_type in ('meshes', 'nodes'): self.handle_basic_data(message[action_type]) elif action_type == 'msg': - print(json.dumps(message, indent=4)) + temp_responses_list.append(message) + response_counter += 1 # Increment response counter - if not args.silent: + if not args.silent or args.information: print("Current Batch: {}".format(math.ceil(response_counter/len(target_ids)))) + print("Current response number: {}".format(response_counter)) print("Current Calculation: {} % {} = {}".format(response_counter, len(target_ids), response_counter % len(target_ids))) if response_counter % len(target_ids) == 0: + batch_name = "Batch {}".format(math.ceil(response_counter / len(target_ids))) + responses_dict[batch_name] = temp_responses_list + temp_responses_list = [] ready_for_next.set() elif action_type == 'close': print(message) @@ -224,14 +250,14 @@ class MeshcallerActions: """Processes playbook actions.""" @staticmethod - async def process_arguments(python_client: MeshcallerWebSocket, playbook_path: str): + async def process_arguments(python_client: MeshbookWebsocket, playbook_path: str): """Executes tasks defined in the playbook.""" global response_counter, expected_responses, target_ids await basic_ready_state.wait() # Wait for the basic data to be ready - playbook_yaml = MeshcallerUtilities.read_yaml(playbook_path) - target_ids = MeshcallerUtilities.get_target_ids( + playbook_yaml = MeshbookUtilities.read_yaml(playbook_path) + target_ids = MeshbookUtilities.get_target_ids( company=playbook_yaml.get('company'), device=playbook_yaml.get('device') ) @@ -248,16 +274,18 @@ class MeshcallerActions: 'reply': True } - # Calculate the total expected responses: tasks x target nodes - expected_responses = len(playbook_yaml['tasks']) * len(target_ids) + expected_responses = len(playbook_yaml['tasks']) * len(target_ids) # Calculate the total expected responses: tasks x target nodes # Send commands for all nodes at once for task in playbook_yaml['tasks']: await ready_for_next.wait() run_command_template["cmds"] = task['command'] run_command_template["nodeids"] = target_ids # Send to all target IDs at once - print("Running task:", task) - print("-=-" * 40) + + if not args.silent or args.information: + print("-=-" * 40) + print("Running task:", task) + print("-=-" * 40) # Send the command to all nodes in one go await python_client.ws_send_data(json.dumps(run_command_template)) @@ -268,6 +296,10 @@ class MeshcallerActions: await asyncio.sleep(1) # Exit gracefully + print("-=-" * 40) + updated_response_dict = MeshbookUtilities.translate_nodeids(responses_dict, global_list) + + print(json.dumps(updated_response_dict,indent=4)) raise ScriptEndTrigger("All tasks completed successfully: Expected {} Received {}".format(expected_responses, response_counter)) @@ -276,14 +308,15 @@ async def main(): parser.add_argument("--conf", type=str, help="Path for the API configuration file (default: ./api.conf).") parser.add_argument("-pb", "--playbook", type=str, help="Path to the playbook file.", required=True) parser.add_argument("-s", "--silent", action="store_true", help="Suppress terminal output.") + parser.add_argument("-i", "--information", action="store_true", help="Output the calculations and other informational output.") global args args = parser.parse_args() try: - credentials = MeshcallerUtilities.load_config(args.conf) - python_client = MeshcallerWebSocket() - processor = MeshcallerProcessor() + credentials = MeshbookUtilities.load_config(args.conf) + python_client = MeshbookWebsocket() + processor = MeshbookProcessor() websocket_task = asyncio.create_task(python_client.ws_handler( credentials['websocket_url'],