From cb137831961bc56063e35130d0b3fd470ab81ce8 Mon Sep 17 00:00:00 2001 From: MosbyTheGreat Date: Mon, 30 Sep 2024 13:26:23 +0200 Subject: [PATCH] control_dna component allows deploying and reading dna files --- commands/control_dna.py | 112 ++++++++++++++++++++++++++++++++++++++++ nerve.py | 12 +++++ requirements.txt | 1 + 3 files changed, 125 insertions(+) create mode 100644 commands/control_dna.py diff --git a/commands/control_dna.py b/commands/control_dna.py new file mode 100644 index 0000000..aab1fbf --- /dev/null +++ b/commands/control_dna.py @@ -0,0 +1,112 @@ +from io import BytesIO +from pathlib import Path +from zipfile import ZipFile + +import yaml +from nerveapi import nodes, session +from nerveapi.utils import append_ending + +from .utils import eprint + + +def handle_deploy_dna(args) -> int: + # Read DNA file from input file to dict + filename = append_ending(args.file_name, ".yaml") + try: + dna_file_path = Path(filename) + dna_file_dict = yaml.load(dna_file_path.read_text(), Loader=yaml.FullLoader) + except yaml.YAMLError as e: + eprint(f"Could not read input file: {e}") + return 1 + except FileNotFoundError: + eprint(f"File {filename} not found.") + return 1 + except OSError: + eprint(f"Could not open file {filename}.") + return 1 + + # Get Node from node_name parameter + node_name = args.node_name + node = get_node_by_name_and_online(node_name) + + # Fail if no node is found + if not node: + return 1 + + # Create file from dict to send with request + dna_file_yaml = yaml.dump(data=dna_file_dict, default_flow_style=False) + files = {'file': ('file-from-pipeline.yaml', dna_file_yaml)} + + # restart_all_workloads flag + restart_workloads = args.restart_all_workloads.__str__().lower() + + # Deploy DNA file + response = session.make_request(f'/nerve/dna/{node["serial_number"]}/target?continueInCaseOfRestart=true&restartAllWorkloads={restart_workloads}', method='PUT', data=dna_file_yaml, files=files) + + if response.status_code == 202: + print('Successfully pushed DNA File.') + return 0 + else: + eprint(f'Error while pushing {response.status_code}: {response.text}') + print('See https://docs.nerve.cloud/developer_guide/ms-api/ for info on error codes.') + return 1 + + +def handle_get_dna(args) -> int: + filename = args.output + strip_hash = args.strip_hash + + # Get Node from node_name parameter + node_name = args.node_name + node = get_node_by_name_and_online(node_name) + + # Fail if no node is found + if not node: + return 1 + + # Get DNA file + response = session.make_request(f'/nerve/dna/{node["serial_number"]}/current') + if response.ok: + print("DNA file loaded.") + with ZipFile(BytesIO(response.content)) as myzip: + with myzip.open(myzip.namelist()[0]) as file: + try: + dna_dict = yaml.safe_load(file) + if strip_hash: + strip_hash_from_dict(dna_dict) + with open(filename, 'w') as outfile: + yaml.dump(dna_dict, outfile, default_flow_style=False) + print(f"Successfully written DNA file to {filename}.") + return 0 + except yaml.YAMLError as e: + eprint(f"Could not write output file: {e}") + return 1 + except OSError: + eprint(f"Could not open output file {filename}.") + return 1 + else: + eprint(f'Retrieving deployed DNA File failed: {response.text}') + return 1 + + +def get_node_by_name_and_online(filter_string): + node_list = nodes.create_node_list() + filtered_node_list = nodes.filter_node_list(name_filter=filter_string, node_list=node_list) + online_nodes = [node for node in filtered_node_list if node['connection_status'] == 'online'] + + if len(online_nodes) == 0: + eprint(f'No matching online node found.') + return None + elif len(online_nodes) > 1: + online_nodes_names = [device['name'] for device in online_nodes] + eprint(f'More than one matching online node found, please retry with more restrictive filter: {online_nodes_names}') + return None + else: + node = online_nodes[0] + print(f'Node {node["name"]} selected.') + return node + + +def strip_hash_from_dict(dna_dict): + for workload in dna_dict['workloads']: + workload.pop('hash', None) diff --git a/nerve.py b/nerve.py index 306ad70..f6667e3 100755 --- a/nerve.py +++ b/nerve.py @@ -35,6 +35,7 @@ handle_stop_workloads, handle_restart_workloads ) +from commands.control_dna import handle_deploy_dna, handle_get_dna import sys from argparse import RawTextHelpFormatter @@ -249,6 +250,17 @@ def main(): parser_restart_workloads.add_argument("-V", '--verbose', help='Show debug information.', action="store_true") parser_restart_workloads.set_defaults(func=handle_restart_workloads) + parser_deploy_dna = subparsers.add_parser("deploy_dna", help="Deploy DNA file to node.", description="Deploys DNA file to specified node.") + parser_deploy_dna.add_argument("-f", "--file_name", default="dna.yaml", help="Input filename. Ignore for 'dna.yaml'.") + parser_deploy_dna.add_argument("-n", "--node_name", required=True, help="Node to deploy on, supports regex.") + parser_deploy_dna.add_argument("-r", "--restart_all_workloads", default=False, help="Whether to restart all deployed workloads.") + parser_deploy_dna.set_defaults(func=handle_deploy_dna) + parser_get_dna = subparsers.add_parser('get_dna', help="Get DNA file from node.", description="Gets DNA file from specified node.") + parser_get_dna.add_argument("-o", "--output", default="dna.yaml", help="Output filename. Ignore for 'dna.yaml'.") + parser_get_dna.add_argument("-n", "--node_name", required=True, help="Node to get DNA file from, supports regex.") + parser_get_dna.add_argument("-s", "--strip_hash", default=False, help="Strip hashes from DNA file.") + parser_get_dna.set_defaults(func=handle_get_dna) + if len(sys.argv)==1: parser.print_help(sys.stderr) diff --git a/requirements.txt b/requirements.txt index 548dd0e..7e8ddba 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,4 @@ requests==2.31.0 requests-toolbelt==1.0.0 urllib3==2.2.1 validators==0.24.0 +pyyaml==6.0.1 \ No newline at end of file