#!python
import ctypes
import itertools
import os
import platform
import re
import shlex
import signal
import subprocess
import sys
import traceback
from dataclasses import dataclass, field
from pwd import getpwnam
from typing import Any, List, Tuple

try:
    import bcc
except ImportError:
    print(
        fr"""Please install bcc

On Debian/Ubuntu, try this:

    sudo apt install python3-bpfcc

On Fedora, try this:

    sudo dnf install python3-bcc
"""
    )
    sys.exit(1)

PROJECT = "gdb-pounce"
REQUIRED_SYSTEM = "Linux"
REQUIRED_RELEASE = "5.3"
REQUIRED_BCC_VERSION = "0.15.0"
DEBIAN_BCC_REPO = "https://salsa.debian.org/debian/bpfcc.git"
DEBIAN_BCC_TAG = "debian/0.16.0-3"
TEXT = r"""#include <linux/mm_types.h>
#include <linux/sched.h>
#include <uapi/linux/bpf.h>
#include <uapi/linux/ptrace.h>
#include <uapi/linux/signal.h>

struct key {
    u32 prefixlen;
    char data[TASK_COMM_LEN];
};

struct __attribute__((__packed__)) event {
    u32 pid;
    u8 skipped;
    long bpf_send_signal_ret;
};

#define XARG_MAX 2000

BPF_ARRAY(uid_map, u32, 1);
BPF_LPM_TRIE(comms_map, struct key, u8);
BPF_ARRAY(cmdline_hash_map, u64, 1);
BPF_ARRAY(cmdline_len_map, u32, 1);
BPF_PERF_OUTPUT(events);

static bool cmdline_may_match()
{
    const char __user *arg_start;
    const char __user *arg_end;
    struct task_struct *t;
    struct mm_struct *mm;
    unsigned char c = 0;
    u64 *cmdline_hash;
    u32 *cmdline_len;
    u64 hash = 0;
    int zero = 0;
    int len;
    int i;

    cmdline_hash = cmdline_hash_map.lookup(&zero);
    if (cmdline_hash == NULL)
        return true;
    cmdline_len = cmdline_len_map.lookup(&zero);
    if (cmdline_len == NULL || *cmdline_len > XARG_MAX)
        return true;
    t = (struct task_struct *)bpf_get_current_task();
    if (bpf_probe_read_kernel(&mm, sizeof(mm), &t->mm) != 0)
        return true;
    if (bpf_probe_read_kernel(&arg_start, sizeof(arg_start), &mm->arg_start) !=
        0)
        return true;
    if (bpf_probe_read_kernel(&arg_end, sizeof(arg_end), &mm->arg_end) != 0)
        return true;
    if (arg_end < arg_start)
        return true;
    len = arg_end - arg_start;
    if (len > XARG_MAX)
        return true;

    /* Always process XARG_MAX bytes in order to make the loop analysis easier
     * for the verifier. */
    for (i = 0; i < XARG_MAX; i++) {
        if (i >= *cmdline_len) {
            if (bpf_probe_read_user(&c, 1, arg_start + i - *cmdline_len) != 0)
                return i < len;
            hash -= c;
        }
        if (bpf_probe_read_user(&c, 1, arg_start + i) != 0)
            return i < len;
        hash += c;
        if (hash == *cmdline_hash)
            return i < len;
    }

    return false;
}

static int maybe_pounce(struct pt_regs *ctx)
{
    struct event event = {
        .pid = bpf_get_current_pid_tgid() >> 32,
        .skipped = 1,
        .bpf_send_signal_ret = 0,
    };
    size_t comm_len;
    struct key key;
    int zero = 0;
    u32 *uid;

    if (bpf_get_current_comm(key.data, sizeof(key.data)) != 0)
        return 0;
    for (comm_len = 0; comm_len < TASK_COMM_LEN; comm_len++)
        if (key.data[comm_len] == 0)
            break;
    key.prefixlen = (comm_len + 1) * 8;
    if (comms_map.lookup(&key) == NULL)
        return 0;

    uid = uid_map.lookup(&zero);
    if (uid != NULL && *uid != 0xFFFFFFFF) {
        u32 current_uid = bpf_get_current_uid_gid() & 0xFFFFFFFF;

        if (current_uid != *uid)
            goto submit;
    }

    if (!cmdline_may_match())
        goto submit;

    event.skipped = 0;
    event.bpf_send_signal_ret = bpf_send_signal(SIGSTOP);

submit:
    events.perf_submit(ctx, &event, sizeof(event));
    return 0;
}

int exec_retprobe(struct pt_regs *ctx)
{
    if (PT_REGS_RC(ctx) != 0)
        return 0;

    return maybe_pounce(ctx);
}

int schedule_tail_retprobe(struct pt_regs *ctx)
{
    u64 current_pid_tgid = bpf_get_current_pid_tgid();

    if ((current_pid_tgid & 0xFFFFFFFF) != (current_pid_tgid >> 32))
        return 0;

    return maybe_pounce(ctx);
}
"""

