Merge pull request #90 from koodaamo/fix-54

put cmd-line attch zip generation support where it belongs

HISTORY.rst

@@ -5,6 +5,7 @@ tnefparse 1.4.0 (unreleased)
 - drop Python 2 support
 - drop Python 3.5 support (jugmac00)
 - add Python 3.9 support (jugmac00)
+- command-line support for zipped export of attachments (Beercow)
 - introduce using type annotations (jugmac00)
 - remove deprecated parseFile & raw_mapi functions
 - fix str representation for TNEF class (jugmac00)

README.rst

@@ -34,6 +34,7 @@ found inside them and so on::
    -h, --help             show this help message and exit
    -o, --overview         show (possibly long) overview of TNEF file contents
    -a, --attachments      extract attachments, by default to current dir
+   -z, --zip              extract attachments into a single zip file, by default to current dir
    -p PATH, --path PATH   optional explicit path to extract attachments to
    -b, --body             extract the body to stdout
    -hb, --htmlbody        extract the HTML body to stdout

tests/test_cmdline.py

@@ -1,3 +1,6 @@
+import os
+import io
+import zipfile
 import json
 import shutil
 import sys
@@ -25,6 +28,19 @@ def test_cmdline_attch_extract(script_runner):
     shutil.rmtree(tmpdir)
 
 
+def test_cmdline_zip_extract(script_runner):
+    tmpdir = tempfile.mkdtemp()
+    ret = script_runner.run('tnefparse', '-z', '-p', tmpdir, 'tests/examples/one-file.tnef')
+    assert os.path.isfile(tmpdir + '/attachments.zip')
+    assert ret.stderr == 'Successfully wrote attachments.zip\n'
+    with open(tmpdir + '/attachments.zip', 'rb') as zip_fp:
+        zip_stream = io.BytesIO(zip_fp.read())
+    zip_file = zipfile.ZipFile(zip_stream)
+    assert zip_file.namelist() == ['AUTHORS']
+    assert ret.success
+    shutil.rmtree(tmpdir)
+
+
 def test_cmdline_no_body(script_runner):
     ret = script_runner.run('tnefparse', '-b', 'tests/examples/one-file.tnef')
     assert ret.stderr == 'No body found\n'

tests/test_decoding.py

@@ -105,6 +105,7 @@ def test_decode(tnefspec):
 
 
 def test_zip():
+    # remove this test once tnef.to_zip(bytes) is no longer supported
     with open(datadir + os.sep + 'one-file.tnef', "rb") as tfile:
         zip_data = to_zip(tfile.read())
         with tempfile.TemporaryFile() as out:

tnefparse/cmdline.py

@@ -5,7 +5,7 @@
 import sys
 
 from . import properties
-from .tnef import TNEF
+from .tnef import TNEF, to_zip
 
 logging.basicConfig()
 logger = logging.getLogger(__package__)
@@ -25,6 +25,9 @@
 argument('-a', '--attachments', action='store_true',
          help='extract attachments, by default to current dir')
 
+argument('-z', '--zip', action='store_true',
+         help='extract attachments into a single zip file, by default to current dir')
+
 argument('-p', '--path',
          help='optional explicit path to extract attachments to')
 
@@ -97,6 +100,14 @@ def tnefparse() -> None:
             sys.stderr.write("Successfully wrote %i files\n" % len(t.attachments))
             sys.exit()
 
+        elif args.zip:
+            zipped = to_zip(t)
+            pth = args.path.rstrip(os.sep) + os.sep if args.path else ''
+            with open(pth + 'attachments.zip', "wb") as afp:
+                afp.write(zipped)
+            sys.stderr.write("Successfully wrote attachments.zip\n")
+            sys.exit()
+
         def print_body(attr: str, description: str) -> None:
             body = getattr(t, attr)
             if body is None:

tnefparse/tnef.py

@@ -2,6 +2,10 @@
 """
 import logging
 import os
+import warnings
+from typing import Union
+from zipfile import ZipFile, ZIP_DEFLATED, ZIP_STORED
+from io import BytesIO
 from datetime import datetime
 from uuid import UUID
 
@@ -391,12 +395,15 @@ def triples(data):
     return sender.rstrip(b'\x00'), etype, email.rstrip(b'\x00')
 
 
-def to_zip(data, default_name='no-name', deflate=True):
-    "Convert attachments in TNEF data to zip format. Accepts and returns str type."
-    # Parse the TNEF data
-    tnef = TNEF(data)
+def to_zip(tnef: Union[TNEF, bytes], default_name='no-name', deflate=True):
+    "Convert attachments in TNEF data to zip format."
 
-    # Convert the TNEF file to an equivalent ZIP file
+    if isinstance(tnef, bytes):
+        msg = "passing bytes to tnef.to_zip will be deprecated, pass a TNEF object instead"
+        warnings.warn(msg, DeprecationWarning)
+        tnef = TNEF(tnef)
+
+    # Extract attachments found in the TNEF object
     tozip = {}
     for attachment in tnef.attachments:
         filename = attachment.name or default_name
@@ -408,18 +415,13 @@ def to_zip(data, default_name='no-name', deflate=True):
         else:
             tozip[filename] = [(attachment.data, filename)]
 
-    # Add each attachment in the TNEF file to the zip file
-    from zipfile import ZipFile, ZIP_DEFLATED, ZIP_STORED
-    from io import BytesIO
-    import contextlib
-
+    # Add each attachment to the zip file
     sfp = BytesIO()
-    zf = ZipFile(sfp, "w", ZIP_DEFLATED if deflate else ZIP_STORED)
-    with contextlib.closing(zf) as z:
+    with ZipFile(sfp, "w", ZIP_DEFLATED if deflate else ZIP_STORED) as zf:
         for filename, entries in list(tozip.items()):
             for entry in entries:
                 data, name = entry
-                z.writestr(name, data)
+                zf.writestr(name, data)
 
     # Return the binary data for the zip file
     return sfp.getvalue()