refactoring for bigger use

Author: Avi-Robusta

src/java_toolkit/utils/pod_ps.py …c/java_toolkit/common/utils/processes.pysrc/java_toolkit/utils/pod_ps.py renamed to src/java_toolkit/common/utils/processes.py src/java_toolkit/utils/pod_ps.py renamed to src/java_toolkit/common/utils/processes.py

@@ -1,6 +1,7 @@
 #!/usr/local/bin/python3
 import re
 from typing import List
+from .remote_ns_cmd import run_cmd_in_proc_namespace
 
 import psutil
 import typer
@@ -14,6 +15,8 @@
     r"\d+:.+:/kubepods\.slice/kubepods-[^/]+\.slice/kubepods-[^/]+-pod([^/]+)\.slice/docker-([0-9a-f]{64})"
 )
 
+ps_regex = re.compile(r"\s*(?P<pid>.*?)\s+(?P<cmd>.*?)\s+(?P<cmd_args>.*?)\n")
+
 app = typer.Typer()
 
 
@@ -61,3 +64,25 @@ def get_pod_processes(pod_uid: str) -> ProcessList:
             processes.append(Process(pid=pid, exe=proc.exe(), cmdline=proc.cmdline()))
     return ProcessList(processes=processes)
 
+def get_ns_processes(pid_in_ns: int, verbose: bool):
+    cmd = "ps -o pid,comm,args | tail -n +2 "
+    processes = run_cmd_in_proc_namespace(pid_in_ns, cmd, verbose)
+    processes = processes.split('\n')
+    return_proc_list=[]
+    processes.remove(processes[0]) #title column
+    for proc_line in processes:
+        if "tail -n" in proc_line or "ps -o pid,comm,args" in proc_line:
+            continue
+        match = ps_regex.match(proc_line)
+        pid = int(match.group("pid"))
+        exe = match.group("cmd")
+        cmdline = match.group("cmd_args")
+        if cmdline is None:
+            cmdline = []
+        else:
+            cmdline = cmdline.split(' ')
+        return_proc_list.append(Process(pid=pid, exe=exe, cmdline=cmdline))
+    return return_proc_list
+
+
+

src/java_toolkit/common/utils/remote_ns_cmd.py

@@ -0,0 +1,24 @@
+import subprocess
+import typer
+
+
+def run_command(cmd: str, verbose: bool):
+    if verbose:
+        typer.echo(f"Running {cmd}")
+    proc = subprocess.Popen(
+        cmd, shell=True, stdin=subprocess.PIPE, stderr=subprocess.STDOUT
+    )
+    (output, err) = proc.communicate()
+
+    proc.stdout.close()
+    return_code = proc.wait()
+    if return_code:
+        raise subprocess.CalledProcessError(return_code, cmd)
+    if verbose:
+        typer.echo(output)
+    return output, err
+
+def run_cmd_in_proc_namespace(pid, command_to_run, verbose):
+    NSENTER_CMD = "nsenter -t {} -p -m {}"
+    nsenter_cmd_formatted = NSENTER_CMD.format(pid, command_to_run)
+    return run_command(nsenter_cmd_formatted, verbose)

src/java_toolkit/common/utils/tmp_mount.py

@@ -0,0 +1,32 @@
+from .remote_ns_cmd import run_command
+from os.path import join, basename
+
+class TmpRemotePodMounter(object):
+    remote_mnt_path =""
+    local_mnt_path =""
+    src_dir =""
+    verbose = False
+    MKDIR_POD_CMD="mkdir -p {}"
+    RMDIR_POD_CMD="rm -R {}"
+    COPY_CMD="cp -R {} {}"
+    ROOT_PROC_NS="/proc/{}/cwd",
+
+    def __init__(self, pid: str, src_dir: str, local_mnt_path: str, verbose: bool):
+        self.verbose = verbose
+        self.src_dir = src_dir
+        self.local_mnt_path = local_mnt_path
+        self.remote_mnt_path = join(self.ROOT_PROC_NS.format(pid), local_mnt_path)
+
+    def __enter__(self):
+        mkdir_cmd = self.MKDIR_POD_CMD.format(self.remote_mnt_path)
+        run_command(mkdir_cmd, self.verbose)
+        cp_cmd=self.COPY_CMD.format(self.src_dir, self.remote_mnt_path)
+        run_command(cp_cmd, self.verbose)
+        return self
+
+    def get_mounted_jdk_dir(self):
+            return join(self.local_mnt_path, basename(self.src_dir) )
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        rm_dir_cmd = self.RMDIR_POD_CMD.format(self.remote_mnt_path)
+        run_command(rm_dir_cmd, self.verbose)

src/java_toolkit/configs.json

