pygrader/extract_mime.py

   1 # Copyright (C) 2012 W. Trevor King <wking@drexel.edu>
   2 #
   3 # This file is part of pygrader.
   4 #
   5 # pygrader is free software: you can redistribute it and/or modify it under the
   6 # terms of the GNU General Public License as published by the Free Software
   7 # Foundation, either version 3 of the License, or (at your option) any later
   8 # version.
   9 #
  10 # pygrader is distributed in the hope that it will be useful, but WITHOUT ANY
  11 # WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
  12 # A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
  13 #
  14 # You should have received a copy of the GNU General Public License along with
  15 # pygrader.  If not, see <http://www.gnu.org/licenses/>.
  16
  17 """Extract message parts with a given MIME type from a mailbox.
  18 """
  19
  20 from __future__ import absolute_import
  21
  22 import email.utils as _email_utils
  23 import hashlib as _hashlib
  24 import mailbox as _mailbox
  25 import os as _os
  26 import os.path as _os_path
  27 import time as _time
  28
  29 from . import LOG as _LOG
  30 from .color import color_string as _color_string
  31 from .color import standard_colors as _standard_colors
  32
  33
  34 def message_time(message, use_color=None):
  35     """Get the Unix time when ``message`` was received.
  36
  37     >>> from email.utils import formatdate
  38     >>> from pgp_mime.email import encodedMIMEText
  39     >>> msg = encodedMIMEText('Ping!')
  40     >>> msg['Received'] = (
  41     ...     'from smtp.home.net (smtp.home.net [123.456.123.456]) '
  42     ...     'by smtp.mail.uu.edu (Postfix) with ESMTP id 5BA225C83EF '
  43     ...     'for <wking@tremily.us>; Sun, 09 Oct 2011 11:50:46 -0400 (EDT)')
  44     >>> time = message_time(msg)
  45     >>> time
  46     1318175446.0
  47     >>> formatdate(time)
  48     'Sun, 09 Oct 2011 15:50:46 -0000'
  49     """
  50     highlight,lowlight,good,bad = _standard_colors(use_color=use_color)
  51     received = message['Received']  # RFC 822
  52     if received is None:
  53         mid = message['Message-ID']
  54         _LOG.debug(_color_string(
  55                 string='no Received in {}'.format(mid), color=lowlight))
  56         return None
  57     date = received.split(';', 1)[1]
  58     return _time.mktime(_email_utils.parsedate(date))
  59
  60 def extract_mime(message, mime_type=None, output='.', dry_run=False):
  61     _LOG.debug('parsing {}'.format(message['Subject']))
  62     time = message_time(message=message)
  63     for part in message.walk():
  64         fname = part.get_filename()
  65         if not fname:
  66             continue  # don't extract parts without filenames
  67         ffname = _os_path.join(output, fname)  # full file name
  68         ctype = part.get_content_type()
  69         if mime_type is None or ctype == mime_type:
  70             contents = part.get_payload(decode=True)
  71             count = 0
  72             base_ffname = ffname
  73             is_copy = False
  74             while _os_path.exists(ffname):
  75                 old = _hashlib.sha1(open(ffname, 'rb').read())
  76                 new = _hashlib.sha1(contents)
  77                 if old.digest() == new.digest():
  78                     is_copy = True
  79                     break
  80                 count += 1
  81                 ffname = '{}.{}'.format(base_ffname, count)
  82             if is_copy:
  83                 _LOG.debug('{} already extracted as {}'.format(fname, ffname))
  84                 continue
  85             _LOG.debug('extract {} to {}'.format(fname, ffname))
  86             if not dry_run:
  87                 with open(ffname, 'wb') as f:
  88                     f.write(contents)
  89                 if time is not None:
  90                     _os.utime(ffname, (time, time))