#!/usr/bin/env python

"""learn_spam_and_ham

Train spam assassin on the messages in particular Maildir folders,
then move the spam to the Junk folder, and (some of) the ham to the
Trash.

On success produces no output, making it suitable to be run by cron.

Eric Wollesen <ericw at xmtp dot net>
2008-09-04
"""

import os
import pwd
import sys
import re
from mailbox import Maildir
from tempfile import NamedTemporaryFile
from subprocess import Popen, PIPE, STDOUT

__version__ = "0.2"

DEBUG = False
INBOX_FOLDER = "/home/ericw/Maildir"
SPAM_FOLDER = "/home/ericw/Maildir/.Junk"
TRASH_FOLDER = "/home/ericw/Maildir/.Trash"
SPAM = [
    "Learn as Spam",
]
HAM_THEN_TRASH = [
    "Learn as Ham",
]


def GenerateFoldersFile(inbox):
    """Generate a "folders" file

    The folders file is a temporary file appropriate for use with the
    --folders flag of sa-learn.  See sa-learn(1p) for more info.
    """

    folders_file = NamedTemporaryFile(prefix="learn_spam_and_ham-")
    for folder in sorted(inbox.list_folders()):
        dir = "%s/.%s" % (INBOX_FOLDER, folder)
        if folder in SPAM:
            folders_file.write("spam:detect:%s/cur\n" % dir)
            folders_file.write("spam:detect:%s/new\n" % dir)
        elif folder in HAM_THEN_TRASH:
            folders_file.write("ham:detect:%s/cur\n" % dir)
            folders_file.write("ham:detect:%s/new\n" % dir)
    folders_file.flush()
    if DEBUG:
        folders_file.seek(0)
        print folders_file.read()

    return folders_file

def RunSaLearn(folders_file):
    """Run the sa-learn program

    Runs sa-learn, passing a file that lists each of the Maildir
    directories to be learned as spam or ham.

    If sa-learn returns non-zero, redirect its output to our stderr.

    Returns sa-learn's return value.
    """

    args = ["sa-learn", "--folders=%s" % folders_file.name]
    if DEBUG:
        args.insert(0, "echo")
    p = Popen(args, stdin=None, stdout=PIPE, stderr=STDOUT)
    ret = p.wait()
    if ret:
        print >>sys.stderr, "Error running sa-learn:"
        print >>sys.stderr, p.stdout.read()

    return ret

def MoveAndClearMessages(source, dest, dest_str="the bit bucket"):
    """Move messages, and clear their source Maildir

    Move the messages in the source Maildir to the destination
    Maildir, then remove all messages from the source Maildir.
    """

    for idx, msg in source.items():
        if DEBUG:
            print "Move to %s: %s" % (dest_str, msg["subject"])
        else:
            dest.add(msg)
    if not DEBUG:
        source.clear()

def MatchingFolders(inbox, list):
    """Return a list of matching Maildir directories

    Finds which of the folders in the list exist within the inbox
    Maildir.
    """

    ret = []

    for folder in inbox.list_folders():
        if folder in list:
            ret.append(inbox.get_folder(folder))

    return ret


if "__main__" == __name__:
    inbox = Maildir(INBOX_FOLDER, factory=None, create=False)
    junk = Maildir(SPAM_FOLDER, factory=None, create=False)
    trash = Maildir(TRASH_FOLDER, factory=None, create=False)

    ret = RunSaLearn(GenerateFoldersFile(inbox))
    if ret:
        print "Errors running sa-learn, not moving any messages."
        sys.exit(ret)
    for folder in MatchingFolders(inbox, SPAM):
        MoveAndClearMessages(folder, junk, "Junk")
    for folder in MatchingFolders(inbox, HAM_THEN_TRASH):
        MoveAndClearMessages(folder, trash, "Trash")