/*
 * ----------------------------------------------------------------------------
 * "THE BEER-WARE LICENSE" (Revision 42):
 * David Bouman (pql) wrote this file.  As long as you retain this notice you
 * can do whatever you want with this stuff. If we meet some day, and you think
 * this stuff is worth it, you can buy me a beer in return.   Signed, David.
 * ----------------------------------------------------------------------------
 */
#define _GNU_SOURCE
#include <stdlib.h>
#include <time.h>
#include <string.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#include <sys/types.h>
#include <sched.h>
#include <signal.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/rule.h>
#include <libnftnl/expr.h>
#include "helpers.h"

static uint64_t default_batch_req_handler(struct mnl_socket* nl, int portid, int table_seq)
{
    char buf[MNL_SOCKET_BUFFER_SIZE];

    int ret = mnl_socket_recvfrom(nl, buf, sizeof(buf));

    while (ret > 0) {
        ret = mnl_cb_run(buf, ret, table_seq, portid, NULL, NULL);
        if (ret <= 0) break;
        ret = mnl_socket_recvfrom(nl, buf, sizeof(buf));
    }
    return ret;
}

int64_t send_batch_request(struct mnl_socket* nl, uint16_t msg, uint16_t msg_flags, uint16_t family, void** object, int* seq, uint64_t (*result_handler)(struct mnl_socket*, int, int))
{
    
    char buf[MNL_SOCKET_BUFFER_SIZE];
    struct mnl_nlmsg_batch* batch = mnl_nlmsg_batch_start(buf, sizeof buf);

    uint8_t msg_type = msg & 0xff;
    uint8_t nft_type = (msg >> 8) & 0xff;
    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), (*seq)++);
    mnl_nlmsg_batch_next(batch);
    int table_seq = *seq;
    struct nlmsghdr* nlh;

    if (result_handler == NULL) {
        result_handler = default_batch_req_handler;
    }
   
    nlh = nftnl_nlmsg_build_hdr(
        mnl_nlmsg_batch_current(batch),
        msg_type, family,
        msg_flags | NLM_F_ACK, (*seq)++
    );
    
    switch(nft_type) {
        case NFT_TYPE_TABLE:
            nftnl_table_nlmsg_build_payload(nlh, *object);
            nftnl_table_free(*object);
            break;
        case NFT_TYPE_CHAIN:
            nftnl_chain_nlmsg_build_payload(nlh, *object);
            nftnl_chain_free(*object);
            break;
        case NFT_TYPE_RULE:
            nftnl_rule_nlmsg_build_payload(nlh, *object);
            nftnl_rule_free(*object);
            break;
        default:
            return -1; // will increment seq wrongly... no prob i guess
    }  

    *object = NULL;

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

    int ret = mnl_socket_sendto(
        nl,
        mnl_nlmsg_batch_head(batch),
        mnl_nlmsg_batch_size(batch)
    );

    if (ret < 0) {
        perror("mnl_socket_send");
        return -1;
    }

    int portid = mnl_socket_get_portid(nl);

    mnl_nlmsg_batch_stop(batch);

    result_handler(nl, portid, table_seq);
}

struct nftnl_table* build_table(char* name, uint16_t family)
{
    struct nftnl_table* t = nftnl_table_alloc();
    
    nftnl_table_set_u32(t, NFTNL_TABLE_FAMILY, family);
    nftnl_table_set_str(t, NFTNL_TABLE_NAME, name);

    return t;
}

struct nftnl_chain* build_chain(char* table_name, char* chain_name, struct unft_base_chain_param* base_param)
{
    struct nftnl_chain* c;

    c = nftnl_chain_alloc();

    nftnl_chain_set_str(c, NFTNL_CHAIN_NAME, chain_name);
    nftnl_chain_set_str(c, NFTNL_CHAIN_TABLE, table_name);

    if (base_param) {
        nftnl_chain_set_u32(c, NFTNL_CHAIN_HOOKNUM, base_param->hook_num);
        nftnl_chain_set_u32(c, NFTNL_CHAIN_PRIO, base_param->prio);
    }

    return c;
    
}


struct nftnl_rule* build_rule(char* table_name, char* chain_name, uint16_t family, uint64_t* handle)
{
    struct nftnl_rule* r = NULL;
    uint8_t proto;
    
    r = nftnl_rule_alloc();

