sb_pop3dnd.py

用python实现的邮件过滤器

class SpambayesIMAPServer(IMAP4Server):
    IDENT = "Spambayes IMAP Server IMAP4rev1 Ready"

    def __init__(self, user_account):
        IMAP4Server.__init__(self)
        self.account = user_account

    def authenticateLogin(self, user, passwd):
        """Lookup the account associated with the given parameters."""
        if user == options["imapserver", "username"] and \
           passwd == options["imapserver", "password"]:
            return (IAccount, self.account, None)
        raise cred.error.UnauthorizedLogin()

    def connectionMade(self):
        state.activeIMAPSessions += 1
        state.totalIMAPSessions += 1
        IMAP4Server.connectionMade(self)

    def connectionLost(self, reason):
        state.activeIMAPSessions -= 1
        IMAP4Server.connectionLost(self, reason)

    def do_CREATE(self, tag, args):
        """Creating new folders on the server is not permitted."""
        self.sendNegativeResponse(tag, \
                                  "Creation of new folders is not permitted")
    auth_CREATE = (do_CREATE, IMAP4Server.arg_astring)
    select_CREATE = auth_CREATE

    def do_DELETE(self, tag, args):
        """Deleting folders on the server is not permitted."""
        self.sendNegativeResponse(tag, \
                                  "Deletion of folders is not permitted")
    auth_DELETE = (do_DELETE, IMAP4Server.arg_astring)
    select_DELETE = auth_DELETE


class OneParameterFactory(ServerFactory):
    """A factory that allows a single parameter to be passed to the created
    protocol."""
    def buildProtocol(self, addr):
        """Create an instance of a subclass of Protocol, passing a single
        parameter."""
        if self.parameter is not None:
            p = self.protocol(self.parameter)
        else:
            p = self.protocol()
        p.factory = self
        return p