TASK_COMM_LEN = 16
COMM_T = ctypes.c_uint8 * TASK_COMM_LEN


class Key(ctypes.Structure):
    _fields_ = [("prefixlen", ctypes.c_uint32), ("data", COMM_T)]


def version_is_less(v1, v2):
    for w1, w2 in itertools.zip_longest(v1.split("."), v2.split("."), fillvalue="0"):
        try:
            w1, w2 = int(w1), int(w2)
        except ValueError:
            pass
        if w1 < w2:
            return True
        elif w1 > w2:
            return False
    return False


def check_system_requirements():
    actual_system = platform.system()
    if actual_system != REQUIRED_SYSTEM:
        print(
            f"{PROJECT} can only run on {REQUIRED_SYSTEM}, "
            f"but you have {actual_system}",
            file=sys.stderr,
        )
        sys.exit(1)

    actual_release = platform.release()
    if version_is_less(actual_release, REQUIRED_RELEASE):
        print(
            f"Please upgrade kernel: you have {actual_release}, "
            f"but {PROJECT} requires {REQUIRED_RELEASE}",
            file=sys.stderr,
        )
        sys.exit(1)

    actual_bcc_version = bcc.__version__
    if version_is_less(actual_bcc_version, REQUIRED_BCC_VERSION):
        print(
            rf"""Please upgrade bcc: you have {actual_bcc_version}, but {PROJECT} requires {REQUIRED_BCC_VERSION}

On Debian/Ubuntu, try this:

    git clone --branch={DEBIAN_BCC_TAG} {DEBIAN_BCC_REPO}
    cd bpfcc
    # Add -d to the below in order to ignore the missing debhelper-compat.
    gbp buildpackage \
        --git-builder="debuild --set-envvar=DPKG_GENSYMBOLS_CHECK_LEVEL=0 -i -I" \
        --git-export-dir=build \
        --git-ignore-branch \
        --unsigned-changes \
        --unsigned-source
    sudo dpkg -i build/*.deb
""",
            file=sys.stderr,
        )
        sys.exit(1)


class SigintHandler:
    def __init__(self):
        self.raise_keyboard_interrupt = True
        signal.signal(signal.SIGINT, self)

    def __call__(self, signum, frame):
        if self.raise_keyboard_interrupt:
            raise KeyboardInterrupt()


def print_help(file):
    print(
        """usage: gdb-pounce [options] [gdb_arg [gdb_arg ...]] name [arg [arg ...]]

Attach gdb to a new process before it had a chance to run.

positional arguments:
  gdb_arg     gdb arguments
  name        basename of an executable to wait for
  arg         arguments to search in executable's cmdline

optional arguments:
  -h, --help  show this help message and exit
  --async     allow multiple gdbs in separate terminals
  --fork      match fork()ed instead of execve()d processes
  --strace    use strace instead of gdb
  --uid=UID   match only processes with the specified real user id""",
        file=file,
    )


def add_dashes(*names):
    result = set()
    for name in names:
        result.add(f"-{name}")
        result.add(f"--{name}")
    return result


def linux_ns_id(pid, ns):
    try:
        st = os.stat(f"/proc/{pid}/ns/{ns}")
    except FileNotFoundError:
        return None
    return st.st_dev, st.st_ino


