Fix a minor misspelling

[mimedecode.git] / mimedecode.py

diff --git a/mimedecode.py b/mimedecode.py
index fdea039d04f0fc6ee9716137d86dc3763f72e0d4..48a6aba65ff01450825798e128319e0479d4e917 100755 (executable)
--- a/mimedecode.py
+++ b/mimedecode.py
@@ -1,10 +1,9 @@
 #! /usr/bin/env python
 """Decode MIME message"""
 
-from mimedecode_version import __version__, __author__, __copyright__, __license__
-
 import sys, os
-import email
+from mimedecode_version import __version__, \
+    __author__, __copyright__, __license__
 
 me = os.path.basename(sys.argv[0])
 
@@ -263,6 +262,18 @@ def totext(msg, instring):
     return s
 
 
+mimetypes = None
+
+def _guess_extension(ctype):
+    global mimetypes
+    if mimetypes is None:
+        import mimetypes
+        mimetypes.init()
+        user_mime_type = os.path.expanduser('~/.mime.types')
+        if os.path.exists(user_mime_type):
+            mimetypes._db.read(user_mime_type)
+    return mimetypes.guess_extension(ctype)
+
 def _save_message(msg, outstring, save_headers=False, save_body=False):
     for header, param in (
         ("Content-Disposition", "filename"),
@@ -270,16 +281,27 @@ def _save_message(msg, outstring, save_headers=False, save_body=False):
     ):
         fname = msg.get_param(param, header=header)
         if fname:
+            if isinstance(fname, tuple):
+                fname = fname[2] # Do not recode if it isn't recoded yet
+            try:
+                    for forbidden in chr(0), '/', '\\':
+                        if forbidden in fname:
+                            raise ValueError
+            except ValueError:
+                continue
             fname = '-' + fname
             break
     else:
         fname = ''
     g.save_counter += 1
     fname = str(g.save_counter) + fname
+    if '.' not in fname:
+        ext = _guess_extension(msg.get_content_type())
+        if ext: fname += ext
 
     global output
     save_output = output
-    outfile = open(os.path.join(g.destination_dir, fname), 'w')
+    outfile = open_output_file(fname)
     output = outfile.write
     if save_headers:
         output_headers(msg)
@@ -299,8 +321,7 @@ def decode_part(msg):
     ctype = msg.get_content_type()
     if ctype:
         masks.append(ctype)
-    mtype = msg.get_content_maintype()
-    if mtype:
+        mtype = ctype.split('/')[0]
         masks.append(mtype + '/*')
     masks.append('*/*')
 
@@ -395,6 +416,19 @@ def decode_message(msg):
         output(msg.as_string())
 
 
+def open_output_file(filename):
+    fullpath = os.path.abspath(os.path.join(g.destination_dir, filename))
+    full_dir = os.path.dirname(fullpath)
+    create = not os.path.isdir(full_dir)
+    if create:
+        os.makedirs(full_dir)
+    try:
+        return open(fullpath, 'w')
+    except:
+        if create:
+            os.removedirs(full_dir)
+
+
 class GlobalOptions:
     from m_lib.defenc import default_encoding
     recode_charset = 1 # recode charset of message body
@@ -521,7 +555,7 @@ if __name__ == "__main__":
         g.input_filename = '-'
         infile = sys.stdin
         if g.output_filename:
-            outfile = open(os.path.join(g.destination_dir, g.output_filename), 'w')
+            outfile = open_output_file(g.output_filename)
         else:
             g.output_filename = '-'
             outfile = sys.stdout
@@ -534,7 +568,7 @@ if __name__ == "__main__":
             infile = open(arguments[0], 'r')
         if la == 1:
             if g.output_filename:
-                outfile = open(os.path.join(g.destination_dir, g.output_filename), 'w')
+                outfile = open_output_file(g.output_filename)
             else:
                 g.output_filename = '-'
                 outfile = sys.stdout
@@ -546,7 +580,7 @@ if __name__ == "__main__":
                 outfile = sys.stdout
             else:
                 g.output_filename = arguments[1]
-                outfile = open(os.path.join(g.destination_dir, g.output_filename), 'w')
+                outfile = open_output_file(g.output_filename)
     else:
         usage(1, 'Too many arguments')
 
@@ -562,10 +596,11 @@ if __name__ == "__main__":
     g.outfile = outfile
     output = outfile.write
 
+    import email
     msg = email.message_from_file(infile)
 
     for header, value in g.set_header_value:
-        msg[header] = value
+        set_header(msg, header, value)
 
     for header, param, value in g.set_header_param:
         if header in msg: