reorg

Author: willcl-ark

src/warnet/cli/network.py

@@ -19,19 +19,7 @@ def network():
     """Network commands"""
 
 
-def set_kubectl_context(namespace: str):
-    """
-    Set the default kubectl context to the specified namespace.
-    """
-    command = f"kubectl config set-context --current --namespace={namespace}"
-    result = run_command(command, stream_output=True)
-    if result:
-        print(f"Kubectl context set to namespace: {namespace}")
-    else:
-        print(f"Failed to set kubectl context to namespace: {namespace}")
-    return result
-
-
+# High-level network operations
 @network.command()
 @click.argument("graph_file", default=DEFAULT_GRAPH_FILE, type=click.Path())
 @click.option("--network", default="warnet", show_default=True)
@@ -40,53 +28,27 @@ def start(graph_file: Path, logging: bool, network: str):
     """
     Start a warnet with topology loaded from a <graph_file> into [network]
     """
-    # Generate the Kubernetes YAML
     graph = read_graph_file(graph_file)
     kubernetes_yaml = generate_kubernetes_yaml(graph)
 
-    # Write the YAML to a temporary file
     with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as temp_file:
         yaml.dump_all(kubernetes_yaml, temp_file)
         temp_file_path = temp_file.name
 
     try:
-        # Deploy base configurations
-        base_configs = [
-            "namespace.yaml",
-            "rbac-config.yaml",
-        ]
-
-        for config in base_configs:
-            command = f"kubectl apply -f {WAR_MANIFESTS}/{config}"
-            result = run_command(command, stream_output=True)
-            if not result:
-                print(f"Failed to apply {config}")
-                return
-
-        # Apply the YAML using kubectl
-        command = f"kubectl apply -f {temp_file_path}"
-        result = run_command(command, stream_output=True)
-
-        if result:
+        if deploy_base_configurations() and apply_kubernetes_yaml(temp_file_path):
             print(f"Warnet '{network}' started successfully.")
-
-            # Set kubectl context to the warnet namespace
-            context_result = set_kubectl_context(network)
-            if not context_result:
+            if not set_kubectl_context(network):
                 print(
                     "Warning: Failed to set kubectl context. You may need to manually switch to the warnet namespace."
                 )
-
-            if logging:
-                helm_result = setup_logging_helm()
-                if helm_result:
-                    print("Helm charts installed successfully.")
-                else:
-                    print("Failed to install Helm charts.")
+            if logging and setup_logging_helm():
+                print("Helm charts installed successfully.")
+            else:
+                print("Failed to install Helm charts.")
         else:
             print(f"Failed to start warnet '{network}'.")
     finally:
-        # Clean up the temporary file
         Path(temp_file_path).unlink()
 
 
@@ -96,30 +58,20 @@ def down(network: str):
     """
     Bring down a running warnet named [network]
     """
-    # Delete the namespace
-    command = f"kubectl delete namespace {network}"
-    result = run_command(command, stream_output=True)
-    # TODO: Fix this
-    command = "kubectl delete namespace warnet-logging"
-    result = run_command(command, stream_output=True)
-
-    if result:
-        print(f"Warnet '{network}' has been successfully brought down and the namespace deleted.")
+    if delete_namespace(network) and delete_namespace("warnet-logging"):
+        print(f"Warnet '{network}' has been successfully brought down and the namespaces deleted.")
     else:
-        print(f"Failed to bring down warnet '{network}' or delete the namespace.")
+        print(f"Failed to bring down warnet '{network}' or delete the namespaces.")
 
 
 @network.command()
 @click.option("--follow", "-f", is_flag=True, help="Follow logs")
 def logs(follow: bool):
     """Get Kubernetes logs from the RPC server"""
     command = "kubectl logs rpc-0"
-    stream_output = False
     if follow:
         command += " --follow"
-        stream_output = True
-
-    run_command(command, stream_output=stream_output)
+    run_command(command, stream_output=follow)
 
 
 @network.command()
@@ -129,44 +81,52 @@ def generate_yaml(graph_file: Path, output: str):
     """
     Generate a Kubernetes YAML file from a graph file for deploying warnet nodes.
     """