def linux_ns_pid(pid):
    prefix = "NSpid:"
    with open(f"/proc/{pid}/status") as fp:
        for s in fp:
            if s.startswith(prefix):
                return int(s[len(prefix) :].split()[-1])
    return None


def proc_state(pid):
    try:
        with open(f"/proc/{pid}/stat") as fp:
            s = fp.readline()
    except FileNotFoundError:
        return None
    m = re.search(r"\) (.) ", s)
    if m is None:
        return m
    return m.group(1)


class Gdb:
    DISPLAY_NAME = "GDB"
    # From gdb --help and man gdb.
    FLAGS = add_dashes(
        "batch",
        "batch-silent",
        "configuration",
        "dbx",
        "f",
        "fullname",
        "n",
        "nh",
        "nw",
        "nx",
        "q",
        "quiet",
        "readnever",
        "readnow",
        "return-child-result",
        "silent",
        "tui",
        "w",
        "write",
    )
    OPTIONS = add_dashes(
        "D",
        "b",
        "c",
        "cd",
        "command",
        "core",
        "core",
        "d",
        "data-directory",
        "directory",
        "e",
        "eval-command",
        "ex",
        "exec",
        "iex",
        "init-command",
        "init-eval-command",
        "interpreter",
        "ix",
        "l",
        "pid",
        "s",
        "se",
        "symbols",
        "tty",
        "x",
    )

    def gen_argv(self, args, pid):
        argv = []
        if linux_ns_id(os.getpid(), "pid") != linux_ns_id(pid, "pid"):
            # Try to prevent "warning: Target and debugger are in different PID
            # namespaces".
            ns_pid = linux_ns_pid(pid)
            if ns_pid is None:
                print(
                    "Kernel does not support NSpid - attaching with the host GDB...",
                    file=sys.stderr,
                )
            else:
                nsenter = "nsenter", "-a", "-t", str(pid)
                try:
                    subprocess.check_call(
                        nsenter + ("gdb", "-batch"),
                        stdin=subprocess.DEVNULL,
                        stdout=subprocess.DEVNULL,
                        stderr=subprocess.DEVNULL,
                    )
                except FileNotFoundError:
                    print(
                        "nsenter is not installed on the host - attaching with the host GDB...",
                        file=sys.stderr,
                    )
                except subprocess.CalledProcessError:
                    print(
                        "GDB is not installed in the container - attaching with the host GDB...",
                        file=sys.stderr,
                    )
                else:
                    argv.extend(nsenter)
                    pid = ns_pid

                    # Use bash -i instead of setsid, since it restores the
                    # controlling tty.
                    setsid = "bash", "-i", "-c", '"$@"', "--"
                    try:
                        subprocess.check_call(
                            nsenter + setsid + ("true",),
                            stdin=subprocess.DEVNULL,
                            stdout=subprocess.DEVNULL,
                            stderr=subprocess.DEVNULL,
                        )
                    except subprocess.CalledProcessError:
                        print(
                            "bash is not installed in the container - Ctrl+C will terminate GDB...",
                            file=sys.stderr,
                        )
                    else:
                        argv.extend(setsid)

        argv.extend(
            (
                "gdb",
                "-p",
                str(pid),
                # When attaching to a SIGSTOPped process, gdb receives an event
                # that indicates that the process got a SIGSTOP. It then passes
                # this signal to the process after each resumption, making
                # debugging impossible. Maybe this should be reported?
                "-ex",
                "handle SIGSTOP nostop noprint nopass",
            )
        )
        return argv

    def resume(self, pid):
        pass

    def cleanup(self, pid):
        if proc_state(pid) == "T":
            print("GDB left the process stopped - sending SIGCONT...", file=sys.stderr)
            os.kill(pid, signal.SIGCONT)


