cringebot/bogofilter.py

import subprocess
from email.utils import format_datetime
from datetime import datetime
import quopri

BOGOFILTER_DB_DIR = "."
BOGOFILTER_COMMAND = ["bogofilter", "-T", "-d", BOGOFILTER_DB_DIR]

# Categories
SPAM   = "S"
HAM    = "H"
UNSURE = "U"
categories = [SPAM, HAM, UNSURE]

# Actions
CLASSIFY     = []
REGISTER     = ["-u"]
LEARN_SPAM   = ["-s"]
UNLEARN_SPAM = ["-S"]
LEARN_HAM    = ["-n"]
UNLEARN_HAM  = ["-N"]

ACTIONS = {
    REGISTER:     ["-u"],
    LEARN_SPAM:   ["-s"],
    UNLEARN_SPAM: ["-S"],
    LEARN_HAM:    ["-n"],
    UNLEARN_HAM:  ["-N"],

    SPAM:         ["-s"],
    HAM:          ["-n"],
    UNSURE:       []
}

class BogofilterResult:
    def __init__(self, category, score):
        self.category = category
        self.score = score

class Mail:
    def __init__(self, headers = {}, body = None):
        self.headers = {**{
            "Date": datetime.now(),
            "Content-Type": "text/html; charset=\"UTF-8\""}, **headers}
        self.body = body

    def format(self):
        text = str()
        for key, value in self.headers.items():
            if key == "Subject":
                value = "=?utf-8?Q?{}?=".format(quopri.encodestring(bytes(value, "utf-8"), header = True).decode("utf-8"))

            if key == "Date":
                value = format_datetime(value)

            text += "{key}: {value}\n".format(key = key, value = value)

        text += "\n"

        if self.body:
            text += self.body

        text += "\n"

        return text

def run(text, actions = [CLASSIFY], category = UNSURE):
    args = []
    for action in actions:
        args.extend(ACTIONS[action])
    
    cp = subprocess.run(BOGOFILTER_COMMAND + args, capture_output = True, encoding = "utf-8", input = text)
    arr = cp.stdout.strip().split(" ")
    if len(arr) == 2:
        (category, score) = arr
        return BogofilterResult(category, score)
    else:
        if cp.stderr:
            print("Bogofilter:")
            print(cp.stderr.strip())
        return None