SylFilter

#include <glib.h>

#include <stdio.h>

#include <string.h>

#include "filter.h"

#include "filter-kvs.h"

#include "bayes-filter.h"

#define N_TOKENS 15

static XFilterKVS *junk_kvs;

static XFilterKVS *clean_kvs;

static XFilterKVS *prob_kvs;

/* Test */

typedef struct _XFilterBayesProbData

{

        GArray *array;

        XFilterBayesLearnStatus status;

} XFilterBayesProbData;

typedef struct _XFilterKeyCount

{

        const char *key;

        int count;

        double prob;

} XFilterKeyCount;

static void xfilter_bayes_content_word_freq(GHashTable *table, const char *prefix, const char *text)

{

        const char *bp = text, *p = text;

        char *word;

        int count;

        if (!text)

                return;

        while (*p != '\0') {

                while (*p == ' ')

                        p++;

                bp = p;

                while (*p != '\0' && *p != ' ')

                        p++;

                if (p > bp) {

                        word = g_strndup(bp, p - bp);

                        if (prefix) {

                                char *bword = word;

                                word = g_strconcat(prefix, "*", bword, NULL);

                                g_free(bword);

                        }

                        count = GPOINTER_TO_INT(g_hash_table_lookup(table, word));

                        count++;

                        g_hash_table_insert(table, word, GINT_TO_POINTER(count));

                }

        }

}

static GHashTable *xfilter_bayes_word_freq(const XMessageData *data)

{

        GHashTable *table;

        const char *content;

        table = g_hash_table_new_full(g_str_hash, g_str_equal, g_free, NULL);

        content = xfilter_message_data_get_attribute(data, XM_FROM);

        xfilter_bayes_content_word_freq(table, "From", content);

        content = xfilter_message_data_get_attribute(data, XM_TO);

        xfilter_bayes_content_word_freq(table, "To", content);

        content = xfilter_message_data_get_attribute(data, XM_CC);

        xfilter_bayes_content_word_freq(table, "Cc", content);

        content = xfilter_message_data_get_attribute(data, XM_SUBJECT);

        xfilter_bayes_content_word_freq(table, "Subject", content);

        content = xfilter_message_data_get_content(data);

        xfilter_bayes_content_word_freq(table, NULL, content);

        return table;

}

static char *get_degenerated_word(const char *word)

{

        const char *p;

        if (!word)

                return NULL;

        if ((p = strchr(word, '*'))) {

                return g_strdup(p + 1);

        }

        if ((p = strchr(word, '!'))) {

                if (*(p + 1) == '!')

                        return g_strndup(word, p + 1 - word);

                else

                        return g_strndup(word, p - word);

        }

        for (p = word; *p != '\0'; p++) {

                if (g_ascii_isupper(*p))

                        return g_ascii_strdown(word, -1);

        }

        return NULL;

}

static double xfilter_get_prob(const char *key, XFilterBayesLearnStatus *status, gboolean do_degeneration)

{

        int n_junk;

        int n_clean;

        int n_junk_learn;

        int n_clean_learn;

        double prob = -1.0;

        n_junk = xfilter_kvs_fetch_int(junk_kvs, key);

        n_clean = xfilter_kvs_fetch_int(clean_kvs, key) * 2;

        if (n_junk + n_clean < 5) {

                if (n_junk > 0 && n_clean == 0) {

                        g_print("* %s: (%4f) (j: %d c: %d)\n", (gchar *)key, 0.6, n_junk, n_clean);

                        return 0.6;

                }

                if (do_degeneration) {

                        char *deg_key;

                        deg_key = get_degenerated_word(key);

                        if (deg_key) {

                                g_print("[degen] %s -> %s\n", key, deg_key);

                                prob = xfilter_get_prob(deg_key, status, TRUE);

                                g_free(deg_key);

                        }

                }

                return prob;

        } 

        n_junk_learn = status->junk_learned_num;

        if (n_junk_learn < 1)

                n_junk_learn = 1;

        n_clean_learn = status->nojunk_learned_num;

        if (n_clean_learn < 1)

                n_clean_learn = 1;

        prob = ((double)n_junk / n_junk_learn) /

                (((double)n_clean / n_clean_learn) + ((double)n_junk / n_junk_learn));

        if (prob < 0.01)

                prob = 0.01;

        else if (prob > 0.99) {

                if (n_clean == 0) {

                        if (n_junk > 10)

                                prob = 0.999;

                        else

                                prob = 0.998;

                } else

                        prob = 0.99;

        }

        g_print("%s: %4f (j: %d c: %d)\n", (gchar *)key, prob, n_junk, n_clean);

        return prob;

}

static void test_walk_func(gpointer key, gpointer val, gpointer data)