    nftnl_rule_set_str(r, NFTNL_RULE_TABLE, table_name);
    nftnl_rule_set_str(r, NFTNL_RULE_CHAIN, chain_name);
    nftnl_rule_set_u32(r, NFTNL_RULE_FAMILY, family);
    
    if (handle) {
        nftnl_rule_set_u64(r, NFTNL_RULE_POSITION, *handle);
    }

    return r;

}

// for some reason my editor does not recognize these
#define NFTA_BITWISE_OP NFTA_BITWISE_XOR + 1
#define NFTA_BITWISE_DATA NFTA_BITWISE_OP + 1


void rule_add_bit_shift(
    struct nftnl_rule* r, uint32_t shift_type, uint32_t bitwise_len,
    uint32_t bitwise_sreg, uint32_t bitwise_dreg, void* data, uint32_t data_len)
{
    
    if(bitwise_len > 0xff) {
        puts("bitwise_len > 0xff");
        exit(EXIT_FAILURE);
    }

    struct nftnl_expr* e;
    e = nftnl_expr_alloc("bitwise");

    nftnl_expr_set_u32(e, NFTA_BITWISE_SREG, bitwise_sreg);
    nftnl_expr_set_u32(e, NFTA_BITWISE_DREG, bitwise_dreg);
    nftnl_expr_set_u32(e, NFTA_BITWISE_OP, shift_type);
    nftnl_expr_set_u32(e, NFTA_BITWISE_LEN, bitwise_len);
    nftnl_expr_set_data(e, NFTA_BITWISE_DATA, data, data_len);

    nftnl_rule_add_expr(r, e);
}

void rule_add_memcpy(struct nftnl_rule* r, uint32_t len, uint32_t sreg, uint32_t dreg)
{
    uint32_t data = 0;
    rule_add_bit_shift(r, NFT_BITWISE_LSHIFT, len, sreg, dreg, &data, sizeof(data));
}

void rule_add_payload(struct nftnl_rule* r, uint32_t base, uint32_t offset, uint32_t len, uint32_t dreg)
{
    struct nftnl_expr* e;
    e = nftnl_expr_alloc("payload");

    nftnl_expr_set_u32(e, NFTNL_EXPR_PAYLOAD_BASE, base);
    nftnl_expr_set_u32(e, NFTNL_EXPR_PAYLOAD_OFFSET, offset);
    nftnl_expr_set_u32(e, NFTNL_EXPR_PAYLOAD_LEN, len);
    nftnl_expr_set_u32(e, NFTNL_EXPR_PAYLOAD_DREG, dreg);

    nftnl_rule_add_expr(r, e);
}

void rule_add_cmp(struct nftnl_rule* r, uint32_t op, uint32_t sreg, void* data, size_t data_len)
{
    struct nftnl_expr* e;
    e = nftnl_expr_alloc("cmp");

    nftnl_expr_set_u32(e, NFTA_CMP_OP, op);
    nftnl_expr_set_u32(e, NFTA_CMP_SREG, sreg);
    nftnl_expr_set_data(e, NFTA_CMP_DATA, data, data_len);

    nftnl_rule_add_expr(r, e);
}

void rule_add_immediate_data(struct nftnl_rule* r, uint32_t dreg, void* data, size_t data_len)
{
    struct nftnl_expr* e;
    
    e = nftnl_expr_alloc("immediate");

    nftnl_expr_set_u32(e, NFTA_IMMEDIATE_DREG, dreg);
    nftnl_expr_set_data(e, NFTA_IMMEDIATE_DATA, data, data_len);

    nftnl_rule_add_expr(r, e);
}

void rule_add_immediate_verdict(struct nftnl_rule* r, uint32_t verdict, char* chain_name)
{
    struct nftnl_expr* e;
    e = nftnl_expr_alloc("immediate");

    // dreg = 0 -> verdict
    nftnl_expr_set_u32(e, NFTA_IMMEDIATE_DREG, 0); 

    nftnl_expr_set_u32(e, NFTNL_EXPR_IMM_VERDICT, verdict);

    if (verdict == NFT_GOTO || verdict == NFT_JUMP) {
        nftnl_expr_set_str(e, NFTNL_EXPR_IMM_CHAIN, chain_name);
    }

    nftnl_rule_add_expr(r, e);
}


