pygrader/extract_mime.py

   1 # Copyright
   2
   3 """Extract message parts with a given MIME type from a mailbox.
   4 """
   5
   6 from __future__ import absolute_import
   7
   8 import email.utils as _email_utils
   9 import hashlib as _hashlib
  10 import mailbox as _mailbox
  11 import os as _os
  12 import os.path as _os_path
  13 import time as _time
  14
  15 from . import LOG as _LOG
  16 from .color import color_string as _color_string
  17 from .color import standard_colors as _standard_colors
  18
  19
  20 def message_time(message, use_color=None):
  21     highlight,lowlight,good,bad = _standard_colors(use_color=use_color)
  22     received = message['Received']  # RFC 822
  23     if received is None:
  24         mid = message['Message-ID']
  25         _LOG.debug(_color_string(
  26                 string='no Received in {}'.format(mid), color=lowlight))
  27         return None
  28     date = received.split(';', 1)[1]
  29     return _time.mktime(_email_utils.parsedate(date))
  30
  31 def extract_mime(message, mime_type=None, output='.', dry_run=False):
  32     _LOG.debug('parsing {}'.format(message['Subject']))
  33     time = message_time(message=message)
  34     for part in message.walk():
  35         fname = part.get_filename()
  36         if not fname:
  37             continue  # don't extract parts without filenames
  38         ffname = _os_path.join(output, fname)  # full file name
  39         ctype = part.get_content_type()
  40         if mime_type is None or ctype == mime_type:
  41             contents = part.get_payload(decode=True)
  42             count = 0
  43             base_ffname = ffname
  44             is_copy = False
  45             while _os_path.exists(ffname):
  46                 old = _hashlib.sha1(open(ffname, 'rb').read())
  47                 new = _hashlib.sha1(contents)
  48                 if old.digest() == new.digest():
  49                     is_copy = True
  50                     break
  51                 count += 1
  52                 ffname = '{}.{}'.format(base_ffname, count)
  53             if is_copy:
  54                 _LOG.debug('{} already extracted as {}'.format(fname, ffname))
  55                 continue
  56             _LOG.debug('extract {} to {}'.format(fname, ffname))
  57             if not dry_run:
  58                 with open(ffname, 'wb') as f:
  59                     f.write(contents)
  60                 if time is not None:
  61                     _os.utime(ffname, (time, time))