{

        XFilterBayesProbData *pdata;

        XFilterKeyCount kc;

        pdata = (XFilterBayesProbData *)data;

        kc.key = (gchar *)key;

        kc.count = GPOINTER_TO_INT(val);

        kc.prob = xfilter_get_prob(kc.key, &pdata->status, TRUE);

        //if (kc.prob > 0)

                //g_print("%s: (this: %d) %4f\n", kc.key, kc.count, kc.prob);

        if (kc.prob < 0)

                kc.prob = 0.4;

        g_array_append_val(pdata->array, kc);

}

static gint key_prob_compare_func(gconstpointer a, gconstpointer b)

{

        const XFilterKeyCount *kc1 = a;

        const XFilterKeyCount *kc2 = b;

        double da, db;

        da = ABS(0.5 - kc1->prob);

        db = ABS(0.5 - kc2->prob);

        return db * 10000 - da * 10000;

}

static XFilterStatus xfilter_bayes_func(XFilter *filter, const XMessageData *data, XFilterResult *result)

{

        const char *type;

        GHashTable *table;

        XFilterBayesProbData pdata;

        int i;

        double prod = 1.0, prod_rev = 1.0;

        double cmb_prob;

        XFilterStatus status;

        g_return_val_if_fail(result != NULL, XF_ERROR);

        type = xfilter_message_data_get_mime_type(data);

        if (!type || g_strncasecmp(type, "text/", 5) != 0) {

                xfilter_result_set_status(result, XF_UNSUPPORTED_TYPE);

                return XF_UNSUPPORTED_TYPE;

        }

        if (!junk_kvs) {

                g_warning("Cannot open junk database");

                xfilter_result_set_status(result, XF_ERROR);

                return XF_ERROR;

        }

        g_print("bayes-guessing message\n");

        table = xfilter_bayes_word_freq(data);

        pdata.array = g_array_sized_new(FALSE, FALSE, sizeof(XFilterKeyCount), 128);

        xfilter_bayes_get_learn_status(&pdata.status);

        g_print("\ncalculating probability for each tokens:\n");

        g_hash_table_foreach(table, test_walk_func, &pdata);

        g_array_sort(pdata.array, key_prob_compare_func);

        g_print("\nmost interesting tokens:\n");

        for (i = 0; i < 15 && i < pdata.array->len; i++) {

                XFilterKeyCount kc = g_array_index(pdata.array, XFilterKeyCount, i);

                prod *= kc.prob;

                prod_rev *= 1 - kc.prob;

                g_print("%s: %d %4f\n", kc.key, kc.count, kc.prob);

        }

        cmb_prob = prod / (prod + prod_rev);

        g_print("\ncombined probability: %4f\n", cmb_prob);

        g_array_free(pdata.array, TRUE);

        g_hash_table_destroy(table);

        xfilter_result_set_probability(result, cmb_prob);

        if (cmb_prob > 0.90)

                status = XF_JUNK;

        else if (cmb_prob < 0.10)

                status = XF_NOJUNK;

        else

                status = XF_UNCERTAIN;

        xfilter_result_set_status(result, status);

        return status;

}

XFilter *xfilter_bayes_new(void)

{

        XFilter *filter;

        filter = xfilter_new(XF_TEST, "bayes-test");

        xfilter_set_test_filter_func(X_TEST_FILTER(filter), xfilter_bayes_func);

        return filter;

}

/* Learning */

static void learn_walk_func(gpointer key, gpointer val, gpointer data)

{

        XFilterKVS *kvs = (XFilterKVS *)data;

        //g_print("%s: %d (%s)\n", (gchar *)key, GPOINTER_TO_INT(val), kvs == junk_kvs ? "j" : "c");

        if (xfilter_kvs_increment(kvs, (gchar *)key, GPOINTER_TO_INT(val)) < 0)

                g_warning("database update error");

}

static void learn_update_prob_func(gpointer key, gpointer val, gpointer data)

{

        XFilterBayesLearnStatus *s = (XFilterBayesLearnStatus *)data;

        double prob;

        char prob_str[7] = "";

        prob = xfilter_get_prob((gchar *)key, s, FALSE);

        if (prob < 0)

                return;

        g_snprintf(prob_str, sizeof(prob_str), "%4f", prob);

        if (xfilter_kvs_update(prob_kvs, (gchar *)key, prob_str, 6) < 0)

                g_warning("prob database update error");

}

static XFilterStatus xfilter_bayes_learn(XFilter *filter, const XMessageData *data, XFilterResult *result, gboolean is_junk)

