transaction: create a new file-like spool

Instead of having proxy methods and handling everything within
the transaction object, create a new Spool object that wraps
either StringIO or a NamedTemporaryFile, providing support for
flushing to disk and extracting the headers from a message.

vsmtpd/transaction.py

@@ -21,10 +21,78 @@
 #
 
 from cStringIO import StringIO
+from email.feedparser import FeedParser
 from tempfile import NamedTemporaryFile
 from vsmtpd.address import Address
 from vsmtpd.util import NoteObject
 
+class Spool(object):
+    """
+    Thin wrapper around either a file-object or a cStringIO object. Allows
+    use of the body attribute without losing functionality by avoiding
+    the class level methods of the transaction.
+    """
+
+    @property
+    def body_start(self):
+        return self._body_start
+
+    @property
+    def name(self):
+        return self._filename
+
+    def __init__(self):
+        self._body_start = 0
+        self._filename = None
+        self._fp       = None
+        self._parser   = None
+        self._rollover = 262144 # 256kb
+
+    def __getattr__(self, key):
+        return getattr(self._fp, key)
+
+    def end_headers(self):
+        """
+        Close off the parser and return the headers.
+        """
+        self._body_start = self._fp.tell()
+        return self._parser.close()
+
+    def flush(self):
+        """
+        Flushes the data held in memory to a temporary file on disk.
+        """
+        # Check to see if the data has been flushed already
+        if self._filename:
+            return
+
+        # Create the named temporary file to write the data to
+        fp = self._fp
+        newfp = self._fp = NamedTemporaryFile(dir='/tmp', prefix='')
+        if fp:
+            newfp.write(fp.getvalue())
+            newfp.seek(fp.tell(), 0)
+            fp.close()
+
+        self._filename = newfp.name
+
+    def write(self, data):
+        """
+        Write data to the end of the email.
+
+        :param data: The data to add to the end of the email
+        :type data: str
+        """
+        if not self._fp:
+            self._fp = StringIO()
+            self._parser = FeedParser()
+
+        if (self._fp.tell() + len(data)) > self._rollover:
+            self.flush()
+
+        self._fp.write(data)
+        self._parser.feed(data)
+
 class Transaction(NoteObject):
     """
     The Transaction class contains all the data relating to a single SMTP
@@ -48,12 +116,12 @@ def body(self):
         been accessed.
         """
         if not self._body:
-            self._body = StringIO()
+            self._body = Spool()
         return self._body
 
     @property
     def body_filename(self):
-        """ 
+        """
         Returns the temporary filename used to store the message contents;
         useful for virus scanners so that an additional copy doesn't need
         to be made.
@@ -65,20 +133,28 @@ def body_filename(self):
         """
         # We have to write the email to disk now into our spool directory
         # and update the body variable.
-        if not self._body_fn:
+        if not self.body.name:
             self.flush()
 
-        return self._body_fn
+        return self.body.name
 
     @property
     def data_size(self):
         """
         Get the current size of the email. Note that this is not the size
-        of the message that will be queued, it si the size of what the
+        of the message that will be queued, it is the size of what the
         client sent after the "DATA" command.
         """
         return self._body.tell() if self._body else 0
 
+    @property
+    def headers(self):
+        """
+        Create an email.message.Message object that contains the contents
+        of the messages headers.
+        """
+        return self._headers
+
     @property
     def recipients(self):
         """
@@ -107,7 +183,8 @@ def __init__(self, connection):
         self._sender     = None
         self._body       = None
         self._body_fn    = None
-        self._rollover   = 262144 # 256kb
+        self._headers    = None
+        self._header_size = 0
 
     def add_recipient(self, recipient):
         """
@@ -120,28 +197,26 @@ def add_recipient(self, recipient):
             recipient = Address(recipient)
         self._recipients.append(recipient)
 
-    def body_write(self, data):
-        """
-        Write data to the end of the email.
-
-        :param data: The data to add to the end of the email
-        :type data: str
-        """
-        if not self._body:
-            self._body = StringIO()
-
-        if (self._body.tell() + len(data)) > self._rollover:
-            self.flush()
-
-        self._body.write(data)
-
     def close(self):
         """
         Dispose all the resources that this transaction has opened.
         """
         if self._body:
             self._body.close()
 
+    def end_headers(self):
+        """
+        Mark the end of the messages headers and the beginning of the
+        body.
+        """
+        self._headers = self.body.end_headers()
+
+    def flush(self):
+        """
+        Flushes the data held in memory to a temporary file on disk.
+        """
+        self.body.flush()
+
     def remove_recipient(self, recipient):
         """
         Remove a recipient from the mail envelope of this message.
@@ -152,22 +227,3 @@ def remove_recipient(self, recipient):
         if not isinstance(recipient, Address):
             recipient = Address(recipient)
         self._recipients.remove(recipient)
-
-    def flush(self):
-        """
-        Flushes the data held in memory to a temporary file on disk.
-        """
-        # Check to see if the data has been flushed already
-        if self._body_fn:
-            return
-
-        # Create the named temporary file to write the data to
-        body = self._body
-        newbody = self._body = NamedTemporaryFile(dir='/tmp', prefix='')
-        if body:
-            newbody.write(body.getvalue())
-            newbody.seek(body.tell(), 0)
-        self._body_fn = newbody.name
-
-    def set_body_start(self):
-        self._header_size = self._body_start = self.data_size