Merge pull request #75 from Moomboh/master

Author: uranusjr

src/shellingham/posix/__init__.py

@@ -4,34 +4,51 @@
 from .._core import SHELL_NAMES, ShellDetectionFailure
 from . import proc, ps
 
-
-def _get_process_mapping():
+# Based on QEMU docs: https://www.qemu.org/docs/master/user/main.html
+QEMU_BIN_REGEX = re.compile(
+    r"""qemu-
+        (alpha
+        |armeb
+        |arm
+        |m68k
+        |cris
+        |i386
+        |x86_64
+        |microblaze
+        |mips
+        |mipsel
+        |mips64
+        |mips64el
+        |mipsn32
+        |mipsn32el
+        |nios2
+        |ppc64
+        |ppc
+        |sh4eb
+        |sh4
+        |sparc
+        |sparc32plus
+        |sparc64
+    )""",
+    re.VERBOSE,
+)
+
+
+def _iter_process_parents(pid, max_depth=10):
     """Select a way to obtain process information from the system.
 
     * `/proc` is used if supported.
     * The system `ps` utility is used as a fallback option.
     """
     for impl in (proc, ps):
         try:
-            mapping = impl.get_process_mapping()
+            iterator = impl.iter_process_parents(pid, max_depth)
         except EnvironmentError:
             continue
-        return mapping
+        return iterator
     raise ShellDetectionFailure("compatible proc fs or ps utility is required")
 
 
-def _iter_process_args(mapping, pid, max_depth):
-    """Traverse up the tree and yield each process's argument list."""
-    for _ in range(max_depth):
-        try:
-            proc = mapping[pid]
-        except KeyError:  # We've reached the root process. Give up.
-            break
-        if proc.args:  # Presumably the process should always have a name?
-            yield proc.args
-        pid = proc.ppid  # Go up one level.
-
-
 def _get_login_shell(proc_cmd):
     """Form shell information from SHELL environ if possible."""
     login_shell = os.environ.get("SHELL", "")
@@ -71,6 +88,12 @@ def _get_shell(cmd, *args):
     if cmd.startswith("-"):  # Login shell! Let's use this.
         return _get_login_shell(cmd)
     name = os.path.basename(cmd).lower()
+    if name == "rosetta" or QEMU_BIN_REGEX.fullmatch(name):
+        # If the current process is Rosetta or QEMU, this likely is a
+        # containerized process. Parse out the actual command instead.
+        cmd = args[0]
+        args = args[1:]
+        name = os.path.basename(cmd).lower()
     if name in SHELL_NAMES:  # Command looks like a shell.
         return (name, cmd)
     shell = _get_interpreter_shell(name, args)
@@ -82,8 +105,7 @@ def _get_shell(cmd, *args):
 def get_shell(pid=None, max_depth=10):
     """Get the shell that the supplied pid or os.getpid() is running in."""
     pid = str(pid or os.getpid())
-    mapping = _get_process_mapping()
-    for proc_args in _iter_process_args(mapping, pid, max_depth):
+    for proc_args, _, _ in _iter_process_parents(pid, max_depth):
         shell = _get_shell(*proc_args)
         if shell:
             return shell

src/shellingham/posix/proc.py

@@ -9,11 +9,9 @@
 # NetBSD: https://man.netbsd.org/NetBSD-9.3-STABLE/mount_procfs.8
 # DragonFlyBSD: https://www.dragonflybsd.org/cgi/web-man?command=procfs
 BSD_STAT_PPID = 2
-BSD_STAT_TTY = 5
 
 # See https://docs.kernel.org/filesystems/proc.html
 LINUX_STAT_PPID = 3
-LINUX_STAT_TTY = 6
 
 STAT_PATTERN = re.compile(r"\(.+\)|\S+")
 
@@ -41,14 +39,14 @@ def _use_bsd_stat_format():
         return False
 
 
-def _get_stat(pid, name):
+def _get_ppid(pid, name):
     path = os.path.join("/proc", str(pid), name)
     with io.open(path, encoding="ascii", errors="replace") as f:
         parts = STAT_PATTERN.findall(f.read())
     # We only care about TTY and PPID -- both are numbers.
     if _use_bsd_stat_format():
-        return parts[BSD_STAT_TTY], parts[BSD_STAT_PPID]
-    return parts[LINUX_STAT_TTY], parts[LINUX_STAT_PPID]
+        return parts[BSD_STAT_PPID]
+    return parts[LINUX_STAT_PPID]
 
 
 def _get_cmdline(pid):
@@ -66,21 +64,13 @@ class ProcFormatError(EnvironmentError):
     pass
 
 
-def get_process_mapping():
+def iter_process_parents(pid, max_depth=10):
     """Try to look up the process tree via the /proc interface."""
     stat_name = detect_proc()
-    self_tty = _get_stat(os.getpid(), stat_name)[0]
-    processes = {}
-    for pid in os.listdir("/proc"):
-        if not pid.isdigit():
-            continue
-        try:
-            tty, ppid = _get_stat(pid, stat_name)
-            if tty != self_tty:
-                continue
-            args = _get_cmdline(pid)
-            processes[pid] = Process(args=args, pid=pid, ppid=ppid)
-        except IOError:
-            # Process has disappeared - just ignore it.
-            continue
-    return processes
+    for _ in range(max_depth):
+        ppid = _get_ppid(pid, stat_name)
+        args = _get_cmdline(pid)
+        yield Process(args=args, pid=pid, ppid=ppid)
+        if ppid == "0":
+            break
+        pid = ppid

src/shellingham/posix/ps.py

@@ -9,7 +9,7 @@ class PsNotAvailable(EnvironmentError):
     pass
 
 
-def get_process_mapping():
+def iter_process_parents(pid, max_depth=10):
     """Try to look up the process tree via the output of `ps`."""
     try:
         cmd = ["ps", "-ww", "-o", "pid=", "-o", "ppid=", "-o", "args="]
@@ -22,12 +22,13 @@ def get_process_mapping():
         # `ps` can return 1 if the process list is completely empty.
         # (sarugaku/shellingham#15)
         if not e.output.strip():
-            return {}
+            return
         raise
     if not isinstance(output, str):
         encoding = sys.getfilesystemencoding() or sys.getdefaultencoding()
         output = output.decode(encoding)
-    processes = {}
+
+    processes_mapping = {}
     for line in output.split("\n"):
         try:
             pid, ppid, args = line.strip().split(None, 2)
@@ -39,5 +40,12 @@ def get_process_mapping():
             args = tuple(a.strip() for a in args.split(" "))
         except ValueError:
             continue
-        processes[pid] = Process(args=args, pid=pid, ppid=ppid)
-    return processes
+        processes_mapping[pid] = Process(args=args, pid=pid, ppid=ppid)
+
+    for _ in range(max_depth):
+        try:
+            process = processes_mapping[pid]
+        except KeyError:
+            return
+        yield process
+        pid = process.ppid