#include "scramble_wlw.h"
#include "patricia.h"
#include <arpa/inet.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>

#define PREFIX_DELIMITER ",; "

static patricia_tree_t *tree4 = NULL;
static patricia_tree_t *tree6 = NULL;
static prefix_t lookup_pfix4;
static prefix_t lookup_pfix6;
static scrambler_ip4_t _scrambler_ip4 = NULL;
static scrambler_ip6_t _scrambler_ip6 = NULL;


static inline int
scramble_wlw_lookup_addr4(const struct in_addr *sin) {
    if (tree4 == NULL)
        return 0;
    memcpy(&lookup_pfix4.add.sin, sin, sizeof(*sin));
    return (patricia_search_best2(tree4, &lookup_pfix4, 1) != NULL);
}

static inline int
scramble_wlw_lookup_addr6(const struct in6_addr *sin6) {
    if (tree6 == NULL)
        return 0;
    memcpy(&lookup_pfix6.add.sin6, sin6, sizeof(*sin6));
    return (patricia_search_best2(tree6, &lookup_pfix6, 1) != NULL);
}

/* Init white list wrapper.  Can be called again to add wl-state.
 * prefix_wlist is a list of ip4 and ip6 prefixes, separated by
 * one of the PREFIX_DELIMITER
 * s4 is an ipv4 address scrambler
 * s6 is an ipv6 address scrambler
 */
void
scramble_wlw_init(const char *prefix_wlist, scrambler_ip4_t s4, scrambler_ip6_t s6) {
    char *plist = strdup(prefix_wlist), *cp;
    int afam, bits;
    patricia_tree_t **tree;

    _scrambler_ip4 = s4;
    _scrambler_ip6 = s6;

    for (cp = strtok(plist, PREFIX_DELIMITER); cp != NULL; cp = strtok(NULL, PREFIX_DELIMITER)) {
        if (index(cp, ':') == NULL) {
            tree = &tree4;
            afam = AF_INET;
            bits = 32;
        } else {
            tree = &tree6;
            afam = AF_INET6;
            bits = 128;
        }
        prefix_t *pref = ascii2prefix(afam, cp);
        if (pref == NULL) {
            fprintf(stderr, "ERROR: can't parse prefix %s\n", cp);
            exit(1);
        }
        if (index(cp, '/') == NULL) {
            pref->bitlen = bits;
        }
        if (*tree == NULL) {
            *tree = New_Patricia(bits);
        }
        patricia_lookup(*tree, pref);
        Deref_Prefix(pref);
    }
    lookup_pfix4.family = AF_INET;
    lookup_pfix4.bitlen = 32;
    lookup_pfix4.ref_count = 0;

    lookup_pfix6.family = AF_INET6;
    lookup_pfix6.bitlen = 128;
    lookup_pfix6.ref_count = 0;
}

/* Clear state */
void
scramble_wlw_clear() {
    if (tree4) {
        Destroy_Patricia(tree4, NULL);
        tree4 = NULL;
    }
    if (tree6) {
        Destroy_Patricia(tree6, NULL);
        tree6 = NULL;
    }
}


uint32_t
scramble_wlw_v4(uint32_t oldsrc, int pass_bits) {
    struct in_addr sin;
    if (tree4) {
        sin.s_addr = oldsrc;
        if (scramble_wlw_lookup_addr4(&sin))
            return oldsrc; /* whitelisted */
    }
    return _scrambler_ip4(oldsrc, pass_bits);
}

void
scramble_wlw_v6(struct in6_addr *sin6, int pass_bits) {
    if (tree6) {
        if (scramble_wlw_lookup_addr6(sin6))
            return; /* whitelisted */
    }
    _scrambler_ip6(sin6, pass_bits);
}