#include "nfq.hpp"
#include <netinet/in.h>
#include <linux/netfilter.h>
#include <libnetfilter_queue/libnetfilter_queue.h>


struct nfq_handle* NFQ::h = NULL;
unsigned NFQ::h_refcount = 0;


NFQ::NFQ(unsigned qid, cb_t _qcb, void** _qdata): qcb(_qcb), qdata(_qdata) {
    qh = nfq_create_queue(h, qid, &cb, this);
    if (!qh) {
        LOG("cannot bind socket to queue '%u'", qid);
        return;
    }
    if (nfq_set_mode(qh, NFQNL_COPY_PACKET, 0xffff) < 0) {
        LOG("cannot set copy_packet mode");
        nfq_destroy_queue(qh);
        qh = NULL;
        return;
    }

    h_refcount++;
}


int NFQ::cb(struct nfq_q_handle* qh, struct nfgenmsg*, struct nfq_data* nfa, void* ctx) {
    NFQ* inst = (NFQ*)ctx;
    const int id = ntohl(nfq_get_msg_packet_hdr(nfa)->packet_id);

    unsigned char* data;
    const int len = nfq_get_payload(nfa, &data);

    result_t result = (result_t){ true, false, 0 };
    if (len > 0) {
        inst->qcb(data, len, result, *(inst->qdata));
    }

    if (!result.mark) {
        return nfq_set_verdict(qh, id, result.verdict? NF_ACCEPT: NF_DROP, result.changed? len: 0, result.changed? data: NULL);
    } else {
#if 1
        return nfq_set_verdict2(qh, id, result.verdict? NF_ACCEPT: NF_DROP, htonl(result.mark), result.changed? len: 0, result.changed? data: NULL);
#else
        // This function is deprecated since it is broken, its use is highly discouraged. Please, use nfq_set_verdict2 instead.
        return nfq_set_verdict_mark(qh, id, result.verdict? NF_ACCEPT: NF_DROP, htonl(result.mark), result.changed? len: 0, result.changed? data: NULL);
#endif
    }
}


NFQ::~NFQ() {
    if (qh) {
        nfq_destroy_queue(qh);
        if (!--h_refcount) {
            DBG("closing handle");
            nfq_close(h);
            h = NULL;
        }
    }
}


NFQ* NFQ::getQueue(unsigned proto, unsigned qid, cb_t cb, void** cbctx) {
    if (proto != AF_INET) {
        return NULL; // TODO: maintain multiple handles
    }

    if (!h) {
        h = nfq_open();
        if (!h) {
            LOG("cannot open library handle");
            return NULL;
        }
        if (nfq_unbind_pf(h, proto) < 0) {
            LOG("cannot unbind existing nf_queue handler for AF_INET (if any)");
            nfq_close(h);
            h = NULL;
            return NULL;
        }
        if (nfq_bind_pf(h, proto) < 0) {
            LOG("cannot bind nfnetlink_queue as nf_queue handler for AF_INET");
            nfq_close(h);
            h = NULL;
            return NULL;
        }
    }

    NFQ* inst = new NFQ(qid, cb, cbctx);
    if (!inst->qh) {
        delete inst;
        if (!h_refcount) {
            nfq_close(h);
            h = NULL;
        }
        return NULL;
    }
    return inst;
}


bool NFQ::loop(volatile int& stop) {
    if (!h_refcount) {
        return false;
    }
    DBG("looping on %u queues...", h_refcount);

    int rv;
    int fd = nfq_fd(h);
    static char buf[4096] __attribute__ ((aligned));

    while (!stop) {
        if ((rv = recv(fd, buf, sizeof(buf), 0)) >= 0) {
            update_now(); // TODO: needed for each packet?
            if (nfq_handle_packet(h, buf, rv) != 0) {
                LOG("cannot handle packet");
                return false;
            }
        } else if (errno == ENOBUFS) {
            LOG("losing packets!");
        } else if (errno == EINTR) {
            break;
        } else {
            LOG_ERRNO("recv failed");
            return false;
        }
    }

    return true;
}