New control_dna component

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)

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)

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