-    # Read and parse the graph file
     graph = read_graph_file(graph_file)
-
-    # Generate the Kubernetes YAML
     kubernetes_yaml = generate_kubernetes_yaml(graph)
 
-    # Write the YAML to a file
     with open(output, "w") as f:
         yaml.dump_all(kubernetes_yaml, f)
 
     print(f"Kubernetes YAML file generated: {output}")
 
 
-def read_graph_file(graph_file: Path) -> nx.Graph:
-    with open(graph_file) as f:
-        return nx.parse_graphml(f.read())
+@network.command()
+@click.argument("graph_file", default=DEFAULT_GRAPH_FILE, type=click.Path(exists=True))
+def connect(graph_file: Path):
+    """
+    Connect nodes based on the edges defined in the graph file.
+    """
+    tree = ET.parse(graph_file)
+    root = tree.getroot()
+    edges = root.findall(".//{http://graphml.graphdrawing.org/xmlns}edge")
 
+    for edge in edges:
+        source = edge.get("source")
+        target = edge.get("target")
+        command = f"kubectl exec -it warnet-node-{source} -- bitcoin-cli -rpcuser=user -rpcpassword=password addnode warnet-node-{target}-service:8333 add"
 
-def generate_kubernetes_yaml(graph: nx.Graph) -> list:
-    kubernetes_objects = []
+        print(f"Connecting node {source} to node {target}")
+        if run_command(command, stream_output=True):
+            print(f"Successfully connected node {source} to node {target}")
+        else:
+            print(f"Failed to connect node {source} to node {target}")
 
-    # Add Namespace object
-    namespace = create_namespace()
-    kubernetes_objects.append(namespace)
+    print("All connections attempted.")
 
-    for node, data in graph.nodes(data=True):
-        # Create a ConfigMap for each node
-        config = generate_node_config(node, data)
-        config_map = create_config_map(node, config)
-        kubernetes_objects.append(config_map)
 
-        # Create a deployment for each node
-        deployment = create_node_deployment(node, data)
-        kubernetes_objects.append(deployment)
+# Kubernetes object generation
+def generate_kubernetes_yaml(graph: nx.Graph) -> list:
+    kubernetes_objects = [create_namespace()]
 
-        # Create a service for each node
-        service = create_node_service(node)
-        kubernetes_objects.append(service)
+    for node, data in graph.nodes(data=True):
+        config = generate_node_config(node, data)
+        kubernetes_objects.extend(
+            [
+                create_config_map(node, config),
+                create_node_deployment(node, data),
+                create_node_service(node),
+            ]
+        )
 
     return kubernetes_objects
 
@@ -221,57 +181,22 @@ def create_node_service(node: int) -> dict:
     }
 
 
-def setup_logging_helm():
-    """
-    Run the required Helm commands for setting up Grafana, Prometheus, and Loki.
-    """
-    helm_commands = [
-        "helm repo add grafana https://grafana.github.io/helm-charts",
-        "helm repo add prometheus-community https://prometheus-community.github.io/helm-charts",
-        "helm repo update",
-        f"helm upgrade --install --namespace warnet-logging --create-namespace --values {WAR_MANIFESTS}/loki_values.yaml loki grafana/loki --version 5.47.2",
-        "helm upgrade --install --namespace warnet-logging promtail grafana/promtail",
-        "helm upgrade --install --namespace warnet-logging prometheus prometheus-community/kube-prometheus-stack --namespace warnet-logging --set grafana.enabled=false",
-        f"helm upgrade --install --namespace warnet-logging loki-grafana grafana/grafana --values {WAR_MANIFESTS}/grafana_values.yaml",
-    ]
-
-    for command in helm_commands:
-        result = run_command(command, stream_output=True)
-        if not result:
-            print(f"Failed to run Helm command: {command}")
-            return False
-    return True
-
-
-@network.command()
-@click.argument("graph_file", default=DEFAULT_GRAPH_FILE, type=click.Path(exists=True))
-def connect(graph_file: Path):
-    """
-    Connect nodes based on the edges defined in the graph file.
-    """
-    # Parse the GraphML file
-    tree = ET.parse(graph_file)
-    root = tree.getroot()
-
-    # Find all edge elements
-    edges = root.findall(".//{http://graphml.graphdrawing.org/xmlns}edge")
-
-    for edge in edges:
-        source = edge.get("source")
-        target = edge.get("target")
-
-        # Construct the kubectl command
-        command = f"kubectl exec -it warnet-node-{source} -- bitcoin-cli -rpcuser=user -rpcpassword=password addnode warnet-node-{target}-service:8333 add"
+def create_config_map(node: int, config: str) -> dict:
+    return {
+        "apiVersion": "v1",
+        "kind": "ConfigMap",
+        "metadata": {
+            "name": f"bitcoin-config-node-{node}",
+            "namespace": "warnet",
+        },
+        "data": {"bitcoin.conf": config},
+    }
 
