gatt

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import argparse
parser = argparse.ArgumentParser(
    description="Process attachments in mail files.")
parser.add_argument('in_file', help='Mail File To Process',
                    metavar='~/Maildir/new/foo', nargs='*')
parser.add_argument('-d', '--dest', help='Directory (or file) to dump the attachment(s) to',
                    metavar='/tmp/attachments/',
                    default=None)
parser.add_argument('-l','--list', help="Just list the attachments",
                    action='store_true', default=False)
parser.add_argument('-a','--all', help="Extract all the attachments, not just "+
                    "ones with 'application' major types",
                    action='store_true', default=False)

# TODO: -b --body, --include, --exclude, --no-clobber
# Support output to stdout if one attatchment
# Maybe support '-' in command line to say read from input (in adition
# to other positional args)

args = parser.parse_args()

def move_to(ofiles, out):
    if out[-1] is os.path.sep and not os.path.isdir(out):
        os.mkdir(out)
    if len(ofiles) == 1:
        shutil.copy(ofiles[0][0], out)
        if os.path.isdir(out):
            # TODO: Maybe don't need to print this out when -l --list action is
            # implemnted?
            print "Saved to %s%s" % (out+'/',ofiles[0][0].split('/')[-1])
    else:
        if os.path.exists(os.path.expanduser(out)) and not os.path.isdir(os.path.expanduser(out)):
                print >> sys.stderr,\
                    "Destination %s is not a directory and there a multiple attachments." % out
                print >> sys.stderr, "Extracted files are in %s" % tout
                exit(1)
        elif not os.path.exists(os.path.expanduser(out)):
            os.mkdir(out)

        for f,a in ofiles:
            shutil.copy(f, out)

def tree_list(z, prefix='└'):
    """
    git clone git://lair.fifthhorseman.net/~dkg/printmimestructure
    """
    if (z.is_multipart()):
        print prefix + '┬╴' + z.get_content_type(), z.as_string().__len__().__str__() + ' bytes'
        if prefix.endswith('└'):
            prefix = prefix.rpartition('└')[0] + ' '
        if prefix.endswith('├'):
            prefix = prefix.rpartition('├')[0] + '│'
        parts = z.get_payload()
        i = 0
        while (i < parts.__len__()-1):
            tree_list(parts[i], prefix + '├')
            i += 1
        tree_list(parts[i], prefix + '└')
        # FIXME: show epilogue?
    else:
        fname = '' if z.get_filename() is None else ' [' + z.get_filename() + ']'
        cset = '' if z.get_charset() is None else ' (' + z.get_charset() + ')'
        disp = z.get_params(None, header='Content-Disposition')
        if (disp is None):
            disposition = ''
        else:
            disposition = ''
            for d in disp:
                if d[0] in [ 'attachment', 'inline' ]:
                    disposition = ' ' + d[0]
        print prefix + '─╴'+ z.get_content_type() + cset + disposition + fname, z.get_payload().__len__().__str__(), 'bytes'


def process(in_file, tout):
    if not os.path.isfile(os.path.expanduser(in_file)):
        print >> sys.stderr, "Input file doesn't exist: %s" % os.path.expanduser(in_file)
        return []

    inf = open(os.path.expanduser(in_file),'r')

    from email.message import Message
    from email.parser import Parser

    msg = Parser().parse(inf)
    atts = msg.get_payload()

    ofiles = []

    if args.list and args.all:
        tree_list(msg)
        return ofiles


    # extract everything to a tempdir
    for a in msg.walk():
        if a.get_content_maintype() == "multipart":
            continue

        # only extract application/* parts by default
        if not args.all and (a.get_content_maintype() != "application"
                             and a.get_params([[""]],
                              header='Content-Disposition')[0][0] != "attachment"):
            continue

        fname = ""
        if a.get_params([]):
            for k, v in a.get_params():
                if k == 'name':
                    fname = v

        if args.list:
            print "%s \t%s" % (fname if fname else "unnamed",
                                  a.get_content_type())

        if not tout:
            continue

        if not fname:
            temp = tempfile.NamedTemporaryFile(suffix="."+a.get_content_subtype(),
                                               dir=tout, delete=False)
        else:
            temp = open(os.path.join(tout, fname),'wb')

        try:
            temp.write(a.get_payload(decode=True))
            ofiles.extend([(temp.name,a)])
        finally:
            temp.close()

    inf.close()
    return ofiles or []

infiles = args.in_file
if len(infiles) < 1:
    # try stdin
    infiles = sys.stdin.read().splitlines()

# TODO: Allow extract *and* list? (like tar -v)
import tempfile
ofiles = []
if args.list:
    tout = None
else:
    tout = tempfile.mkdtemp(prefix='atts')

out = args.dest or ""

for f in infiles:
    ofiles.extend(process(f, tout))

# todo: other actions like list
if out and tout:
    import shutil
    move_to(ofiles, out)
    shutil.rmtree(tout)
elif tout:
    print "Attachments are saved to ",tout