"""
SMTP server base class.

The class in this module does not implement a fully functional SMTP
server; it simply provides a framework that subclasses can use to
build their own special-purpose SMTP servers.
"""


import email
import logging
import SocketServer


###############################################################################
###  class SMTPServer
###############################################################################

class SMTPServer(SocketServer.TCPServer):
  """
  An SMTP server class.
  """
  allow_reuse_address = 1


class ForkingSMTPServer(SocketServer.ForkingMixIn, SMTPServer): pass
class ThreadingSMTPServer(SocketServer.ThreadingMixIn, SMTPServer): pass


###############################################################################
###  class BaseSMTPRequestHandler
###############################################################################

# SMTP state values
_STATE_INIT	= 0
_STATE_MAIL	= 1
_STATE_RCPT	= 2

class BaseSMTPRequestHandler(SocketServer.StreamRequestHandler):
  """
  A base class for SMTP request handlers.

  At minimum, subclasses will need to override the new_message() method.
  They can also override the allow_sender() and allow_recipient() methods
  if they choose to implement sender- or recipient-based access control.
  """

  def handle(self):
    """
    Main loop for the SMTP connection.
    """
    self._close = 0
    self._helo = None
    self.reset_state()

    self.send_response(200, 'hi there')
    while not self._close:
      self.process_smtp_command()

  #############################################################################
  ###  utility methods
  #############################################################################

  def reset_state(self):
    """
    Utility function to reset the state of the SMTP session.
    Should be called before a new SMTP transaction is started.
    """
    self._state = _STATE_INIT
    self._mail_from = None
    self._rcpt_to = [ ]

  def read_command(self):
    """
    Read a command from the client, parse it into arguments,
    and return the argument list.
    """
    line = self.rfile.readline().strip()
    logging.debug('>>> ' + line)
    return line.split(' ')

  def send_response(self, code, msg):
    """
    Sends an SMTP response to the client.
    """
    response = '%d %s' % (code, msg)
    logging.debug('<<< ' + response)
    self.wfile.write(response + '\n')

  def process_smtp_command(self):
    """
    Reads an SMTP command and calls the appropriate handler method.
    """
    args = self.read_command()
    cmd = args.pop(0).upper()
    method_name = 'smtp_' + cmd
    if hasattr(self, method_name):
      method = getattr(self, method_name)
      code, msg = method(args)
    else:
      code = 501
      msg = 'unknown command "%s"' % cmd
    self.send_response(code, msg)

  #############################################################################
  ###  SMTP command handlers
  #############################################################################

  def smtp_HELO(self, args):
    """
    Handler for SMTP HELO command.
    """
    if len(args) != 1:
      return 501, 'usage: HELO hostname'
    if self._helo:
      return 503, 'already said HELO'
    self._helo = args[0]
    return 250, 'ok'

  def smtp_QUIT(self, args):
    """
    Handler for SMTP QUIT command.

    Sets self._close, which tells the main loop in the handler() method
    that the connection should be closed after sending this response.
    """
    self._close = 1
    return 221, 'closing connection'

  def smtp_RSET(self):
    """
    Handler for the SMTP RSET command.
    """
    self.reset_state()
    return 250, 'ok'

  def smtp_MAIL(self, args):
    """
    Handler for the SMTP MAIL command.
    """
    if len(args) == 1 and args[0].upper().startswith('FROM:') and \
       len(args[0]) > 5:
      args = [ 'FROM', args[0][5:] ]
    elif args[0].endswith(':'):
      args[0] = args[0][:len(args[0]) - 1]
    if len(args) != 2 or args[0].upper() != 'FROM':
      return 501, 'usage: MAIL FROM address'
    if self._state != _STATE_INIT:
      return 503, 'transaction already in progress - use RSET to abort'
    if not self.allow_sender(args[1]):
      return 553, 'bad sender'
    self._state = _STATE_MAIL
    self._mail_from = args[1]
    return 250, 'sender ok'

  def smtp_RCPT(self, args):
    """
    Handler for the SMTP RCPT command.
    """
    if len(args) == 1 and args[0].upper().startswith('TO:') and \
       len(args[0]) > 3:
      args = [ 'TO', args[0][3:] ]
    elif args[0].endswith(':'):
      args[0] = args[0][:len(args[0]) - 1]
    if len(args) != 2 or args[0].upper() != 'TO':
      return 501, 'usage: RCPT TO address'
    if self._state not in (_STATE_MAIL, _STATE_RCPT):
      return 503, 'send MAIL command first'
    if not self.allow_recipient(args[1]):
      return 553, 'bad recipient'
    self._state = _STATE_RCPT
    self._rcpt_to.append(args[1])
    return 250, 'recipient ok'

  def smtp_DATA(self, args):
    """
    Handler for the SMTP DATA command.
    """
    if args:
      return 501, 'usage: DATA'
    if self._state != _STATE_RCPT:
      return 503, 'send RCPT command first'
    self.send_response(334, 'end DATA with <CR><LF>.<CR><LF>')
    data = [ ]
    while True:
      line = self.rfile.readline().rstrip('\n\r')
      if line.startswith('.'):
	if len(line) == 1:
	  break
	data.append(line[1:])
	continue
      data.append(line)
    result = self.new_message(self._mail_from, self._rcpt_to, '\n'.join(data))
    self.reset_state()
    #return 250, 'message accepted'
    return result

  #############################################################################
  ###  methods to be implemented in subclasses
  #############################################################################

  def new_message(self, sender, recipients, msg):
    """
    Handler for new messages.  Must be implemented by subclasses.

    This method should return a tuple containing the SMTP response code
    and a one-line human-readable message.  On success, it should return
    something like (250, 'message accepted').
    """
    return 500, 'message handler not implemented'

  def allow_sender(self, sender):
    """
    Access control function to allow or deny the specified sender.
    This implementation will allow all senders.  Subclasses may override
    this method.

    This method should return 1 to allow the sender, or 0 otherwise.
    """
    return 1

  def allow_recipient(self, recipient):
    """
    Access control function to allow or deny the specified recipient.
    This implementation will allow all recipients.  Subclasses may override
    this method.

    This method should return 1 to allow the recipient, or 0 otherwise.
    """
    return 1