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]

SPAM   = "S"
HAM    = "H"
UNSURE = "U"

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

# If run with category == UNSURE, the message is classified
# If run with category == HAM | SPAM, the message is learned
# If learn = True, messages are learned as they are categorised
def run(text, category = UNSURE, learn = False):
    if category == SPAM:
        args = ["-s"]
    elif category == HAM:
        args = ["-n"]
    else:
        args = []

    if learn:
        args.append("-u")

    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