process_tree.py

#!/usr/bin/env python3
"""Render a process tree from bpftrace telemetry (see bpf-monitor.sh).

The input is a flat, noisy, heavily-duplicated event stream with two record
types, one per line:

  TIME: ... | [EXEC] PID: .. UID: .. Parent: .. COMM: .. | ARGS: <command>
  TIME: ... | [SOCK] PID: .. UID: .. Parent: .. COMM: .. | FAMILY: n TYPE: n PROTO: n

This script reconstructs the parent/child hierarchy from the PID/Parent links
and prints it as a tree, annotating each process with the commands it exec'd
and an aggregated summary of the sockets it opened (decoded to names).

Usage:
    python process_tree.py [telemetry_file]   # default: final_telemetry.txt
    python process_tree.py -h
"""

import argparse
import re
import sys
from collections import OrderedDict


# --- Decoding tables --------------------------------------------------------

AF_FAMILY = {
    0: "AF_UNSPEC",
    1: "AF_UNIX",
    2: "AF_INET",
    10: "AF_INET6",
    16: "AF_NETLINK",
    17: "AF_PACKET",
    33: "AF_RXRPC",
    38: "AF_ALG",
    40: "AF_VSOCK",
}

SOCK_BASE_TYPE = {
    1: "SOCK_STREAM",
    2: "SOCK_DGRAM",
    3: "SOCK_RAW",
    4: "SOCK_RDM",
    5: "SOCK_SEQPACKET",
}

# Flags OR'd into the type argument of socket(2).
SOCK_NONBLOCK = 0o0004000   # 0x800
SOCK_CLOEXEC = 0o02000000   # 0x80000

# Protocol numbers, only meaningful for AF_INET / AF_INET6.
IPPROTO = {
    0: "IP",
    1: "ICMP",
    2: "IGMP",
    6: "TCP",
    17: "UDP",
    58: "ICMPv6",
    132: "SCTP",
}


def decode_family(fam):
    name = AF_FAMILY.get(fam)
    return f"{name}({fam})" if name else f"AF_{fam}"


def decode_type(typ):
    base = typ & ~(SOCK_NONBLOCK | SOCK_CLOEXEC)
    parts = [SOCK_BASE_TYPE.get(base, f"TYPE_{base}")]
    if typ & SOCK_NONBLOCK:
        parts.append("NONBLOCK")
    if typ & SOCK_CLOEXEC:
        parts.append("CLOEXEC")
    return "|".join(parts)


def decode_proto(fam, proto):
    if fam in (2, 10):  # AF_INET / AF_INET6
        name = IPPROTO.get(proto)
        if name:
            return f"{name}({proto})"
    return str(proto)


def decode_socket(fam, typ, proto):
    return f"{decode_family(fam):<14} {decode_type(typ):<28} {decode_proto(fam, proto)}"


# --- Parsing ----------------------------------------------------------------

# Common header: PID/UID/Parent/COMM. Anchored after a [EXEC]/[SOCK] marker.
_HEADER = (
    r"PID:\s*(?P<pid>\d+)\s+UID:\s*(?P<uid>\d+)\s+"
    r"Parent:\s*(?P<parent>\d+)\s+COMM:\s*(?P<comm>\S+)"
)
_TIME = r"TIME:\s*(?P<time>\d+:\d+:\d+)"

EXEC_RE = re.compile(
    rf"{_TIME}.*\[EXEC\]\s+{_HEADER}\s*\|\s*ARGS:\s*(?P<args>.*)$"
)
SOCK_RE = re.compile(
    rf"{_TIME}.*\[SOCK\]\s+{_HEADER}\s*\|\s*"
    r"FAMILY:\s*(?P<family>\d+)\s+TYPE:\s*(?P<type>\d+)\s+PROTO:\s*(?P<proto>\d+)"
)


