From 01dd848f5f94d0f19d2483b2415c878dd02adf52 Mon Sep 17 00:00:00 2001
From: StoneArk <stoneark86@gmail.com>
Date: Fri, 9 Jun 2017 18:00:35 +0800
Subject: [PATCH] Support read email content from stdin.

---
 getemailattach.py | 111 +++++++++++++++++++++++++++++++---------------
 1 file changed, 76 insertions(+), 35 deletions(-)

diff --git a/getemailattach.py b/getemailattach.py
index dfd3bdf..c366d14 100644
--- a/getemailattach.py
+++ b/getemailattach.py
@@ -1,48 +1,89 @@
+# getemailattach.py
+# Version: 2.0
+# Date: 2017-6-9
+# By: StoneArk
+# Extract email attachments from MIME.
+
 import sys
 import email
 import os
 
-if len(sys.argv) < 2:
-	print 'Usage:'
-	print 'getemailattach.py <SOURCE_MIME_PATH> [<OUTPUT_FOLDER_PATH>] [-d]'
-	print 'Parameters:'
-	print '<SOURCE_MIME_PATH>\tRequired, support a file path or a folder path.'
-	print '<OUTPUT_FOLDER_PATH>\tOptional, default value is the current path.'
-	print '-d\t\t\tDelete source MIME files after completed.'
-	quit()
-
-arrSrcFile = []
-destFolder = ''
-ifDelete = False
-successCount = 0
-srcMIME = sys.argv[1]
-if len(sys.argv) > 2:
-	if sys.argv[2] == '-d':
-		ifDelete = True
-	else:
-		destFolder = sys.argv[2]
-		if len(sys.argv) > 3 and sys.argv[3] == '-d':
-			ifDelete = True
-
-if os.path.isfile(srcMIME):
-	arrSrcFile = [srcMIME]
-else:
-	arrSrcFile = [os.path.join(srcMIME, f) for f in os.listdir(srcMIME) if os.path.isfile(os.path.join(srcMIME, f))]
+def parseEmailFiles(arrFiles, destFolder, ifDelete):
+	successCount = 0
+	for mimeFileItem in arrFiles:
+		mimeMsg = email.message_from_file(open(mimeFileItem))
+		if mimeMsg.is_multipart() == False:
+			print('Unable to load file: ' + mimeFileItem)
+			continue
+		extractSuccess = extractAttachment(mimeMsg, destFolder)
+		if extractSuccess == True:
+			successCount+=1
+			if ifDelete == True:
+				os.remove(mimeFileItem)
+	return successCount
 
-for mimeFileItem in arrSrcFile:
-	mimeMsg = email.message_from_file(open(mimeFileItem))
+def parseEmailString(strMIME, destFolder):
+	mimeMsg = email.message_from_string(strMIME)
 	if mimeMsg.is_multipart() == False:
-		print('Unable to load file: ' + mimeFileItem)
-		continue
+		print('Unable to load email content')
+		return False
+	extractSuccess = extractAttachment(mimeMsg, destFolder)
+	return extractSuccess
+
+def extractAttachment(mimeMsg, destFolder):
 	attachmentCount = len(mimeMsg.get_payload())
 	if attachmentCount < 2:
 		print('No attachment found in this email: ' + mimeFileItem)
-		continue
+		return False
 	for i in range(1, attachmentCount):
 		attachmentItem = mimeMsg.get_payload()[i]
 		attachmentFilename = attachmentItem.get_filename()
 		open(os.path.join(destFolder, attachmentFilename), 'wb').write(attachmentItem.get_payload(decode=True))
-	if ifDelete == True:
-		os.remove(mimeFileItem)
-	successCount+=1
-print(str(len(arrSrcFile)) + ' files found, ' + str(successCount) + ' successfully extracted.')
+	return True
+
+if __name__ == "__main__":
+	if len(sys.argv) < 2:
+		print 'Usage:'
+		print 'getemailattach.py <SOURCE_MIME_PATH> [<OUTPUT_FOLDER_PATH>] [-d]'
+		print 'Or'
+		print 'getemailattach.py -stdin [<OUTPUT_FOLDER_PATH>]'
+		print '\nParameters:'
+		print '<SOURCE_MIME_PATH>\tRequired, support a file path or a folder path.'
+		print '<OUTPUT_FOLDER_PATH>\tOptional, default value is the current path.'
+		print '-d\t\t\tDelete source MIME files after completed.'
+		print '-stdin\t\t\tRead email from standard input.'
+		quit()
+
+	srcMIME = sys.argv[1]
+	if srcMIME == '-stdin':
+		# From standard input.
+		destFolder = ''
+		if len(sys.argv) > 2 and sys.argv[2] != '-d':
+			destFolder = sys.argv[2]
+
+		strMIME = sys.stdin.read()
+		ifSuccess = parseEmailString(strMIME, destFolder)
+		if ifSuccess == True:
+			print('Attachment extracted successfully from stdin.')
+		else:
+			print('Attachment extracted FAILED from stdin.')
+	else:
+		# From local file or folder.
+		arrSrcFile = []
+		destFolder = ''
+		ifDelete = False
+		if len(sys.argv) > 2:
+			if sys.argv[2] == '-d':
+				ifDelete = True
+			else:
+				destFolder = sys.argv[2]
+				if len(sys.argv) > 3 and sys.argv[3] == '-d':
+					ifDelete = True
+
+		if os.path.isfile(srcMIME):
+			arrSrcFile = [srcMIME]
+		else:
+			arrSrcFile = [os.path.join(srcMIME, f) for f in os.listdir(srcMIME) if os.path.isfile(os.path.join(srcMIME, f))]
+
+		successCount = parseEmailFiles(arrSrcFile, destFolder, ifDelete)
+		print(str(len(arrSrcFile)) + ' files found, ' + str(successCount) + ' successfully extracted.')