#include <stdio.h>
#include <poll.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/select.h>
#include <signal.h>
#include <linux/netlink.h>
#include <linux/connector.h>
#include <linux/cn_proc.h>
#include <unistd.h>
#include <errno.h>
#include <time.h>

typedef struct __attribute__((aligned(NLMSG_ALIGNTO)))
{
    struct nlmsghdr nl_hdr;
    struct __attribute__((__packed__))
    {
        struct cn_msg cn_msg;
        enum proc_cn_mcast_op cn_mcast;
    };
} register_msg_t;

typedef struct __attribute__((aligned(NLMSG_ALIGNTO)))
{
    struct nlmsghdr nl_hdr;
    struct __attribute__((__packed__))
    {
        struct cn_msg cn_msg;
        struct proc_event proc_ev;
    };
} event_msg_t;

event_msg_t proc_msg;

void Now()
{
    struct timespec ts;
    struct tm *tm_info;
    char buffer[64];

    if (clock_gettime(CLOCK_REALTIME, &ts) == -1)
    {
        perror("clock_gettime");
        return;
    }

    tm_info = localtime(&ts.tv_sec);
    strftime(buffer, sizeof(buffer), "%Y-%m-%d %H:%M:%S", tm_info);
    printf("Localtime %s.%03ld ", buffer, ts.tv_nsec / 1000000);
}

void printEvent()
{
    // union unnamed *procEvent = &proc_msg.proc_ev.event_data;
    switch (proc_msg.proc_ev.what)
    {
    case PROC_EVENT_FORK:
        Now();
        printf("Fork\t%6d\t%6d\t%6d\t%6d\n", proc_msg.proc_ev.event_data.fork.parent_pid, proc_msg.proc_ev.event_data.fork.parent_tgid, proc_msg.proc_ev.event_data.fork.child_pid, proc_msg.proc_ev.event_data.fork.child_tgid);
        break;
    case PROC_EVENT_EXIT:
        Now();
        printf("Exit\t%6d\t%6d\t%6d\t%6d\n", proc_msg.proc_ev.event_data.exit.process_pid, proc_msg.proc_ev.event_data.exit.process_tgid, proc_msg.proc_ev.event_data.exit.exit_code, proc_msg.proc_ev.event_data.exit.exit_signal);
        break;
    case PROC_EVENT_EXEC:
    default:
        break;
    }
}

int main()
{
    int s = socket(PF_NETLINK, SOCK_DGRAM, NETLINK_CONNECTOR);

    register_msg_t nlcn_msg;
    struct sockaddr_nl l_local;
    l_local.nl_family = AF_NETLINK;
    l_local.nl_groups = 12345;
    l_local.nl_pid = 0;

    if (bind(s, (struct sockaddr *)&l_local, sizeof(struct sockaddr_nl)) == -1)
    {
        perror(bind);
        close(s);
        return -1;
    }

    // int on = l_local.nl_groups;
    // setsockopt(s,270,1,&on,sizeof(on));
    memset(&nlcn_msg, 0, sizeof(nlcn_msg));
    nlcn_msg.nl_hdr.nlmsg_len = sizeof(nlcn_msg);
    nlcn_msg.nl_hdr.nlmsg_pid = getpid();
    nlcn_msg.nl_hdr.nlmsg_type = NLMSG_DONE;

    nlcn_msg.cn_msg.id.idx = CN_IDX_PROC;
    nlcn_msg.cn_msg.id.val = CN_VAL_PROC;
    nlcn_msg.cn_msg.len = sizeof(enum proc_cn_mcast_op);

    nlcn_msg.cn_mcast = PROC_CN_MCAST_LISTEN;

    if (send(s, &nlcn_msg, sizeof(nlcn_msg), 0) == -1)
    {
        perror("can't register to netlink");
        close(s);
        return -1;
    }

    // 震惊，拿到socket了，开听！
    printf("Hello, kernel-connector!\n");
    // fd_set readfds;
    // struct timeval tv = {
    //     .tv_sec = 5,
    //     .tv_usec = 0};
    struct pollfd fds;

    fds.fd = s;
    fds.events = POLLIN;
    int rc;

    while (1)
    {
        // FD_ZERO(&readfds);
        // FD_SET(s, &readfds);

        // int rc = select(s + 1, &readfds, NULL, NULL, &tv);
        rc = poll(&fds, 1, 5000);

        if (rc == -1)
        {
            if (errno == EINTR)
            {
                continue;
            }
            fprintf(stderr, "Failed to listen to netlink socket: %s\n", strerror(errno));
            return -1;
        }
        else if (rc == 0)
        {
            printf("No message in 5s...\n");
        }
        else
        {
            rc = recv(s, &proc_msg, sizeof(proc_msg), 0);
            if (rc == -1)
            {
                if (errno == EINTR)
                {
                    continue;
                }
                fprintf(stderr, "Failed to listen to netlink socket: %s\n", strerror(errno));
            }
            else
            {
                printEvent();
            }
        }
    }
}