def parse_line(line):
    """Return ('exec'|'sock', dict) or None.

    Robust against interleaved records: if a real record is embedded inside
    another's ARGS field, we anchor on the *last* [EXEC]/[SOCK] marker so the
    inner (complete) record wins.
    """
    line = line.rstrip("\n")
    if "[EXEC]" not in line and "[SOCK]" not in line:
        return None

    # Trim to the last marker so an embedded record is parsed cleanly.
    e = line.rfind("[EXEC]")
    s = line.rfind("[SOCK]")
    cut = max(e, s)
    # Keep the leading "TIME: ..." of the chosen record by searching backward.
    head = line.rfind("TIME:", 0, cut)
    if head != -1:
        line = line[head:]

    m = SOCK_RE.search(line)
    if m:
        d = m.groupdict()
        return "sock", {
            "time": d["time"],
            "pid": int(d["pid"]),
            "uid": int(d["uid"]),
            "parent": int(d["parent"]),
            "comm": d["comm"],
            "family": int(d["family"]),
            "type": int(d["type"]),
            "proto": int(d["proto"]),
        }

    m = EXEC_RE.search(line)
    if m:
        d = m.groupdict()
        return "exec", {
            "time": d["time"],
            "pid": int(d["pid"]),
            "uid": int(d["uid"]),
            "parent": int(d["parent"]),
            "comm": d["comm"],
            "args": d["args"].strip(),
        }
    return None


# --- Data model -------------------------------------------------------------

class Proc:
    __slots__ = ("pid", "ppid", "uid", "comm", "execs", "sockets",
                 "first_time", "synthetic", "children")

    def __init__(self, pid):
        self.pid = pid
        self.ppid = None
        self.uid = None
        self.comm = None
        self.execs = []                  # ordered, consecutive-deduped commands
        self.sockets = OrderedDict()     # (fam, typ, proto) -> count
        self.first_time = None
        self.synthetic = True            # True until a real event is seen
        self.children = []

    def note_common(self, ev):
        self.synthetic = False
        self.uid = ev["uid"]
        self.comm = ev["comm"]           # latest COMM wins (image may change)
        if ev["parent"] is not None:
            self.ppid = ev["parent"]
        if self.first_time is None:
            self.first_time = ev["time"]


def build_procs(events):
    procs = {}

    def get(pid):
        p = procs.get(pid)
        if p is None:
            p = Proc(pid)
            procs[pid] = p
        return p

    for kind, ev in events:
        p = get(ev["pid"])
        p.note_common(ev)
        get(ev["parent"])  # ensure parent node exists (may stay synthetic)

        if kind == "exec":
            cmd = ev["args"] or f"({ev['comm']})"
            if not p.execs or p.execs[-1] != cmd:
                p.execs.append(cmd)
        else:  # sock
            key = (ev["family"], ev["type"], ev["proto"])
            p.sockets[key] = p.sockets.get(key, 0) + 1

    # Wire up children; determine roots.
    roots = []
    for p in procs.values():
        parent = procs.get(p.ppid) if p.ppid is not None else None
        if parent is not None and parent is not p:
            parent.children.append(p)
        else:
            roots.append(p)
    for p in procs.values():
        p.children.sort(key=lambda c: c.pid)
    roots.sort(key=lambda c: c.pid)
    return procs, roots


def iter_subtree(proc):
    """Yield `proc` and every descendant beneath it (depth-first)."""
    stack = [proc]
    while stack:
        p = stack.pop()
        yield p
        stack.extend(p.children)


def exclude_subtrees(procs, roots, exclude_pids):
    """Remove each PID in `exclude_pids` along with all processes below it.

    Returns `(procs, roots, dropped, missing)` where `dropped` is the set of
    PIDs actually removed (the excluded roots plus their descendants) and
    `missing` is the set of requested PIDs not present in the data.
    """
    drop = set()
    missing = set()
    for pid in exclude_pids:
        p = procs.get(pid)
        if p is None:
            missing.add(pid)
            continue
        for descendant in iter_subtree(p):
            drop.add(descendant.pid)

    if drop:
        procs = {pid: p for pid, p in procs.items() if pid not in drop}
        for p in procs.values():
            p.children = [c for c in p.children if c.pid not in drop]
        roots = [r for r in roots if r.pid not in drop]
    return procs, roots, drop, missing


# --- Rendering --------------------------------------------------------------

def label(p, verbose=True):
    comm = p.comm if p.comm else "?"
    if not verbose:
        # Compact: comm (pid N): <main cmd>
        base = f"{comm} (pid {p.pid})"
        if p.execs:
            return f"{base}: {p.execs[-1]}"
        return base
    uid = "?" if p.uid is None else p.uid
    ppid = "?" if p.ppid is None else p.ppid
    tag = " [not exec'd in log]" if p.synthetic else ""
    return f"{comm} (pid {p.pid}, ppid {ppid}, uid {uid}){tag}"


