SylFilter

/* SylFilter - a message filter

 *

 * Copyright (C) 2011 Hiroyuki Yamamoto

 * Copyright (C) 2011 Sylpheed Development Team

 */

#include <glib.h>

#include <stdio.h>

#include <string.h>

#include "filter.h"

#include "filter-kvs.h"

#include "filter-utils.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;

        double upper = 0.999;

        double lower = 0.001;

        //n_junk_learn = status->junk_learned_num;

        n_junk_learn = status->junk_words;

        if (n_junk_learn < 1)

                return -1.0;

        //n_clean_learn = status->nojunk_learned_num;

        n_clean_learn = status->nojunk_words;

        if (n_clean_learn < 1)

                return -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 == 0) {

                if (do_degeneration) {

                        char *deg_key;

                        deg_key = get_degenerated_word(key);

                        if (deg_key) {

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

                                prob = xfilter_get_prob(deg_key, status, TRUE);

                                g_free(deg_key);

                        }

                }

                return prob;

        }

        if (n_junk + n_clean < 5) {

                switch (n_junk + n_clean) {

                case 1:

                        upper = 0.6; lower = 0.4; break;

                case 2:

                        upper = 0.7; lower = 0.3; break;

                case 3:

                        upper = 0.8; lower = 0.2; break;

                case 4:

                        upper = 0.9; lower = 0.1; break;

                }

        } 

        prob = ((double)n_junk / n_junk_learn) /

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

        if (prob < lower) {

                if (n_junk == 0) {

                        if (n_clean > 10)

                                prob = lower;

                        else

                                prob = lower + 0.001;

                } else

                        prob = lower + 0.002;

        } else if (prob > upper) {

                if (n_clean == 0) {

                        if (n_junk > 10)

                                prob = upper;

                        else

                                prob = upper - 0.001;

                } else

                        prob = upper - 0.002;

        }

        xfilter_debug_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;

        }

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

        xfilter_bayes_get_learn_status(&pdata.status);

        if (pdata.status.junk_learned_num < 1) {

                xfilter_debug_print("junk message not learned yet\n");

                cmb_prob = 0.5;

        } else if (pdata.status.nojunk_learned_num < 1) {

                xfilter_debug_print("clean message not learned yet\n");

                cmb_prob = 0.5;

        } else {

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

                table = xfilter_bayes_word_freq(data);

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

                g_hash_table_foreach(table, test_walk_func, &pdata);

                g_array_sort(pdata.array, key_prob_compare_func);

                xfilter_debug_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;

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

                }

                cmb_prob = prod / (prod + prod_rev);

                xfilter_debug_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 unlearn_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_decrement(kvs, (gchar *)key, GPOINTER_TO_INT(val)) < 0)

                g_warning("database update error");

}

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

{

        const char *type;

        GHashTable *table;

        XFilterKVS *kvs;

        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;

        }

        xfilter_debug_print("%slearning %s message\n", is_register ? "" : "un", is_junk ? "junk" : "clean");

        table = xfilter_bayes_word_freq(data);

        if (is_register) {

                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);

        } else {

                g_hash_table_foreach(table, unlearn_walk_func, kvs);

                if (is_junk)

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

                else

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

        }

        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, TRUE);

}

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

{

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

}

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

{

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

}

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

{

        return xfilter_bayes_learn(filter, data, result, FALSE, 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;

}

XFilter *xfilter_bayes_unlearn_junk_new(void)

{

        XFilter *filter;

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

        xfilter_set_content_filter_func(X_CONTENT_FILTER(filter), xfilter_bayes_unlearn_junk_func);

        return filter;

}

XFilter *xfilter_bayes_unlearn_nojunk_new(void)

{

        XFilter *filter;

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

        xfilter_set_content_filter_func(X_CONTENT_FILTER(filter), xfilter_bayes_unlearn_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(const char *path)

{

        char *file;

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

        xfilter_debug_print("xfilter_bayes_db_init: path: %s\n",

                            path ? path : "(current dir)");

        if (path) {

                xfilter_debug_print("xfilter_bayes_db_init: making directory: %s\n", path);

                if (xfilter_utils_mkdir(path) < 0) {

                        g_warning("Making directory failed: %s", path);

                        return -1;

                }

        }

        if (!junk_kvs) {

                if (path)

                        file = g_strconcat(path, G_DIR_SEPARATOR_S, "junk.db",

                                           NULL);

                else

                        file = g_strdup("junk.db");

                xfilter_debug_print("xfilter_bayes_db_init: opening db: %s\n", file);

                junk_kvs = xfilter_kvs_open(file);

                if (!junk_kvs) {

                        g_warning("Cannot open database: %s", file);

                        g_free(file);

                        return -1;

                }

                g_free(file);

        }

        if (!clean_kvs) {

                if (path)

                        file = g_strconcat(path, G_DIR_SEPARATOR_S, "clean.db",

                                           NULL);

                else

                        file = g_strdup("clean.db");

                xfilter_debug_print("xfilter_bayes_db_init: opening db: %s\n", file);

                clean_kvs = xfilter_kvs_open(file);

                if (!clean_kvs) {

                        g_warning("Cannot open database: %s", file);

                        xfilter_kvs_close(junk_kvs);

                        g_free(file);

                        return -1;

                }

                g_free(file);

        }

        if (!prob_kvs) {

                if (path)

                        file = g_strconcat(path, G_DIR_SEPARATOR_S, "prob.db",

                                           NULL);

                else

                        file = g_strdup("prob.db");

                xfilter_debug_print("xfilter_bayes_db_init: opening db: %s\n", file);

                prob_kvs = xfilter_kvs_open(file);

                if (!prob_kvs) {

                        g_warning("Cannot open database: %s", file);

                        xfilter_kvs_close(clean_kvs);

                        xfilter_kvs_close(junk_kvs);

                        g_free(file);

                        return -1;

                }

                g_free(file);

        }

        return 0;

}

int xfilter_bayes_db_done(void)

{

        int ret = 0;

        xfilter_debug_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;

}