Merge pull request #626 from mplsgrant/2024-10-fix-logs-container

Add container selection logic

Author: mplsgrant

src/warnet/bitcoin.py

@@ -5,11 +5,11 @@
 from io import BytesIO
 
 import click
-from urllib3.exceptions import MaxRetryError
-
 from test_framework.messages import ser_uint256
 from test_framework.p2p import MESSAGEMAP
+from urllib3.exceptions import MaxRetryError
 
+from .constants import BITCOINCORE_CONTAINER
 from .k8s import get_default_namespace, get_mission
 from .process import run_command
 
@@ -39,10 +39,11 @@ def _rpc(tank: str, method: str, params: str):
     # bitcoin-cli should be able to read bitcoin.conf inside the container
     # so no extra args like port, chain, username or password are needed
     namespace = get_default_namespace()
+
     if params:
-        cmd = f"kubectl -n {namespace} exec {tank} -- bitcoin-cli {method} {' '.join(map(str, params))}"
+        cmd = f"kubectl -n {namespace} exec {tank} --container {BITCOINCORE_CONTAINER} -- bitcoin-cli {method} {' '.join(map(str, params))}"
     else:
-        cmd = f"kubectl -n {namespace} exec {tank} -- bitcoin-cli {method}"
+        cmd = f"kubectl -n {namespace} exec {tank} --container {BITCOINCORE_CONTAINER} -- bitcoin-cli {method}"
     return run_command(cmd)
 
 

src/warnet/constants.py

@@ -16,6 +16,9 @@
 INGRESS_NAMESPACE = "ingress"
 HELM_COMMAND = "helm upgrade --install --create-namespace"
 
+BITCOINCORE_CONTAINER = "bitcoincore"
+COMMANDER_CONTAINER = "commander"
+
 # Directories and files for non-python assets, e.g., helm charts, example scenarios, default configs
 SRC_DIR = files("warnet")
 RESOURCES_DIR = files("resources")

src/warnet/control.py

@@ -16,11 +16,17 @@
 from rich.prompt import Confirm, Prompt
 from rich.table import Table
 
-from .constants import COMMANDER_CHART, LOGGING_NAMESPACE
+from .constants import (
+    BITCOINCORE_CONTAINER,
+    COMMANDER_CHART,
+    COMMANDER_CONTAINER,
+    LOGGING_NAMESPACE,
+)
 from .k8s import (
     delete_pod,
     get_default_namespace,
     get_mission,
+    get_pod,
     get_pods,
     pod_log,
     snapshot_bitcoin_datadir,
@@ -329,9 +335,28 @@ def _logs(pod_name: str, follow: bool):
             return  # cancelled by user
 
     try:
-        stream = pod_log(pod_name, container_name=None, follow=follow)
-        for line in stream.stream():
-            print(line.decode("utf-8"), end=None)
+        pod = get_pod(pod_name)
+        eligible_container_names = [BITCOINCORE_CONTAINER, COMMANDER_CONTAINER]
+        available_container_names = [container.name for container in pod.spec.containers]
+        container_name = next(
+            (
+                container_name
+                for container_name in available_container_names
+                if container_name in eligible_container_names
+            ),
+            None,
+        )
+        if not container_name:
+            print("Could not determine primary container.")
+            return
+    except Exception as e:
+        print(f"Error getting pods. Could not determine primary container: {e}")
+        return
+
+    try:
+        stream = pod_log(pod_name, container_name=container_name, follow=follow)
+        for line in stream:
+            click.echo(line.decode("utf-8").rstrip())
     except Exception as e:
         print(e)
     except KeyboardInterrupt:

src/warnet/k8s.py

@@ -4,10 +4,11 @@
 import tempfile
 from pathlib import Path
 from time import sleep
+from typing import Optional
 
 import yaml
 from kubernetes import client, config, watch
-from kubernetes.client.models import CoreV1Event, V1PodList
+from kubernetes.client.models import CoreV1Event, V1Pod, V1PodList
 from kubernetes.client.rest import ApiException
 from kubernetes.dynamic import DynamicClient
 from kubernetes.stream import stream
@@ -41,6 +42,13 @@ def get_pods() -> V1PodList:
     return pod_list
 
 
+def get_pod(name: str, namespace: Optional[str] = None) -> V1Pod:
+    sclient = get_static_client()
+    if not namespace:
+        namespace = get_default_namespace()
+    return sclient.read_namespaced_pod(name=name, namespace=namespace)
+
+
 def get_mission(mission: str) -> list[V1PodList]:
     pods = get_pods()
     crew = []