def render(p, prefix, is_last, lines, verbose, show_sockets):
    connector = "└─ " if is_last else "├─ "
    lines.append(prefix + connector + label(p, verbose))

    # Continuation prefix for this node's detail lines and children.
    child_prefix = prefix + ("   " if is_last else "│  ")

    if verbose:
        for cmd in p.execs:
            lines.append(child_prefix + "  exec: " + cmd)

        if show_sockets and p.sockets:
            lines.append(child_prefix + "  sockets:")
            for (fam, typ, proto), count in p.sockets.items():
                suffix = f"  ×{count}" if count > 1 else ""
                lines.append(child_prefix + "    " + decode_socket(fam, typ, proto) + suffix)

    kids = p.children
    for i, child in enumerate(kids):
        render(child, child_prefix, i == len(kids) - 1, lines, verbose, show_sockets)


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("file", nargs="?", default="final_telemetry.txt",
                    help="telemetry file (default: final_telemetry.txt)")
    ap.add_argument("--no-sockets", action="store_true",
                    help="omit the per-process socket summary (verbose mode)")
    ap.add_argument("-o", "--output", metavar="FILE",
                    help="write the tree to FILE as UTF-8 instead of stdout; "
                         "avoids console/redirect encoding mangling on Windows")
    ap.add_argument("-x", "--exclude", action="append", type=int, default=[],
                    metavar="PID",
                    help="hide this PID and all processes below it; "
                         "repeat to exclude several (e.g. -x 6086 -x 6190)")
    mode = ap.add_mutually_exclusive_group()
    mode.add_argument("-v", "--verbose", action="store_true",
                      help="full detail: exec list, decoded sockets, uid/ppid (default)")
    mode.add_argument("-s", "--simple", action="store_true",
                      help="compact one-line-per-process tree: 'comm (pid N): <cmd>'")
    args = ap.parse_args(argv)

    # Verbose is the default; --simple opts into the reduced view.
    verbose = not args.simple

    # Tree uses box-drawing glyphs; force UTF-8 so Windows cp1252 doesn't choke.
    # Note: when output is redirected with PowerShell's `>`, the shell still
    # re-decodes our bytes with the console codepage and mangles them -- use
    # --output (or `chcp 65001` / Out-File -Encoding utf8) to avoid that.
    try:
        sys.stdout.reconfigure(encoding="utf-8")
    except (AttributeError, ValueError):
        pass

    try:
        with open(args.file, "r", encoding="utf-8", errors="replace") as fh:
            raw_lines = fh.readlines()
    except OSError as exc:
        print(f"error: cannot read {args.file!r}: {exc}", file=sys.stderr)
        return 1

    events = []
    unparsed = 0
    for line in raw_lines:
        if not line.strip():
            continue
        parsed = parse_line(line)
        if parsed is None:
            unparsed += 1
        else:
            events.append(parsed)

    procs, roots = build_procs(events)
    total_procs = len(procs)

    excluded_note = ""
    if args.exclude:
        procs, roots, dropped, missing = exclude_subtrees(procs, roots, args.exclude)
        if missing:
            print(f"warning: excluded PID(s) not found: "
                  f"{', '.join(str(p) for p in sorted(missing))}", file=sys.stderr)
        if dropped:
            excluded_note = f"; excluded {len(dropped)} process(es) via -x"

    n_exec = sum(1 for k, _ in events if k == "exec")
    n_sock = sum(1 for k, _ in events if k == "sock")

    lines = []
    for i, root in enumerate(roots):
        render(root, "", i == len(roots) - 1, lines, verbose, not args.no_sockets)

    summary = (f"Parsed {len(events)} events ({n_exec} exec, {n_sock} sock) "
               f"across {total_procs} processes{excluded_note}; "
               f"{unparsed} line(s) unparsed.\n")
    text = summary + "\n" + "\n".join(lines)

    if args.output:
        try:
            with open(args.output, "w", encoding="utf-8", newline="\n") as fh:
                fh.write(text + "\n")
        except OSError as exc:
            print(f"error: cannot write {args.output!r}: {exc}", file=sys.stderr)
            return 1
    else:
        print(text)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())