// ----------------------------------------------------------------------------------;
// michael.mill@metacom.net --- (07/09/2018)
//
// install : libipset-dev

#include <ipset_t_header.h>
#include <unistd.h>
#include <assert.h>
#include <libipset/types.h>
#include <libipset/session.h>
#include <libipset/data.h>
#include <stdint.h>

// #include <libnetfilter_conntrack/libnetfilter_conntrack.h>

static struct ipset_session *session;

static bool try_ipset_cmd (enum ipset_cmd cmd, 
                           const char    *setname, 
                           const struct   in_addr *addr, 
                           uint32_t       timeout) {

    const struct ipset_type *type;
    uint8_t family;
    int r;

    r = ipset_session_data_set(session, IPSET_SETNAME, setname);

    // since the IPSET_SETNAME option is valid, this should never fail 
    assert(r == 0);

    type = ipset_type_get(session, cmd);

    if (type == NULL) {
    // possible reasons for failure: set name does not exist
        return false;
    }

    family = NFPROTO_IPV4;
    ipset_session_data_set(session, IPSET_OPT_FAMILY, &family);
    ipset_session_data_set(session, IPSET_OPT_IP, addr);

    if (timeout) {
        ipset_session_data_set(session, IPSET_OPT_TIMEOUT, &timeout);
    }

    r = ipset_cmd(session, cmd, 0);
    // assume that errors always occur if NOT in set. To do it otherwise,
    // see lib/session.c for IPSET_CMD_TEST in ipset_cmd

    return r == 0;
}

static bool try_ipset_create(const char *setname, 
                             const char *_typename) {

    const struct ipset_type *type;
    uint32_t timeout;
    uint8_t family;
    int r;

    r = ipset_session_data_set(session, IPSET_SETNAME, setname);

    /* since the IPSET_SETNAME option is valid, this should never fail */
    assert(r == 0);

    ipset_session_data_set(session, IPSET_OPT_TYPENAME, _typename);

    type = ipset_type_get(session, IPSET_CMD_CREATE);

    if (type == NULL) {
        return false;
    }

    timeout = 0; /* timeout support, but default to infinity */
    ipset_session_data_set(session, IPSET_OPT_TIMEOUT, &timeout);
    ipset_session_data_set(session, IPSET_OPT_TYPE, type);
    family = NFPROTO_IPV4;
    ipset_session_data_set(session, IPSET_OPT_FAMILY, &family);

    r = ipset_cmd(session, IPSET_CMD_CREATE, /*lineno*/ 0);
    return r == 0;
}

static bool has_ipset_setname(const char *setname) {
    ipset_session_data_set(session, IPSET_SETNAME, setname);
    return ipset_cmd(session, IPSET_CMD_HEADER, 0) == 0;
}

void do_block(const struct in_addr addr) {
    try_ipset_cmd(IPSET_CMD_ADD, SETNAME_BLACKLIST, &addr, BLOCK_TIME);
}

void do_unblock(const struct in_addr addr) {
    try_ipset_cmd(IPSET_CMD_DEL, SETNAME_BLACKLIST, &addr, 0);
}

void do_whitelist(const struct in_addr addr) {
    try_ipset_cmd(IPSET_CMD_ADD, SETNAME_WHITELIST, &addr, WHITELIST_TIME);
}

bool is_whitelisted(const struct in_addr addr) {
    return try_ipset_cmd(IPSET_CMD_TEST, SETNAME_WHITELIST, &addr, 0);
}

bool blocker_init(void) {

    ipset_load_types();

    session = ipset_session_init(printf);

    if (!session) {
        fprintf(stderr, "Cannot initialize ipset session.\n");
        return false;
    }

    /* return success on attempting to add an existing / remove an
     * non-existing rule */
    ipset_envopt_parse(session, IPSET_ENV_EXIST, NULL);

    return true;
}

void blocker_fini(void) {

    ipset_session_fini(session);
}