class RedirectingBayesProxy(POP3ProxyBase):
    """Proxies between an email client and a POP3 server, redirecting
    mail to the imap server as necessary.  It acts on the following
    POP3 commands:

     o RETR:
        o Adds the judgement header based on the raw headers and body
          of the message.
    """

    # This message could be a bit more informative - it could at least
    # say whether it's the spam or unsure folder.  It could give
    # information about who the message was from, or what the subject
    # was, if people thought that would be a good idea.
    intercept_message = 'From: "Spambayes" <[e-mail]>\r\n' \
                        'Subject: Spambayes Intercept\r\n\r\nA message ' \
                        'was intercepted by Spambayes (it scored %s).\r\n' \
                        '\r\nYou may find it in the Spam or Unsure ' \
                        'folder.\r\n\r\n'

    def __init__(self, clientSocket, serverName, serverPort, spam, unsure):
        POP3ProxyBase.__init__(self, clientSocket, serverName, serverPort)
        self.handlers = {'RETR': self.onRetr}
        state.totalSessions += 1
        state.activeSessions += 1
        self.isClosed = False
        self.spam_folder = spam
        self.unsure_folder = unsure

    def send(self, data):
        """Logs the data to the log file."""
        if options["globals", "verbose"]:
            state.logFile.write(data)
            state.logFile.flush()
        try:
            return POP3ProxyBase.send(self, data)
        except socket.error:
            self.close()

    def recv(self, size):
        """Logs the data to the log file."""
        data = POP3ProxyBase.recv(self, size)
        if options["globals", "verbose"]:
            state.logFile.write(data)
            state.logFile.flush()
        return data

    def close(self):
        # This can be called multiple times by async.
        if not self.isClosed:
            self.isClosed = True
            state.activeSessions -= 1
            POP3ProxyBase.close(self)

    def onTransaction(self, command, args, response):
        """Takes the raw request and response, and returns the
        (possibly processed) response to pass back to the email client.
        """
        handler = self.handlers.get(command, self.onUnknown)
        return handler(command, args, response)

    def onRetr(self, command, args, response):
        """Classifies the message.  If the result is ham, then simply
        pass it through.  If the result is an unsure or spam, move it
        to the appropriate IMAP folder."""
        # XXX This is all almost from sb_server!  We could just
        # XXX extract that out into a function and call it here.

        # Use '\n\r?\n' to detect the end of the headers in case of
        # broken emails that don't use the proper line separators.
        if re.search(r'\n\r?\n', response):
            # Remove the trailing .\r\n before passing to the email parser.
            # Thanks to Scott Schlesier for this fix.
            terminatingDotPresent = (response[-4:] == '\n.\r\n')
            if terminatingDotPresent:
                response = response[:-3]

            # Break off the first line, which will be '+OK'.
            ok, messageText = response.split('\n', 1)

            try:
                msg = email.message_from_string(messageText,
                                                _class=message.SBHeaderMessage)
                # Now find the spam disposition and add the header.
                (prob, clues) = state.bayes.spamprob(msg.tokenize(),\
                                 evidence=True)

                # Note that the X-SpamBayes-MailID header will be worthless
                # because we don't know the message id at this point.  It's
                # not necessary for anything anyway, so just don't set the
                # [Headers] add_unique_id option.
                msg.addSBHeaders(prob, clues)

                # Check for "RETR" or "TOP N 99999999" - fetchmail without
                # the 'fetchall' option uses the latter to retrieve messages.
                if (command == 'RETR' or
                    (command == 'TOP' and
                     len(args) == 2 and args[1] == '99999999')):
                    cls = msg.GetClassification()
                    dest_folder = None
                    if cls == options["Headers", "header_ham_string"]:
                        state.numHams += 1
                        headers = []
                        for name, value in msg.items():
                            header = "%s: %s" % (name, value)
                            headers.append(re.sub(r'\r?\n', '\r\n', header))
                        body = re.split(r'\n\r?\n', messageText, 1)[1]
                        messageText = "\r\n".join(headers) + "\r\n\r\n" + body
                    elif prob > options["Categorization", "spam_cutoff"]:
                        dest_folder = self.spam_folder
                        state.numSpams += 1
                    else:
                        dest_folder = self.unsure_folder
                        state.numUnsure += 1
                    if dest_folder:
                        msg = StringIO.StringIO(msg.as_string())
                        date = imaplib.Time2Internaldate(time.time())[1:-1]
                        dest_folder.addMessage(msg, (), date)

                        # We have to return something, because the client
                        # is expecting us to.  We return a short message
                        # indicating that a message was intercepted.
                        messageText = self.intercept_message % (prob,)
            except:
                messageText, details = \
                             message.insert_exception_header(messageText)

                # Print the exception and a traceback.
                print >>sys.stderr, details
            retval = ok + "\n" + messageText
            if terminatingDotPresent:
                retval += '.\r\n'
            return retval
        else:
            # Must be an error response.
            return response

    def onUnknown(self, command, args, response):
        """Default handler; returns the server's response verbatim."""
        return response


class RedirectingBayesProxyListener(Dibbler.Listener):
    """Listens for incoming email client connections and spins off
    RedirectingBayesProxy objects to serve them.
    """
    def __init__(self, serverName, serverPort, proxyPort, spam, unsure):
        proxyArgs = (serverName, serverPort, spam, unsure)
        Dibbler.Listener.__init__(self, proxyPort, RedirectingBayesProxy,
                                  proxyArgs)
        print 'Listener on port %s is proxying %s:%d' % \
               (_addressPortStr(proxyPort), serverName, serverPort)


