#include "libiptc.h"

int x_get_rule(const char *chain, struct xtc_handle *handle, char *ipv4)
{
    int r = 1;
    const struct ipt_entry *entry;
    struct ipt_entry_match *entry_match;

    for (entry = iptc_first_rule(chain, handle); entry; entry = iptc_next_rule(entry, handle)) {
        const char *t = iptc_get_target(entry, handle);
        entry_match = (struct ipt_entry_match *)entry->elems;
        //printf("u.user.name: %s\n", entry_match->u.user.name);
        char addr[64];
        memset(addr, 0, 64);
        inet_ntop(AF_INET, &(entry->ip.src), addr, sizeof(addr));
        //printf("%s\n", addr);
        //printf("%s\n", t);
        if (0 == strcmp(ipv4, addr) && 0 == strcmp(t, "DROP") && 0 == strcmp(entry_match->u.user.name, "tcp")) {
            r = 0;
            break;
        }
    }

    return r;

}

int show_all_rule(char *ipv4)
{
    int r = 0;
    struct xtc_handle *handle;
    const char *chain = NULL;
    struct ipt_counters counter;

    handle = iptc_init("filter");
    
    if ( !handle ) {
        printf("Error initializing: %s\n", iptc_strerror(errno));
        exit(errno);
    }

    for (chain = iptc_first_chain(handle); chain; chain = iptc_next_chain(handle)) {
        if (chain != NULL && handle != NULL) {
            iptc_get_policy(chain, &counter, handle);
        }

        r = x_get_rule(chain, handle, ipv4);
        if (r == 0) {
            break;
        }
    }
    
    iptc_free(handle);

    return r;
}

static u_int16_t parse_port(const char *port)
{
    return atoi(port);
}

static void parse_ports(const char *portstring, u_int16_t * ports)
{
    char *buffer;
    char *cp;

    buffer = strdup(portstring);
    if ((cp = strchr(buffer, ':')) == NULL)
        ports[0] = ports[1] = parse_port(buffer);
    else {
        *cp = '\0';
        cp++;

        ports[0] = buffer[0] ? parse_port(buffer) : 0;
        ports[1] = cp[0] ? parse_port(cp) : 0xFFFF;
    }
    
    free(buffer);
}

struct ipt_entry_match *get_match(const char *sports, const char *dports, unsigned int *nfcache, char *protocol)
{
    struct ipt_entry_match *match;
    struct ipt_udp *udpinfo;
    size_t size;

    size = IPT_ALIGN(sizeof(*match)) + IPT_ALIGN(sizeof(*udpinfo));
    match = (struct ipt_entry_match *)calloc(1, size);
    match->u.match_size = size;
    strncpy(match->u.user.name, protocol, IPT_FUNCTION_MAXNAMELEN - 2);

    udpinfo = (struct ipt_udp *)match->data;
    udpinfo->spts[1] = udpinfo->dpts[1] = 0xFFFF;
    //printf("sports=%s,dports=%s\n", sports, dports);
    if (sports) {
        *nfcache |= NFC_IP_SRC_PT;
        parse_ports(sports, udpinfo->spts);
    }
    if (dports) {
        *nfcache |= NFC_IP_DST_PT;
        parse_ports(dports, udpinfo->dpts);
    }

    return match;
}

int iptc_add_rule(const char *table, const char *chain, int protocol, const char *iniface, const char *outiface, \
    unsigned int src, unsigned int dest, const char *srcports, const char *destports, const char *target, const char *dnat_to, const int append)
{
    struct xtc_handle *handle;
    struct ipt_entry *chain_entry;
    struct ipt_entry_match *entry_match = NULL;
    struct ipt_entry_target *entry_target = NULL;
    ipt_chainlabel labelit;
    long match_size;
    int result = 0;

    chain_entry = (struct ipt_entry *)calloc(1, sizeof(*chain_entry));
    if (chain_entry == NULL) {
        perror("calloc failed");
        return 1;
    }

    if (src != 0) {
        chain_entry->ip.src.s_addr = src;
        chain_entry->ip.smsk.s_addr = 0xFFFFFFFF;
    }
    if (dest != 0) {
        chain_entry->ip.dst.s_addr = dest;
        chain_entry->ip.dmsk.s_addr = 0xFFFFFFFF;
    }

    if (iniface)
        strncpy(chain_entry->ip.iniface, iniface, IFNAMSIZ);
    if (outiface)
        strncpy(chain_entry->ip.outiface, outiface, IFNAMSIZ);

    chain_entry->ip.proto = protocol;
    if (IPPROTO_UDP == protocol)
        entry_match = get_match(srcports, destports, &chain_entry->nfcache, "udp");
    if (IPPROTO_TCP == protocol)
        entry_match = get_match(srcports, destports, &chain_entry->nfcache, "tcp");

    if (strcmp(target, "") == 0 || strcmp(target, IPTC_LABEL_ACCEPT) == 0 || strcmp(target, IPTC_LABEL_DROP) == 0 || \
            strcmp(target, IPTC_LABEL_QUEUE) == 0 || strcmp(target, IPTC_LABEL_RETURN) == 0) {
        size_t size;
        size = IPT_ALIGN(sizeof(struct ipt_entry_target)) + IPT_ALIGN(sizeof(int));
        entry_target = (struct ipt_entry_target *)calloc(1, size);
        entry_target->u.user.target_size = size;
        strncpy(entry_target->u.user.name, target, XT_EXTENSION_MAXNAMELEN);
    }
    if (entry_match) {
        match_size = entry_match->u.match_size;
    } else
        match_size = 0;

    struct ipt_entry *tmp_ipt = chain_entry;
    chain_entry = (struct ipt_entry *)realloc(chain_entry, sizeof(*chain_entry) + match_size + entry_target->u.target_size);
    if (chain_entry == NULL) {
        free(tmp_ipt);
    }

    memcpy(chain_entry->elems + match_size, entry_target, entry_target->u.target_size);
    chain_entry->target_offset = sizeof(*chain_entry) + match_size;
    chain_entry->next_offset = sizeof(*chain_entry) + match_size + entry_target->u.target_size;
    if (entry_match) {
        memcpy(chain_entry->elems, entry_match, match_size);
    }

    handle = iptc_init(table);
    if (!handle) {
        printf("libiptc error: Can't initialize table %s, %s\n", table, iptc_strerror(errno));
        free(chain_entry);
        free(entry_target);
        if (entry_match)
            free(entry_match);
        return 1;
    }

    strncpy(labelit, chain, sizeof(ipt_chainlabel));
    //printf("Chain name:%s\n", chain);
    result = iptc_is_chain(chain, handle);
    if (!result) {
        printf("libiptc error: Chain %s does not exist!\n", chain);
        free(chain_entry);
        free(entry_target);
        if (entry_match)
            free(entry_match);
        return 1;
    }
    
    if (append)
    {
        result = iptc_append_entry(labelit, chain_entry, handle);
    }
    else
    {
        result = iptc_insert_entry(labelit, chain_entry, 0, handle);
    }

    
    if (!result) {
        printf("libiptc error: Can't add, %s\n", iptc_strerror(errno));
        free(chain_entry);
        free(entry_target);
        if (entry_match)
            free(entry_match);
        return 1;
    }
    
    result = iptc_commit(handle);
    if (!result) {
        printf("libiptc error: Commit error, %s\n", iptc_strerror(errno));
        free(chain_entry);
        free(entry_target);
        if (entry_match)
            free(entry_match);
        return 1;
    }

    if (entry_match)
        free(entry_match);
    free(entry_target);
    free(chain_entry);
    
    iptc_free(handle);

    return 0;
}