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 <unistd.h>

#include <errno.h>

#include "filter.h"

#include "filter-kvs.h"

#include "filter-utils.h"

#include "bayes-filter.h"

#include "libsylph/utils.h"

#define N_TOKENS 15

#undef USE_STATUS_KVS

static XFilterKVS *junk_kvs;

static XFilterKVS *clean_kvs;

#ifdef USE_STATUS_KVS

static XFilterKVS *prob_kvs;

#else

static XFilterBayesLearnStatus learn_status;

static FILE *status_fp;

#endif

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

                xfilter_kvs_begin(junk_kvs);

                xfilter_kvs_begin(clean_kvs);

                g_hash_table_foreach(table, test_walk_func, &pdata);

                xfilter_kvs_end(junk_kvs);

                xfilter_kvs_end(clean_kvs);

                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 */

typedef struct _XFilterLearnWalkData

{

        XFilterKVS *kvs;

        int sum;

} XFilterLearnWalkData;

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

{

        XFilterLearnWalkData *lwd = (XFilterLearnWalkData *)data;

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

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

                g_warning("database update error");

        lwd->sum += GPOINTER_TO_INT(val);

}

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 void xfilter_update_status(gboolean is_junk, gboolean is_register, int sum_add)

{

#ifdef USE_STATUS_KVS

        xfilter_kvs_begin(prob_kvs);

        if (is_register) {

                if (is_junk) {

                        xfilter_kvs_increment(prob_kvs, "@junk_words_sum", sum_add);

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

                } else {

                        xfilter_kvs_increment(prob_kvs, "@clean_words_sum", sum_add);

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

                }

        } else {

                if (is_junk) {

                        xfilter_kvs_set_int(prob_kvs, "@junk_words_sum", sum_add);

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

                } else {

                        xfilter_kvs_set_int(prob_kvs, "@clean_words_sum", sum_add);

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

                }

        }

        xfilter_kvs_end(prob_kvs);

#else /* !USE_STATUS_KVS */

        if (is_register) {

                if (is_junk) {

                        learn_status.junk_words += sum_add;

                        learn_status.junk_learned_num++;

                } else {

                        learn_status.nojunk_words += sum_add;

                        learn_status.nojunk_learned_num++;

                }

        } else {

                if (is_junk) {

                        learn_status.junk_words = sum_add;

                        if (learn_status.junk_learned_num > 0)

                                learn_status.junk_learned_num--;

                } else {

                        learn_status.nojunk_words = sum_add;

                        if (learn_status.nojunk_learned_num > 0)

                                learn_status.nojunk_learned_num--;

                }

        }

        ftruncate(fileno(status_fp), 0);

        rewind(status_fp);

        fprintf(status_fp, "@junk_words_sum=%d\n", learn_status.junk_words);

        fprintf(status_fp, "@junk_learn_count=%d\n", learn_status.junk_learned_num);

        fprintf(status_fp, "@clean_words_sum=%d\n", learn_status.nojunk_words);

        fprintf(status_fp, "@clean_learn_count=%d\n", learn_status.nojunk_learned_num);

        xfilter_debug_print("xfilter_update_status: writing status to file\n");

        if (fflush(status_fp) < 0) {

                perror("fflush");

                return;

        }

#if HAVE_FSYNC

        if (fsync(fileno(status_fp)) < 0) {

                perror("fsync");

        }

#elif defined(G_OS_WIN32)

        if (_commit(_fileno(status_fp)) < 0) {

                perror("_commit");

        }

#endif

        xfilter_debug_print("xfilter_update_status: done\n");

#endif /* !USE_STATUS_KVS */

}

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;

        int sum_add;

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

        xfilter_kvs_begin(kvs);

        if (is_register) {

                XFilterLearnWalkData lwd = {kvs, 0};

                g_hash_table_foreach(table, learn_walk_func, &lwd);

                sum_add = lwd.sum;

        } else {

                g_hash_table_foreach(table, unlearn_walk_func, kvs);

                sum_add = xfilter_kvs_count_sum(kvs);

        }

        xfilter_kvs_end(kvs);

        g_hash_table_destroy(table);

        xfilter_update_status(is_junk, is_register, sum_add);

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

#ifdef USE_STATUS_KVS

        status->junk_words = xfilter_kvs_fetch_int(prob_kvs, "@junk_words_sum");

        status->nojunk_words = xfilter_kvs_fetch_int(prob_kvs, "@clean_words_sum");

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

#else

        *status = learn_status;

#endif

        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(int verbose)

{

        XFilterBayesLearnStatus status = {0};

        if (!junk_kvs || !clean_kvs) {

                g_warning("Database not ready");

                return -1;

        }

        if (verbose) {

                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;

}

#ifndef USE_STATUS_KVS

int xfilter_read_status_file(FILE *fp)

{

        char buf[1024];

        int n;

        while (fgets(buf, sizeof(buf), fp) != NULL) {

                if (sscanf(buf, "@junk_words_sum=%d", &n) == 1)

                        learn_status.junk_words = n;

                else if (sscanf(buf, "@junk_learn_count=%d", &n) == 1)

                        learn_status.junk_learned_num = n;

                else if (sscanf(buf, "@clean_words_sum=%d", &n) == 1)

                        learn_status.nojunk_words = n;

                else if (sscanf(buf, "@clean_learn_count=%d", &n) == 1)

                        learn_status.nojunk_learned_num = n;

        }

        return 0;

}

#endif

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

        }

#ifdef USE_STATUS_KVS

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

        }

#else

        if (!status_fp) {

                if (path)

                        file = g_strconcat(path, G_DIR_SEPARATOR_S, "status.dat",

                                           NULL);

                else

                        file = g_strdup("status.dat");

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

                status_fp = g_fopen(file, "rb");

                if (!status_fp) {

                        if (ENOENT == errno)

                                status_fp = g_fopen(file, "w+b");

                        if (!status_fp) {

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

                                xfilter_kvs_close(clean_kvs);

                                xfilter_kvs_close(junk_kvs);

                                g_free(file);

                                return -1;

                        }

                } else {

                        xfilter_read_status_file(status_fp);

                        status_fp = freopen(file, "r+b", status_fp);

                        if (!status_fp) {

                                g_warning("Cannot reopen data file: %s", file);

                                xfilter_kvs_close(clean_kvs);

                                xfilter_kvs_close(junk_kvs);

                                g_free(file);

                                return -1;

                        }

                }

                g_free(file);

        }

#endif

        return 0;

}

int xfilter_bayes_db_done(void)

{

        int ret = 0;

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

#ifdef USE_STATUS_KVS

        if (prob_kvs)

                ret |= xfilter_kvs_close(prob_kvs);

#else

        if (status_fp)

                ret |= fclose(status_fp);

#endif

        if (clean_kvs)

                ret |= xfilter_kvs_close(clean_kvs);

        if (junk_kvs)

                ret |= xfilter_kvs_close(junk_kvs);

        return ret;

}