@@ -1,8 +1,5 @@
 {
   "JDK_PATH": "/app/openjdk",
-  "MKDIR_POD_CMD": "mkdir {DST_MOUNT_PATH}",
-  "RMDIR_POD_CMD": "rm -R {DST_MOUNT_PATH}",
-  "DST_MOUNT_PATH": "/proc/{PID}/cwd/mnt/robusta",
   "COPY_CMD": "cp -R {JDK_PATH} {DST_MOUNT_PATH}",
   "JDK_NAME": "openjdk",
   "JMAP_CMD": "{JDK_PATH}/bin/jmap -J-Xshareclasses:nonfatal -histo {LOCAL_PID}",

src/java_toolkit/configs.py

@@ -1,10 +1,5 @@
 JDK_PATH= "/app/openjdk"
-MKDIR_POD_CMD= 'mkdir -p {}'
-RMDIR_POD_CMD= "rm -R {}"
-DST_MOUNT_PATH= "/proc/{}/cwd/mnt/robusta"
 LOCAL_MOUNT_PATH= "/mnt/robusta"
-COPY_CMD= "cp -R {} {}"
 JDK_NAME= "openjdk"
 JMAP_CMD= "{}/bin/jmap -J-Xshareclasses:nonfatal -histo:live {}"
 JSTACK_CMD= "{}/bin/jstack -J-Xshareclasses:nonfatal -l {}"
-NSENTER_CMD= "nsenter -t {} -p -m {}"

src/java_toolkit/main.py

@@ -1,73 +1,32 @@
 #!/usr/local/bin/python3
 import subprocess
 import typer
-from .utils.pod_ps import get_pod_processes
+from .utils.process import get_pod_processes, get_ns_processes
+from .utils.remote_ns_cmd import run_cmd_in_proc_namespace
+from .utils.tmp_mount import TmpRemotePodMounter
 from .configs import *
 from os.path import join
 
 app = typer.Typer()
 
-
-def run_command(cmd: str, verbose: bool):
-    if verbose:
-        typer.echo(f"running {cmd}")
-    output = subprocess.Popen(
-        cmd, shell=True, stdin=subprocess.PIPE, stderr=subprocess.STDOUT
-    )
-    return_code = output.wait()
-    if return_code:
-        raise subprocess.CalledProcessError(return_code, cmd)
-    if verbose:
-        pass
-        #typer.echo(output.decode())
-
-
-class JDKMounter(object):
-    mnt_path =""
-    verbose = False
-
-    def __init__(self, pid: str, verbose: bool):
-        self.verbose = verbose
-        self.mnt_path = DST_MOUNT_PATH.format(pid)
-
-    def __enter__(self):
-        mkdir_cmd = MKDIR_POD_CMD.format(self.mnt_path)
-        run_command(mkdir_cmd, self.verbose)
-        cp_cmd=COPY_CMD.format(JDK_PATH, self.mnt_path)
-        run_command(cp_cmd, self.verbose)
-        return self
-
-    def get_mounted_jdk_dir(self):
-            return join(LOCAL_MOUNT_PATH, JDK_NAME )
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        rm_dir_cmd = RMDIR_POD_CMD.format(self.mnt_path)
-        run_command(rm_dir_cmd, self.verbose)
-
-def run_cmd_in_proc_namespace(pid, command_to_run, verbose):
-    nsenter_cmd_formatted = NSENTER_CMD.format(pid, command_to_run)
-    run_command(nsenter_cmd_formatted, verbose)
-
 @app.command()
 def pod_ps(pod_uid: str):
-    typer.echo(get_pod_processes(pod_uid).json())
+    app.echo(get_pod_processes(pod_uid).json())
 
 @app.command()
-def find_pid(pod_uid: str, cmdline: str, exe: str):
-    for proc in get_pod_processes(pod_uid).processes:
-        if cmdline in " ".join(proc.cmdline) and exe in proc.exe:
-            typer.echo(proc.pid)
+def pod_ns_ps(pid: str, verbose: bool = True):
+    app.echo(get_ns_processes(pid, verbose).json())
 
 @app.command()
 def jmap(pid: int, verbose: bool = True):
-    with JDKMounter(pid, verbose) as jdk_mounter:
+    with TmpRemotePodMounter(pid, JDK_PATH, LOCAL_MOUNT_PATH, verbose) as jdk_mounter:
         jstack_cmd = JMAP_CMD.format(jdk_mounter.get_mounted_jdk_dir(), 1)
         run_cmd_in_proc_namespace(pid, jstack_cmd, verbose)
 
 
 @app.command()
 def jstack(pid: int, verbose: bool = True):
-    with JDKMounter(pid, verbose) as jdk_mounter:
+    with TmpRemotePodMounter(pid, JDK_PATH, LOCAL_MOUNT_PATH, verbose) as jdk_mounter:
         jstack_cmd = JSTACK_CMD.format(jdk_mounter.get_mounted_jdk_dir(), 1)
         run_cmd_in_proc_namespace(pid, jstack_cmd, verbose)