class Strace:
    DISPLAY_NAME = "strace"
    # From strace --help and man strace.
    FLAGS = add_dashes(
        "A",
        "C",
        "D",
        "DD",
        "DDD",
        "F",
        "T",
        "Z",
        "c",
        "d",
        "debug",
        "f",
        "ff",
        "i",
        "instruction-pointer",
        "k",
        "no-abbrev",
        "output-append-mode",
        "q",
        "qq",
        "r",
        "seccomp-bpf",
        "stack-traces",
        "summary",
        "summary-only",
        "summary-wall-clock",
        "t",
        "tt",
        "ttt",
        "v",
        "w",
        "x",
        "xx",
        "y",
        "yy",
        "z",
    )
    OPTIONS = add_dashes(
        "E",
        "I",
        "O",
        "P",
        "S",
        "X",
        "a",
        "abbrev",
        "attach",
        "b",
        "columns",
        "const-print-style",
        "detach-on",
        "e",
        "env",
        "fault",
        "inject",
        "kvm",
        "o",
        "output",
        "p",
        "raw",
        "read",
        "s",
        "s",
        "signal",
        "status",
        "string-limit",
        "summary-sort-by",
        "trace",
        "trace-path",
        "u",
        "user",
        "verbose",
        "write",
    )

    def gen_argv(self, args, pid):
        return ["strace", "-p", str(pid)]

    def resume(self, pid):
        while True:
            state = proc_state(pid)
            if state is None:
                break
            if state == "t":
                os.kill(pid, signal.SIGCONT)
                break

    def cleanup(self, pid):
        pass


def get_val(i):
    if i + 2 > len(sys.argv):
        print(
            f"""{sys.argv[0]}: option '{sys.argv[i]}' requires an argument
Use `{sys.argv[0]} --help' for a complete list of options.""",
            file=sys.stderr,
        )
        sys.exit(1)
    return sys.argv[i + 1]


def resolve_uid(s):
    try:
        return int(s)
    except ValueError:
        return getpwnam(s)[2]


@dataclass
class Args:
    async_: bool = False
    fork: bool = False
    uid: int = 0xFFFFFFFF
    tool_args: List[str] = field(default_factory=list)
    tool: Any = field(default_factory=Gdb)


def parse_argv() -> Tuple[Args, str, List[str]]:
    args = Args()
    i = 1
    while i < len(sys.argv):
        arg = sys.argv[i]
        if arg in ("-h", "--help"):
            print_help(sys.stdout)
            sys.exit(0)
        elif arg == "--async":
            args.async_ = True
            i += 1
        elif arg == "--fork":
            args.fork = True
            i += 1
        elif arg.startswith("--uid="):
            args.uid = resolve_uid(arg[6:])
            i += 1
        elif arg == "--uid":
            args.uid = resolve_uid(get_val(i))
            i += 2
        elif arg in args.tool.FLAGS or (
            "=" in arg and arg[: arg.index("=")] in args.tool.OPTIONS
        ):
            args.tool_args.append(arg)
            i += 1
        elif arg in args.tool.OPTIONS:
            args.tool_args.append(arg)
            args.tool_args.append(get_val(i))
            i += 2
        elif arg == "--strace":
            args.tool = Strace()
            i += 1
        else:
            return args, arg, sys.argv[i + 1 :]
    print_help(sys.stderr)
    sys.exit(1)


def print_skip_non_matching(pid, filtered_by):
    print(
        f"Skipping non-matching pid {pid} (filtered by {filtered_by})...",
        file=sys.stderr,
    )