class IMAPState(State):
    def __init__(self):
        State.__init__(self)

        # Set up the extra statistics.
        self.totalIMAPSessions = 0
        self.activeIMAPSessions = 0

    def createWorkers(self):
        """There aren't many workers in an IMAP State - most of the
        work is done elsewhere.  We do need to load the classifier,
        though, and build the status strings."""
        # Load token and message databases.
        if not hasattr(self, "DBName"):
            self.DBName, self.useDB = storage.database_type([])
        self.bayes = storage.open_storage(self.DBName, self.useDB)
        if not hasattr(self, "MBDName"):
            self.MDBName, self.useMDB = message.database_type()
        self.mdb = message.open_storage(self.MDBName, self.useMDB)

        # Load stats manager.
        self.stats = Stats(options, self.mdb)

        # Build status strings.
        self.buildStatusStrings()

    def buildServerStrings(self):
        """After the server details have been set up, this creates string
        versions of the details, for display in the Status panel."""
        self.serverPortString = str(self.imap_port)
        # Also build proxy strings
        State.buildServerStrings(self)

state = IMAPState()

# ===================================================================
# __main__ driver.
# ===================================================================

def prepare():
    # Setup state, server, boxes, trainers and account.
    state.imap_port = options["imapserver", "port"]
    state.createWorkers()
    proxyListeners = []

    spam_box = SpambayesMailbox("Spam", 0,
                                options["Storage", "spam_cache"])
    unsure_box = SpambayesMailbox("Unsure", 1,
                                  options["Storage", "unknown_cache"])
    ham_train_box = SpambayesMailbox("TrainAsHam", 2,
                                     options["Storage", "ham_cache"])
    # We don't have a third cache location in the directory, so make one up.
    spam_train_cache = os.path.join(options["Storage", "ham_cache"], "..",
                                    "spam_to_train")
    spam_train_box = SpambayesMailbox("TrainAsSpam", 3, spam_train_cache)
    inbox = SpambayesInbox(4, state)

    spam_trainer = Trainer(spam_train_box, True)
    ham_trainer = Trainer(ham_train_box, False)
    spam_train_box.addListener(spam_trainer)
    ham_train_box.addListener(ham_trainer)

    user_account = SpambayesAccount(options["imapserver", "username"],
                                    ham_train_box, spam_box, unsure_box,
                                    spam_train_box, inbox)

    # Add IMAP4 server.
    f = OneParameterFactory()
    f.protocol = SpambayesIMAPServer
    f.parameter = user_account
    reactor.listenTCP(state.imap_port, f)

    # Add POP3 proxy.
    for (server, serverPort), proxyPort in zip(state.servers,
                                               state.proxyPorts):
        listener = RedirectingBayesProxyListener(server, serverPort,
                                                 proxyPort, spam_box,
                                                 unsure_box)
        proxyListeners.append(listener)
    state.prepare()

def start():
    assert state.prepared, "Must prepare before starting"
    # The asyncore stuff doesn't play nicely with twisted (or vice-versa),
    # so put them in separate threads.
    thread.start_new_thread(Dibbler.run, ())
    reactor.run()

def stop():
    # Save the classifier, although that should not be necessary.
    state.bayes.store()
    # Explicitly closing the db is a good idea, though.
    state.bayes.close()
    
    # Stop the POP3 proxy.
    if state.proxyPorts:
        killer = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            killer.connect(('localhost', state.proxyPorts[0][1]))
            killer.send('KILL\r\n')
            killer.close()
        except socket.error:
            # Well, we did our best to shut down gracefully.  Warn the user
            # and just die when the thread we are in does.
            print "Could not shut down POP3 proxy gracefully."
    # Stop the IMAP4 server.
    reactor.stop()

def run():
    # Read the arguments.
    try:
        opts, args = getopt.getopt(sys.argv[1:], 'ho:')
    except getopt.error, msg:
        print >>sys.stderr, str(msg) + '\n\n' + __doc__
        sys.exit()

    for opt, arg in opts:
        if opt == '-h':
            print >>sys.stderr, __doc__
            sys.exit()
        elif opt == '-o':
            options.set_from_cmdline(arg, sys.stderr)

    # Let the user know what they are using...
    v = get_current_version()
    print v.get_long_version()
    from twisted.copyright import version as twisted_version
    print "Twisted version %s.\n" % (twisted_version,)

    # Setup everything.
    prepare()

    # Kick things off.
    start()


if __name__ == "__main__":
    run()