{

        const char *type;

        GHashTable *table;

        XFilterKVS *kvs;

        XFilterBayesLearnStatus status = {0};

        g_return_val_if_fail(result != NULL, XF_ERROR);

        type = xfilter_message_data_get_mime_type(data);

        if (!type || g_strncasecmp(type, "text/", 5) != 0) {

                xfilter_result_set_status(result, XF_UNSUPPORTED_TYPE);

                return XF_UNSUPPORTED_TYPE;

        }

        if (is_junk)

                kvs = junk_kvs;

        else

                kvs = clean_kvs;

        if (!kvs) {

                g_warning("xfilter_bayes_learn: Cannot open database");

                xfilter_result_set_status(result, XF_ERROR);

                return XF_ERROR;

        }

        g_print("learning message\n");

        table = xfilter_bayes_word_freq(data);

        g_hash_table_foreach(table, learn_walk_func, kvs);

        if (is_junk)

                xfilter_kvs_increment(prob_kvs, "@junk_learn_count", 1);

        else

                xfilter_kvs_increment(prob_kvs, "@clean_learn_count", 1);

        xfilter_bayes_get_learn_status(&status);

        g_hash_table_foreach(table, learn_update_prob_func, &status);

        g_hash_table_destroy(table);

        xfilter_result_set_status(result, XF_NONE);

        return XF_NONE;

}

static XFilterStatus xfilter_bayes_learn_junk_func(XFilter *filter, const XMessageData *data, XFilterResult *result)

{

        return xfilter_bayes_learn(filter, data, result, TRUE);

}

static XFilterStatus xfilter_bayes_learn_nojunk_func(XFilter *filter, const XMessageData *data, XFilterResult *result)

{

        return xfilter_bayes_learn(filter, data, result, FALSE);

}

XFilter *xfilter_bayes_learn_junk_new(void)

{

        XFilter *filter;

        filter = xfilter_new(XF_CONTENT, "bayes-learn-junk");

        xfilter_set_content_filter_func(X_CONTENT_FILTER(filter), xfilter_bayes_learn_junk_func);

        return filter;

}

XFilter *xfilter_bayes_learn_nojunk_new(void)

{

        XFilter *filter;

        filter = xfilter_new(XF_CONTENT, "bayes-learn-clean");

        xfilter_set_content_filter_func(X_CONTENT_FILTER(filter), xfilter_bayes_learn_nojunk_func);

        return filter;

}

int xfilter_bayes_get_learn_status(XFilterBayesLearnStatus *status)

{

        g_return_val_if_fail(status != NULL, -1);

        status->junk_words = xfilter_kvs_count_sum(junk_kvs);

        status->nojunk_words = xfilter_kvs_count_sum(clean_kvs);

        status->junk_learned_num = xfilter_kvs_fetch_int(prob_kvs, "@junk_learn_count");

        status->nojunk_learned_num = xfilter_kvs_fetch_int(prob_kvs, "@clean_learn_count");

        return 0;

}

int xfilter_bayes_reset_learn_count(void)

{

        return 0;

}

int xfilter_bayes_reset_all(void)

{

        return 0;

}

static int show_walk_func(XFilterKVS *kvs, const char *key, void *value, int size, void *data)

{

        int ival;

        if (size == 4) {

                ival = *(int *)value;

                printf("%s: %d\n", key, ival);

        }

        return 0;

}

int xfilter_bayes_db_show_contents(void)

{

        XFilterBayesLearnStatus status = {0};

        if (!junk_kvs || !clean_kvs || !prob_kvs) {

                g_warning("Database not ready");

                return -1;

        }

        printf("Junk tokens:\n");

        xfilter_kvs_foreach(junk_kvs, show_walk_func, NULL);

        printf("\nClean tokens:\n");

        xfilter_kvs_foreach(clean_kvs, show_walk_func, NULL);

        printf("\nStatus:\n");

        xfilter_bayes_get_learn_status(&status);

        printf("junk_words: %d\n", status.junk_words);

        printf("nojunk_words: %d\n", status.nojunk_words);

        printf("junk_learned_num: %d\n", status.junk_learned_num);

        printf("nojunk_learned_num: %d\n", status.nojunk_learned_num);

        return 0;

}

int xfilter_bayes_db_init(void)

{

        g_print("xfilter_bayes_db_init: init database\n");

        if (!junk_kvs) {

                junk_kvs = xfilter_kvs_open("junk.db");

                if (!junk_kvs) {

                        g_warning("Cannot open database: junk.db");

                        return -1;

                }

        }

        if (!clean_kvs) {

                clean_kvs = xfilter_kvs_open("clean.db");

                if (!clean_kvs) {

                        g_warning("Cannot open database: clean.db");

                        xfilter_kvs_close(junk_kvs);

                        return -1;

                }

        }

        if (!prob_kvs) {

                prob_kvs = xfilter_kvs_open("prob.db");

                if (!prob_kvs) {

                        g_warning("Cannot open database: prob.db");

                        xfilter_kvs_close(clean_kvs);

                        xfilter_kvs_close(junk_kvs);

                        return -1;

                }

        }

        return 0;

}

int xfilter_bayes_db_done(void)

{

        int ret = 0;

        g_print("xfilter_bayes_db_init: close database\n");

        if (prob_kvs)

                ret |= xfilter_kvs_close(prob_kvs);

        if (clean_kvs)

                ret |= xfilter_kvs_close(clean_kvs);

        if (junk_kvs)

                ret |= xfilter_kvs_close(junk_kvs);

        return ret;

}