def main():
    check_system_requirements()

    args, name, rest = parse_argv()

    cmdline_needle = b"\0"
    for arg in rest:
        cmdline_needle += arg.encode() + b"\0"

    try:
        bpf = bcc.BPF(text=TEXT)
    except Exception:
        print(file=sys.stderr)
        traceback.print_exc()
        if os.getuid() != 0:
            suggestion = " ".join(shlex.quote(arg) for arg in sys.argv)
            print(
                f"""
Are you root? Do you have CAP_SYS_ADMIN / CAP_BPF? Try this:
sudo env \"PATH=$PATH\" {suggestion}""",
                file=sys.stderr,
            )
        sys.exit(1)
    bpf["uid_map"][0] = ctypes.c_uint32(args.uid)
    comm_bytes = name[: TASK_COMM_LEN - 1].encode() + b"\0"
    key = Key(prefixlen=len(comm_bytes) * 8, data=COMM_T(*comm_bytes))
    bpf["comms_map"][key] = ctypes.c_uint8(0)
    bpf["cmdline_len_map"][0] = ctypes.c_uint32(len(cmdline_needle))
    bpf["cmdline_hash_map"][0] = ctypes.c_uint64(sum(cmdline_needle))
    if args.fork:
        # schedule_tail - first thing a freshly forked thread must call.
        bpf.attach_kretprobe(event="schedule_tail", fn_name="schedule_tail_retprobe")
    else:
        for syscall in ("execve", "execveat"):
            bpf.attach_kretprobe(
                event=bpf.get_syscall_fnname(syscall), fn_name="exec_retprobe"
            )
    pids = []

    def callback(cpu, data, size):
        # Do not put bpf["events"] into a variable, this will cause:
        # Exception ignored in: <function PerfEventArray.__del__ at 0x7f55b11c2f70>
        # Traceback (most recent call last):
        #   File "/usr/lib/python3.8/site-packages/bcc/table.py", line 584, in __del__
        #     del self[key]
        #   File "/usr/lib/python3.8/site-packages/bcc/table.py", line 590, in __delitem__
        #     super(PerfEventArray, self).__delitem__(key)
        #   File "/usr/lib/python3.8/site-packages/bcc/table.py", line 493, in __delitem__
        #     super(ArrayBase, self).__delitem__(key)
        #   File "/usr/lib/python3.8/site-packages/bcc/table.py", line 262, in __delitem__
        #     raise KeyError
        # Report this to https://github.com/iovisor/bcc?
        event = bpf["events"].event(data)
        if event.skipped:
            print_skip_non_matching(event.pid, "BPF")
        elif event.bpf_send_signal_ret == 0:
            pids.append(event.pid)
        else:
            print(
                f"Skipping pid {event.pid}, because it could not be stopped"
                f" ({event.bpf_send_signal_ret})...",
                file=sys.stderr,
            )

    bpf["events"].open_perf_buffer(callback)
    sigint_handler = SigintHandler()
    print("Running, press Ctrl+C to stop...", file=sys.stderr)
    while True:
        try:
            bpf.perf_buffer_poll()
        except KeyboardInterrupt:
            sys.exit()
        for pid in pids:
            try:
                exe = os.readlink(f"/proc/{pid}/exe")
                with open(f"/proc/{pid}/cmdline", "rb") as cmdline_fp:
                    cmdline = cmdline_fp.read()
            except FileNotFoundError:
                print(f"Skipping nonexistent pid {pid}...", file=sys.stderr)
                continue
            cmdline_i = cmdline.index(b"\0")
            argv0, cmdline = cmdline[:cmdline_i], cmdline[cmdline_i:]
            if (
                os.path.basename(exe) != name
                and os.path.basename(argv0) != name.encode()
            ) or cmdline_needle not in cmdline:
                print_skip_non_matching(pid, "Python")
                os.kill(pid, signal.SIGCONT)
                continue
            tool_argv = args.tool.gen_argv(args, pid) + args.tool_args
            tool_cmdline = " ".join(shlex.quote(gdb_arg) for gdb_arg in tool_argv)
            print(f"Starting {tool_cmdline}...", file=sys.stderr)
            if args.async_:
                if "TMUX" in os.environ:
                    subprocess.check_call(
                        ["tmux", "new-window", "-n", str(pid), tool_cmdline]
                    )
                else:
                    subprocess.Popen(["x-terminal-emulator", "--"] + tool_argv)
                args.tool.resume(pid)
            else:
                sigint_handler.raise_keyboard_interrupt = False
                try:
                    p = subprocess.Popen(tool_argv)
                    args.tool.resume(pid)
                    p.wait()
                finally:
                    print(f"{args.tool.DISPLAY_NAME} exited.", file=sys.stderr)
                    args.tool.cleanup(pid)
                    sigint_handler.raise_keyboard_interrupt = True
        pids.clear()


if __name__ == "__main__":
    main()