-        print(f"Connecting node {source} to node {target}")
-        result = run_command(command, stream_output=True)
 
-        if result:
-            print(f"Successfully connected node {source} to node {target}")
-        else:
-            print(f"Failed to connect node {source} to node {target}")
-
-    print("All connections attempted.")
+# Utility functions
+def read_graph_file(graph_file: Path) -> nx.Graph:
+    with open(graph_file) as f:
+        return nx.parse_graphml(f.read())
 
 
 def generate_node_config(node: int, data: dict) -> str:
@@ -300,13 +225,59 @@ def generate_node_config(node: int, data: dict) -> str:
     return f"{base_config}\n{node_specific_config}"
 
 
-def create_config_map(node: int, config: str) -> dict:
-    return {
-        "apiVersion": "v1",
-        "kind": "ConfigMap",
-        "metadata": {
-            "name": f"bitcoin-config-node-{node}",
-            "namespace": "warnet",
-        },
-        "data": {"bitcoin.conf": config},
-    }
+def set_kubectl_context(namespace: str):
+    """
+    Set the default kubectl context to the specified namespace.
+    """
+    command = f"kubectl config set-context --current --namespace={namespace}"
+    result = run_command(command, stream_output=True)
+    if result:
+        print(f"Kubectl context set to namespace: {namespace}")
+    else:
+        print(f"Failed to set kubectl context to namespace: {namespace}")
+    return result
+
+
+def deploy_base_configurations():
+    base_configs = [
+        "namespace.yaml",
+        "rbac-config.yaml",
+    ]
+
+    for config in base_configs:
+        command = f"kubectl apply -f {WAR_MANIFESTS}/{config}"
+        if not run_command(command, stream_output=True):
+            print(f"Failed to apply {config}")
+            return False
+    return True
+
+
+def apply_kubernetes_yaml(yaml_file: str):
+    command = f"kubectl apply -f {yaml_file}"
+    return run_command(command, stream_output=True)
+
+
+def delete_namespace(namespace: str):
+    command = f"kubectl delete namespace {namespace}"
+    return run_command(command, stream_output=True)
+
+
+def setup_logging_helm():
+    """
+    Run the required Helm commands for setting up Grafana, Prometheus, and Loki.
+    """
+    helm_commands = [
+        "helm repo add grafana https://grafana.github.io/helm-charts",
+        "helm repo add prometheus-community https://prometheus-community.github.io/helm-charts",
+        "helm repo update",
+        f"helm upgrade --install --namespace warnet-logging --create-namespace --values {WAR_MANIFESTS}/loki_values.yaml loki grafana/loki --version 5.47.2",
+        "helm upgrade --install --namespace warnet-logging promtail grafana/promtail",
+        "helm upgrade --install --namespace warnet-logging prometheus prometheus-community/kube-prometheus-stack --namespace warnet-logging --set grafana.enabled=false",
+        f"helm upgrade --install --namespace warnet-logging loki-grafana grafana/grafana --values {WAR_MANIFESTS}/grafana_values.yaml",
+    ]
+
+    for command in helm_commands:
+        if not run_command(command, stream_output=True):
+            print(f"Failed to run Helm command: {command}")
+            return False
+    return True