int create_table(struct mnl_socket* nl, char* name, uint16_t family, int* seq, uint64_t (*result_handler)(struct mnl_socket*, int, int))
{
    
    struct nftnl_table* t = build_table(name, family);

    return send_batch_request(
        nl,
        NFT_MSG_NEWTABLE | (NFT_TYPE_TABLE << 8),
        NLM_F_CREATE, family, (void**)&t, seq,
        result_handler
    );
}

int create_chain(struct mnl_socket* nl, char* chain_name, char* table_name, uint16_t family, struct unft_base_chain_param* base_param, int* seq, uint64_t (*result_handler)(struct mnl_socket*, int, int))
{
    struct nftnl_chain* c = build_chain(chain_name, table_name, base_param);

    return send_batch_request(
        nl,
        NFT_MSG_NEWCHAIN | (NFT_TYPE_CHAIN << 8),
        NLM_F_CREATE, family, (void**)&c, seq,
        result_handler  
    );
}

/*
int update_chain(struct mnl_socket* nl, char* chain_name, char* table_name, uint16_t family, struct unft_base_chain_param* base_param, int* seq, uint64_t (*result_handler)(struct mnl_socket*, int, int))
{
    struct nftnl_chain* c = build_chain(chain_name, table_name, base_param);

    return send_batch_request(
        nl,
        NFT_MSG_NEWCHAIN | (NFT_TYPE_CHAIN << 8),
        NLM_F_CREATE | NLM_F_REPLACE, family, (void**)&c, seq,
        result_handler
    ;)
}
*/

struct child_proc {
    struct child_proc* next;
    pid_t pid;
};

static struct child_proc *children;


static void add_child(pid_t pid)
{
    struct child_proc* child = malloc(sizeof *child);
    child->pid = pid;
    child->next = children;
    children = child;
}

static void kill_children(int sig)
{
    //printf("[pid=%d] killing children!\n", getpid());

    struct child_proc* current_child = children;
    while (current_child) {
        kill(current_child->pid, SIGTERM);
        current_child = current_child->next;
    }

    exit(EXIT_SUCCESS);
}

pid_t setup_listener(char* ip_string, uint16_t port, int (*handler)(int))
{
    
    int err;

    int s = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);

    if (s < 0) {
        perror("socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)");
        exit(EXIT_FAILURE);
    }

    int reuse_addr = 1;

    setsockopt(s, SOL_SOCKET, SO_REUSEADDR, &reuse_addr, sizeof reuse_addr);

    struct sockaddr_in addr;
    inet_aton(ip_string, &addr.sin_addr);
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);

    err = bind(s, (struct sockaddr*)&addr, sizeof(addr));
    
    if (err < 0) {
        perror("bind");
        exit(EXIT_FAILURE);
    }

    printf("Started listener on [%s:%d] (udp)\n", ip_string, port);

    pid_t pid = fork();
    if (pid) { 
        // parent process
        add_child(pid);   
        return pid;
    }

    handler(s);

    exit(EXIT_SUCCESS);

}

int stop_listener(pid_t pid)
{
    
    if (kill(pid, SIGTERM)) {
        perror("kill");
        return -1;
    };

    struct child_proc* next_child = children;
    struct child_proc* current_child = NULL;

    while (next_child) {    
        
        if (next_child->pid == pid) {
            
            struct child_proc** prev = current_child == NULL ? &children : &current_child;
            if (current_child == NULL) {
                prev = &children;
            } else {
                prev = &current_child;
            }

            (*prev)->next = next_child->next;
            break;

        }

        current_child = next_child;
        next_child = next_child->next;
    }

    return 0;
}

int connect_to(char* ip_string, uint16_t port)
{
    int s = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);

    if (s < 0) {
        perror("socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)");
        return -1;
    }
    struct sockaddr_in conn_addr;
    conn_addr.sin_port = htons(port);
    inet_aton(ip_string, &conn_addr.sin_addr);
    conn_addr.sin_family = AF_INET;

    int err = connect(s, (struct sockaddr*)&conn_addr, sizeof conn_addr);
    if (err < 0) {
        perror("connect");
        return -1;
    }
    
    printf("Successfully connected to [%s:%hd] (udp)\n", ip_string, port);

    return s;
}

void hexdump(void* data, size_t len, unsigned int n_columns)
{

    uint8_t* bdata = data;

    for (int i = 0; i < len; ++i) {
        printf("%.2hhx ", bdata[i]);

        if ( (i+1) % n_columns == 0) {
            putchar('\n');
        }
    }
}