# Copyright (C) 2002-2003 by James Henstridge <james@daa.com.au>
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, US

'''Module used to communicate with a SpamAssassin spamd process to check
or tag messages.

Usage is as follows:
  >>> conn = spamd.SpamdConnection()
  >>> conn.addheader('User', 'username')
  >>> conn.check(spamd.SYMBOLS, 'From: user@example.com\n...')
  >>> print conn.getspamstatus()
  (True, 4.0)
  >>> print conn.response_message
  ...
'''

import socket
import StringIO
import mimetools

import __builtin__
if not hasattr(__builtin__, 'True'):
    __builtin__.True = (1 == 1)
    __builtin__.False = (1 != 1)
del __builtin__


class error(Exception):
    "exception"
    pass

SPAMD_PORT = 783

# available methods
SKIP = 'SKIP'
PROCESS = 'PROCESS'
CHECK = 'CHECK'
SYMBOLS = 'SYMBOLS'
REPORT = 'REPORT'
REPORT_IFSPAM = 'REPORT_IFSPAM'

# error codes
EX_OK = 0
EX_USAGE = 64
EX_DATAERR = 65
EX_NOINPUT = 66
EX_NOUSER = 67
EX_NOHOST = 68
EX_UNAVAILABLE = 69
EX_SOFTWARE = 70
EX_OSERR = 71
EX_OSFILE = 72
EX_CANTCREAT = 73
EX_IOERR = 74
EX_TEMPFAIL = 75
EX_PROTOCOL = 76
EX_NOPERM = 77
EX_CONFIG = 78


class SpamdConnection:
    '''Class to handle talking to SpamAssassin spamd servers.'''
    # default spamd
    host = 'localhost'
    port = SPAMD_PORT

    PROTOCOL_VERSION = 'SPAMC/1.3'

    def __init__(self, host='', port=0, socket='/tmp/spamd.sock'):
        if not port and ':' in host:
            host, port = host.split(':', 1)
            port = int(port)
        if host:
            self.host = host
        if port:
            self.port = port
        self.usock = socket

        # message structure to hold request headers
        self.request_headers = mimetools.Message(StringIO.StringIO(),
                                seekable=False)
        self.request_headers.fp = None

        # stuff that will be filled in after check()
        self.server_version = None
        self.result_code = None
        self.response_message = None
        self.response_headers = mimetools.Message(StringIO.StringIO(),
                                seekable=False)

    def addheader(self, header, value):
        '''Adds a header to the request.'''
        self.request_headers[header] = value

    def check(self, method='PROCESS', message=''):
        '''Sends a request to the spamd process.'''
        try:
            # sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            # sock.connect((self.host, self.port))
            sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            sock.connect(self.usock)
        except socket.error:
            raise error('could not connect to spamd on %s' % self.host)

        # set content length request header
        del self.request_headers['Content-length']
        self.request_headers['Content-length'] = str(len(message))

        request = '%s %s\r\n%s\r\n' % \
                  (method, self.PROTOCOL_VERSION,
                   str(self.request_headers).replace('\n', '\r\n'))

        try:
            sock.send(request)
            sock.send(message)
            sock.shutdown(1)  # shut down the send half of the socket
        except (socket.error, IOError):
            raise error('could not send request to spamd')

        fp = sock.makefile('rb')
        try:
            response = fp.readline()
        except IOError:
            raise error('could not read response code from spamd')
        words = response.split(None, 2)
        if len(words) != 3:
            raise error('not enough words in response header')
        if words[0][:6] != 'SPAMD/':
            raise error('bad protocol name in response string')
        self.server_version = float(words[0][6:])
        if self.server_version < 1.0 or self.server_version >= 2.0:
            raise error('incompatible server version')
        self.result_code = int(words[1])
        if self.result_code != 0:
            raise error('spamd server returned error %s' % words[2])

        try:
            # parse header
            self.response_headers = mimetools.Message(fp, seekable=False)
            self.response_headers.fp = None
        except IOError:
            raise error('could not read in response headers')

        try:
            # read in response message
            self.response_message = fp.read()
        except IOError:
            raise error('could not read in response message')

        fp.close()
        sock.close()

    def getspamstatus(self):
        '''Decode the "Spam" response header.'''
        if 'Spam' not in self.response_headers:
            raise error('Spam header not found in response')

        isspam, score = self.response_headers['Spam'].split(';', 1)
        isspam = (isspam.strip() != 'False')
        score = float(score.split('/', 1)[0])
        return isspam, score
