# SPDX-License-Identifier: GPL-3.0-or-later
"""eBPF source for security-relevant syscall telemetry."""
from __future__ import annotations

from collections.abc import Callable

from .bcc_source import _decode, available
from .syscall_event import SyscallEvent

_BPF_PROGRAM = r"""
#include <uapi/linux/ptrace.h>
#include <linux/sched.h>

#define TEXTLEN 80
#define SYS_MPROTECT 1
#define SYS_MMAP 2
#define SYS_MEMFD_CREATE 3
#define SYS_PTRACE 4
#define SYS_PRCTL 5
#define SYS_SECCOMP 6
#define SYS_PROCESS_VM_READV 7
#define SYS_PROCESS_VM_WRITEV 8
#define SYS_MLOCK 9
#define SYS_MLOCK2 10
#define SYS_MLOCKALL 11

struct data_t {
    u32 pid;
    u32 ppid;
    u32 uid;
    char comm[TASK_COMM_LEN];
    u32 syscall_id;
    u64 arg0;
    u64 arg1;
    u64 arg2;
    u64 arg3;
    u64 arg4;
    u64 arg5;
    char text[TEXTLEN];
};
BPF_PERF_OUTPUT(events);

static int submit(struct pt_regs *ctx, u32 syscall_id,
                  u64 a0, u64 a1, u64 a2, u64 a3, u64 a4, u64 a5,
                  const char __user *textp)
{
    struct data_t data = {};
    struct task_struct *task = (struct task_struct *)bpf_get_current_task();
    data.pid = bpf_get_current_pid_tgid() >> 32;
    data.ppid = task->real_parent->tgid;
    data.uid = bpf_get_current_uid_gid() & 0xffffffff;
    bpf_get_current_comm(&data.comm, sizeof(data.comm));
    data.syscall_id = syscall_id;
    data.arg0 = a0;
    data.arg1 = a1;
    data.arg2 = a2;
    data.arg3 = a3;
    data.arg4 = a4;
    data.arg5 = a5;
    if (textp)
        bpf_probe_read_user_str(&data.text, sizeof(data.text), textp);
    events.perf_submit(ctx, &data, sizeof(data));
    return 0;
}

int syscall__mprotect(struct pt_regs *ctx, unsigned long start, unsigned long len,
                      unsigned long prot)
{
    return submit(ctx, SYS_MPROTECT, start, len, prot, 0, 0, 0, 0);
}

int syscall__mmap(struct pt_regs *ctx, unsigned long addr, unsigned long len,
                  unsigned long prot, unsigned long flags, unsigned long fd,
                  unsigned long off)
{
    return submit(ctx, SYS_MMAP, addr, len, prot, flags, fd, off, 0);
}

int syscall__memfd_create(struct pt_regs *ctx, const char __user *name,
                          unsigned int flags)
{
    return submit(ctx, SYS_MEMFD_CREATE, flags, 0, 0, 0, 0, 0, name);
}

int syscall__ptrace(struct pt_regs *ctx, long request, long pid,
                    unsigned long addr, unsigned long data)
{
    return submit(ctx, SYS_PTRACE, request, pid, addr, data, 0, 0, 0);
}

int syscall__prctl(struct pt_regs *ctx, int option, unsigned long arg2,
                   unsigned long arg3, unsigned long arg4, unsigned long arg5)
{
    return submit(ctx, SYS_PRCTL, option, arg2, arg3, arg4, arg5, 0, 0);
}

int syscall__seccomp(struct pt_regs *ctx, unsigned int operation,
                     unsigned int flags, const void __user *args)
{
    return submit(ctx, SYS_SECCOMP, operation, flags, (u64)args, 0, 0, 0, 0);
}

int syscall__process_vm_readv(struct pt_regs *ctx, int pid,
                              const void __user *lvec,
                              unsigned long liovcnt,
                              const void __user *rvec,
                              unsigned long riovcnt,
                              unsigned long flags)
{
    return submit(ctx, SYS_PROCESS_VM_READV, pid, liovcnt, riovcnt, flags, 0, 0, 0);
}

int syscall__process_vm_writev(struct pt_regs *ctx, int pid,
                               const void __user *lvec,
                               unsigned long liovcnt,
                               const void __user *rvec,
                               unsigned long riovcnt,
                               unsigned long flags)
{
    return submit(ctx, SYS_PROCESS_VM_WRITEV, pid, liovcnt, riovcnt, flags, 0, 0, 0);
}

int syscall__mlock(struct pt_regs *ctx, const void __user *addr, unsigned long len)
{
    return submit(ctx, SYS_MLOCK, (u64)addr, len, 0, 0, 0, 0, 0);
}

int syscall__mlock2(struct pt_regs *ctx, const void __user *addr, unsigned long len,
                    int flags)
{
    return submit(ctx, SYS_MLOCK2, (u64)addr, len, flags, 0, 0, 0, 0);
}

int syscall__mlockall(struct pt_regs *ctx, int flags)
{
    return submit(ctx, SYS_MLOCKALL, flags, 0, 0, 0, 0, 0, 0);
}
"""

