#!/usr/bin/python3

# SPDX-License-Identifier: GPL-3.0-or-later
# Copyright (c) 2022 Red Hat, Inc.

import collections
import ctypes
import os
import struct
import sys

import bcc


LIBS = ('libcrypto.so.3', 'libssl.so.3')
LIBDIR = '/usr/lib64' if os.path.exists('/usr/lib64') else '/usr/lib'
NIDS = {
    64:  'sha1',
    114: 'md5_sha1',
    115: 'sha1WithRSA',
}

# https://github.com/iovisor/bcc/issues/3240
CODE_HEADER = r'''
#include <linux/sched.h>

#pragma pack(1)
struct event {
    unsigned pid;
    int nid;  // int EVP_MD_get_type(const EVP_MD *md);
    char process_name[64];
    char probe_name[64];
};
BPF_PERF_OUTPUT(events);

static void _strncpy(char* dst, char* src, unsigned n) {
    char* dst_start = dst;
    while (*src && dst - dst_start < n)
        *dst++ = *src++;
    dst_start[n - 1] = '\0';
}
'''

CODE_TEMPLATE = r'''
int trace_<PROBENAME>(struct pt_regs *ctx) {
    struct event e = {};
    e.pid = bpf_get_current_pid_tgid();
    bpf_usdt_readarg(1, ctx, &e.nid);
    _strncpy(e.probe_name, "<PROBENAME>", 64);
    char comm[TASK_COMM_LEN];
    bpf_get_current_comm(&comm, sizeof(comm));
    _strncpy(e.process_name, comm, TASK_COMM_LEN > 64 ? 64 : TASK_COMM_LEN);
    events.perf_submit(ctx, &e, sizeof(e));
    return 0;
}
'''

process_eventset_map = collections.defaultdict(set)


class Event(ctypes.Structure):
    _pack_ = 1
    _fields_ = [
        ("pid", ctypes.c_uint),
        ("nid", ctypes.c_int),
        ("process_name", ctypes.c_char * 64),
        ("probe_name", ctypes.c_char * 64),
    ]


def setup_tracing(libname):
    usdt = bcc.USDT(path=os.path.join(LIBDIR, libname))
    code_fragments = []
    for probe in sorted(usdt.enumerate_probes(), key=lambda p: p.name):
        probe_name = probe.name.decode()
        if not probe_name.startswith('fedora_'):
            print(f'{libname} probe ignored: {probe_name}',
                  file=sys.stderr, flush=True)
        usdt.enable_probe(probe=probe_name, fn_name=f'trace_{probe_name}')
        print(f'{libname} probe attached: {probe_name}',
              file=sys.stderr, flush=True)
        code_fragments.append(CODE_TEMPLATE.replace('<PROBENAME>', probe_name))
    return usdt, code_fragments


def trace_event(_unused_cpu, data, _unused_size):
    global process_eventset_map
    e = Event.from_address(data)
    nid = NIDS.get(e.nid, e.nid)
    print(f'{e.process_name.decode()}[{e.pid}] '
          f'uses mdnid={nid} in {e.probe_name.decode()}', flush=True)
    process_eventset_map[(e.process_name, e.pid)].add((e.probe_name, nid))


def main():
    print('# attaching...', file=sys.stderr, flush=True)
    usdt_contexts, code_fragments = [], []
    for libname in LIBS:
        usdt, _code_fragments = setup_tracing(libname)
        usdt_contexts.append(usdt)
        code_fragments.extend(_code_fragments)
    if not code_fragments:
        print('could not attach any probes, is your library instrumented?',
              file=sys.stderr, flush=True)
        sys.exit(1)
    code = CODE_HEADER + '\n'.join(code_fragments)
    b = bcc.BPF(text=code, usdt_contexts=usdt_contexts)
    b["events"].open_perf_buffer(trace_event)

    print('# tracing...', file=sys.stderr, flush=True)
    try:
        while True:
            b.perf_buffer_poll()
    except KeyboardInterrupt:
        pass

    print('# summarizing...', file=sys.stderr, flush=True)
    events_processes_map = {
            tuple(sorted(eventset)): sorted(p for p, es
                                            in process_eventset_map.items()
                                            if es == eventset)
            for eventset in process_eventset_map.values()
    }
    for eventset, processes in events_processes_map.items():
        process_names = sorted(set((pname.decode() for pname, _ in processes)))
        if len(process_names) == 1:
            print(f'`{process_names[0]}` has triggered the following probes:',
                  file=sys.stderr, flush=True)
        else:
            print(f'{", ".join(f"`{p}`" for p in process_names)} \n'
                  'have triggered the following probes:',
                  file=sys.stderr, flush=True)
        for probe_name, nid in eventset:
            print(f'* {probe_name.decode()} mdnid={nid}',
                  file=sys.stderr, flush=True)
    print('# done', file=sys.stderr, flush=True)


if __name__ == '__main__':
    main()
