#include "./helpers_nfqueue.h"

// copied from https://git.netfilter.org/libnetfilter_queue/tree/examples/nf-queue.c
int setup_nf_queue() {
    char buf[MNL_SOCKET_BUFFER_SIZE];
    int ret = 1;
    nlmsghdr nlh;

    nlsock_queue = mnl_socket_open(NETLINK_NETFILTER);
    if (nlsock_queue == NULL) {
        perror("mnl_socket_open queue");
        exit(EXIT_FAILURE);
    }

    if (mnl_socket_bind(nlsock_queue, 0, MNL_SOCKET_AUTOPID) < 0) {
        perror("mnl_socket_bind");
        exit(EXIT_FAILURE);
    }

    queue_socket_portid = mnl_socket_get_portid(nlsock_queue);

    nlh = nfq_nlmsg_put(buf, NFQNL_MSG_CONFIG, QUEUE_NUM);
    nfq_nlmsg_cfg_put_cmd(nlh, AF_INET, NFQNL_CFG_CMD_BIND);

    if (mnl_socket_sendto(nlsock_queue, nlh, nlh->nlmsg_len) < 0) {
        perror("mnl_socket_send");
        exit(EXIT_FAILURE);
    }

    nlh = nfq_nlmsg_put(buf, NFQNL_MSG_CONFIG, QUEUE_NUM);
    nfq_nlmsg_cfg_put_params(nlh, NFQNL_COPY_PACKET, 0xffff);

    // enable NFQA_CFG_F_CONNTRACK
    mnl_attr_put_u32(nlh, NFQA_CFG_FLAGS, htonl(NFQA_CFG_F_GSO | NFQA_CFG_F_CONNTRACK));
    mnl_attr_put_u32(nlh, NFQA_CFG_MASK, htonl(NFQA_CFG_F_GSO | NFQA_CFG_F_CONNTRACK));

    if (mnl_socket_sendto(nlsock_queue, nlh, nlh->nlmsg_len) < 0) {
        perror("mnl_socket_send");
        exit(EXIT_FAILURE);
    }

    mnl_socket_setsockopt(nlsock_queue, NETLINK_NO_ENOBUFS, &ret, sizeof(int));

    return ret;
}

// Extract the packet ID from the Netfilter queue message.
int queue_get_id_cb(const struct nlmsghdr *nlh, void *data) {
    struct nfqnl_msg_packet_hdr *ph = NULL;
    struct nlattr *attr[NFQA_MAX + 1] = {};

    if (nfq_nlmsg_parse(nlh, attr) < 0) {
        perror("problems parsing");
        return MNL_CB_ERROR;
    }

    if (attr[NFQA_PACKET_HDR] == NULL) {
        fputs("metaheader not set\n", stderr);
        return MNL_CB_ERROR;
    }

    ph = mnl_attr_get_payload(attr[NFQA_PACKET_HDR]);

    return ntohl(ph->packet_id);
}

// copied from https://git.netfilter.org/libnetfilter_queue/tree/examples/nf-queue.c
void nfq_send_verdict(int queue_num, uint32_t id, int verdict) {
    char buf[MNL_SOCKET_BUFFER_SIZE];
    struct nlmsghdr *nlh;

    nlh = nfq_nlmsg_put(buf, NFQNL_MSG_VERDICT, queue_num);
    nfq_nlmsg_verdict_put(nlh, id, verdict);

    if (mnl_socket_sendto(nlsock_queue, nlh, nlh->nlmsg_len) < 0) {
        perror("mnl_socket_send");
        exit(EXIT_FAILURE);
    }
}

uint64_t queue_recv(void *cb_data) {
    char buf[MNL_SOCKET_BUFFER_SIZE];

    int numbytes = mnl_socket_recvfrom(nlsock_queue, buf, sizeof(buf));
    if (numbytes == -1) {
        perror("mnl_socket_recvfrom");
        return -1;
    }

    int ret = mnl_cb_run(buf, numbytes, 0, queue_socket_portid, cb_data, NULL);
    if (ret < 0) {
        perror("mnl_cb_run");
        return -1;
    }

    return ret;
}