Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix #74 #75

Merged
merged 5 commits into from
May 18, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 40 additions & 18 deletions src/shellingham/posix/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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", "")
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down
32 changes: 11 additions & 21 deletions src/shellingham/posix/proc.py
Original file line number Diff line number Diff line change
Expand Up @@ -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+")

Expand Down Expand Up @@ -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):
Expand All @@ -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
18 changes: 13 additions & 5 deletions src/shellingham/posix/ps.py
Original file line number Diff line number Diff line change
Expand Up @@ -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="]
Expand All @@ -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)
Expand All @@ -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