_PROBES = {
    "mprotect": "syscall__mprotect",
    "mmap": "syscall__mmap",
    "memfd_create": "syscall__memfd_create",
    "ptrace": "syscall__ptrace",
    "prctl": "syscall__prctl",
    "seccomp": "syscall__seccomp",
    "process_vm_readv": "syscall__process_vm_readv",
    "process_vm_writev": "syscall__process_vm_writev",
    "mlock": "syscall__mlock",
    "mlock2": "syscall__mlock2",
    "mlockall": "syscall__mlockall",
}

_SYSCALL_NAMES = {
    1: "mprotect",
    2: "mmap",
    3: "memfd_create",
    4: "ptrace",
    5: "prctl",
    6: "seccomp",
    7: "process_vm_readv",
    8: "process_vm_writev",
    9: "mlock",
    10: "mlock2",
    11: "mlockall",
}


class BccSyscallSource:
    def __init__(self, on_event: Callable[[SyscallEvent], None]) -> None:
        self._on_event = on_event
        self._bpf = None
        self._running = False
        self.attach_errors: list[str] = []

    def start(self) -> None:
        from bcc import BPF

        self._bpf = BPF(text=_BPF_PROGRAM)
        attached = 0
        for syscall, fn_name in _PROBES.items():
            try:
                event = self._bpf.get_syscall_fnname(syscall)
                self._bpf.attach_kprobe(event=event, fn_name=fn_name)
                attached += 1
            except Exception as exc:
                self.attach_errors.append(f"{syscall}: {exc!r}")
        if attached == 0:
            raise RuntimeError("no syscall probes attached")
        self._bpf["events"].open_perf_buffer(self._handle, page_cnt=64)
        self._running = True

    def poll(self, timeout_ms: int = 200) -> None:
        if self._bpf is not None:
            self._bpf.perf_buffer_poll(timeout=timeout_ms)

    def stop(self) -> None:
        self._running = False
        if self._bpf is not None:
            try:
                self._bpf.cleanup()
            except Exception:
                pass
            self._bpf = None

    @property
    def running(self) -> bool:
        return self._running

    def _handle(self, cpu, data, size) -> None:
        event = self._bpf["events"].event(data)
        ev = SyscallEvent(
            pid=event.pid,
            ppid=event.ppid,
            uid=event.uid,
            comm=_decode(event.comm),
            syscall=_SYSCALL_NAMES.get(int(event.syscall_id), "unknown"),
            arg0=int(event.arg0),
            arg1=int(event.arg1),
            arg2=int(event.arg2),
            arg3=int(event.arg3),
            arg4=int(event.arg4),
            arg5=int(event.arg5),
            text=_decode(event.text),
        )
        try:
            self._on_event(ev)
        except Exception:
            pass