control: fetch and stream pod logs without kubectl

Author: pinheadmz

src/warnet/control.py

@@ -21,6 +21,7 @@
     get_mission,
     get_pods,
     snapshot_bitcoin_datadir,
+    pod_log
 )
 from .process import run_command, stream_command
 
@@ -235,46 +236,39 @@ def run(scenario_file: str, additional_args: tuple[str]):
 @click.option("--follow", "-f", is_flag=True, default=False, help="Follow logs")
 def logs(pod_name: str, follow: bool):
     """Show the logs of a pod"""
-    follow_flag = "--follow" if follow else ""
     namespace = get_default_namespace()
 
-    if pod_name:
+    if pod_name == "":
         try:
-            command = f"kubectl logs pod/{pod_name} -n {namespace} {follow_flag}"
-            stream_command(command)
-            return
+            pods = get_pods()
+            pod_list = [item.metadata.name for item in pods.items]
         except Exception as e:
-            print(f"Could not find the pod {pod_name}: {e}")
+            print(f"Could not fetch any pods in namespace {namespace}: {e}")
+            return
 
-    try:
-        pods = run_command(f"kubectl get pods -n {namespace} -o json")
-        pods = json.loads(pods)
-        pod_list = [item["metadata"]["name"] for item in pods["items"]]
-    except Exception as e:
-        print(f"Could not fetch any pods in namespace {namespace}: {e}")
-        return
+        if not pod_list:
+            print(f"Could not fetch any pods in namespace {namespace}")
+            return
 
-    if not pod_list:
-        print(f"Could not fetch any pods in namespace {namespace}")
-        return
+        q = [
+            inquirer.List(
+                name="pod",
+                message="Please choose a pod",
+                choices=pod_list,
+            )
+        ]
+        selected = inquirer.prompt(q, theme=GreenPassion())
+        if selected:
+            pod_name = selected["pod"]
+        else:
+            return  # cancelled by user
 
-    q = [
-        inquirer.List(
-            name="pod",
-            message="Please choose a pod",
-            choices=pod_list,
-        )
-    ]
-    selected = inquirer.prompt(q, theme=GreenPassion())
-    if selected:
-        pod_name = selected["pod"]
-        try:
-            command = f"kubectl logs pod/{pod_name} -n {namespace} {follow_flag}"
-            stream_command(command)
-        except Exception as e:
-            print(f"Please consider waiting for the pod to become available. Encountered: {e}")
-    else:
-        pass  # 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)
+    except Exception as e:
+        print(e)
 
 
 @click.command()

src/warnet/k8s.py

@@ -9,6 +9,7 @@
 from kubernetes.client.models import CoreV1Event, V1PodList
 from kubernetes.dynamic import DynamicClient
 from kubernetes.stream import stream
+from kubernetes.client.rest import ApiException
 
 from .constants import (
     CADDY_INGRESS_NAME,
@@ -282,3 +283,17 @@ def get_ingress_ip_or_host():
     except Exception as e:
         print(f"Error getting ingress IP: {e}")
         return None
+
+
+def pod_log(pod_name, container_name=None, follow=False):
+    sclient = get_static_client()
+    try:
+        return sclient.read_namespaced_pod_log(
+            name=pod_name,
+            namespace=get_default_namespace(),
+            container=container_name,
+            follow=follow,
+            _preload_content=False
+        )
+    except ApiException as e:
+        raise Exception(json.loads(e.body.decode('utf-8'))["message"])