pskill.py

#!/usr/bin/env python3
"""
Allows the user to select a process with fzf, then performs the specified
action on it (send SIGTERM by default).
"""
import argparse
import os
import re
import shlex
import subprocess
import tempfile

cli = argparse.ArgumentParser()
args, invocation = cli.parse_known_args()
if not invocation:
    invocation = ['kill']

# short notation for specifying signals directly instead of needing to write 'kill'
if invocation[0].startswith('-'):
    invocation = ['kill'] + invocation

# short notation for invoking common utilities which require '-p'
if len(invocation) == 1 and invocation[0] in {'gdb', 'strace', 'ltrace'}:
    invocation.append('-p')

processes = {}
my_pid = os.getpid()
for entry in os.listdir('/proc'):
    try:
        pid = int(entry)
    except ValueError:
        continue
    if pid == my_pid:
        # ignore self
        continue
    try:
        status = open(f'/proc/{pid}/status', 'rb').read().decode(errors='replace')
        match = re.search(r'^Tgid:\s*([1-9][0-9]*)$', status, re.MULTILINE)
        tgid = int(match.group(1))
        if pid != tgid:
            # this is a thread, not a process
            continue
        cmdline = open(f'/proc/{pid}/cmdline', 'rb').read().decode(errors='replace')
        if not cmdline:
            # this is a kernel thread
            continue
    except FileNotFoundError:
        # ???
        continue
    # properly escape the cmdline
    cmdline = ' '.join(shlex.quote(part) for part in cmdline[:-1].split('\0'))
    processes[pid] = cmdline

max_pid_len = len(str(max(processes, default=0)))
with tempfile.NamedTemporaryFile("w") as tempfile:
    for pid, cmdline in reversed(sorted(processes.items())):
        tempfile.write(f'{pid:{max_pid_len}d} {cmdline}\n')
    tempfile.flush()
    selector = subprocess.Popen(
        f"fzf --print0 < " + shlex.quote(tempfile.name),
        shell=True,
        env=os.environ,
        stdout=subprocess.PIPE,
    )
    choice = selector.communicate()[0].decode().rstrip("\0")
print(f'selected process: {choice!r}')
try:
    pid = int(choice.split()[0])
except:
    raise SystemExit(1) from None

for idx in range(1, len(invocation)):
    if '{pid}' in invocation[idx]:
        invocation[idx] = invocation[idx].replace('{pid}', str(pid))
        break
else:
    invocation.append(str(pid))

print(' '.join(shlex.quote(part) for part in invocation))
os.execvp(invocation[0], invocation)