argparse refactor WIP

search.py

@@ -1,3 +1,4 @@
+import argparse
 import re
 import platform
 import subprocess
@@ -18,6 +19,10 @@ class BooleanOperator(Enum):
 class Filter:
 	"""Base class for message filters."""
 	def matches(self, message: EmailMessage) -> bool:
+		"""Returns true if the given email message matches this filter's criteria."""
+		raise "Not implemented"
+	def matches_raw(self, raw_content: str) -> bool:
+		"""Returns true if the given raw, unparsed email content matches this filter's criteria."""
 		raise "Not implemented"
 
 class BodyKeywordFilter(Filter):
@@ -37,6 +42,15 @@ class BodyKeywordFilter(Filter):
 						return True
 		return False
 
+	def matches_raw(self, raw_content: str) -> bool:
+		if self.case_sensitive:
+			if self.keyword in raw_content:
+				return True
+			else:
+				if self.keyword.lower() in raw_content.lower():
+					return True
+		return False
+
 class HeaderFilter(Filter):
 	"""Matches a value in an email header. Can search one filter or multiple.
 	Header names case-insensitive; value is case-insensitive."""
@@ -70,10 +84,19 @@ class BooleanFilter(Filter):
 			return True
 		return False
 
-start_path = '.'
-output_path = TemporaryDirectory(prefix='Email search results (id ', suffix=')').name
+	def matches_raw(self, raw_content: str) -> bool:
+		for subfilter in self.subfilters:
+			result = subfilter.matches_raw(raw_content)
+			if self.operator == BooleanOperator.and_op and not result:
+				return False
+			if self.operator == BooleanOperator.or_op and result:
+				return True
+		if self.operator == BooleanOperator.and_op:
+			return True
+		return False
+
+args: argparse.Namespace = None
 filter = None
-case_sensitive = False
 result_count = 0
 zip_result_count = 0
 zip_count = 0
@@ -130,9 +153,9 @@ def process_zip_file(zip_path: str) -> None:
 				email = parser.parsebytes(data)
 				search_content(email, zip_path, entry)
 			except UnicodeError:
-				print('Unicode error in message. Skipping.')
+				print('Unicode error in message. Searching raw content.', file=sys.stderr)
 			except:
-				print('Error reading message')
+				print('Error parsing message. Searching raw content.', file=sys.stderr)
 	if zip_result_count > 0:
 		print(f"\t{zip_result_count} results in zip")
 
@@ -147,56 +170,133 @@ def search_content(email: EmailMessage, zip_path: str, entry: ZipInfo) -> None:
 			zip_result_count += 1
 			f.write(email.as_bytes())
 
+def search_raw_content(raw_bytes: bytes, zip_path: str, entry: ZipInfo) -> None:
+	global result_count, zip_result_count
+	try:
+		content = raw_bytes.decode('iso-8859-1', errors='ignore')
+	except:
+		try:
+			content = raw_bytes.decode('utf-8', errors='ignore')
+		except:
+			print('Cannot decode email bytes. Skipping.', file=sys.stderr)
+
 def parse_arguments():
-	"""Parses the command-line arguments."""
-	global filter
-	global start_path
-	global output_path
-	global case_sensitive
-	expect = 'script_name'
-	for arg in sys.argv:
-		if arg.startswith('-'):
-			if arg == '-d':
-				expect = 'start_path'
-			elif arg == '-o':
-				expect = 'output_path'
-			elif arg == '-c':
-				case_sensitive = True
-			else:
-				raise f'Unknown argument {arg}'
-		elif expect is not None:
-			if expect == 'script_name':
-				expect = None
-				continue
-			elif expect == 'start_path':
-				start_path = arg
-				expect = None
-			elif expect == 'output_path':
-				output_path = arg
-				expect = None
-			else:
-				raise f'Expected other argument {expect}'
-		else:
-			if filter is None:
-				words = arg.split(' ')
-				word_filters = []
-				for word in words:
-					word = word.strip()
-					if len(word) == 0:
-						continue
-					word_filters.append(BodyKeywordFilter(word, case_sensitive))
-				if len(word_filters) == 0:
-					continue
-				filter = BooleanFilter(BooleanOperator.and_op, word_filters)
-			else:
-				print('Too many arguments')
-				sys.exit(4)
+	"""Parses command-line arguments to `args`."""
+	global args
+	parser = argparse.ArgumentParser(
+		prog='search.py',
+		description='Searches a directory of zipped email messages. ' + \
+			'Messages are assumed to be stored one per file within the zip files (Maildir format). ' + \
+			'Input directories are searched recursively for any zip files contained within.',
+		epilog='Raw mode will skip parsing each email message and treat them like simple text files. ' + \
+			'The headers and body are all searched together without decoding. ' + \
+			'Arguments for searching individual fields will be ignored. ' + \
+			'This option exists for messages with encoding errors that prevent them from being ' + \
+			'parsed correctly. ' + \
+			'Note that various escaping/encoding schemes commonly used in email messages, such ' + \
+			'as base64, may cause keywords to not be found despite being in the decoded message ' + \
+			'because only the raw encoded content is searched. ' + \
+			'Use this option as a last resort.'
+	)
+	parser.add_argument(
+		'keywords',
+		action='append',
+		nargs='+',
+		help='one or more phrases to search for in the message body'
+	)
+	parser.add_argument(
+		'--any',
+		default=False,
+		action='store_true',
+		help='matches messages containing any of the given search phrases (default requires all phrases appear in a message)'
+	)
+	parser.add_argument(
+		'-d', '--dir',
+		action='append',
+		help='directory(s) to search for email zip archives (default is working directory)'
+	)
+	parser.add_argument(
+		'-o', '--output',
+		help='directory to copy matching messages to (default is a temp directory)'
+	)
+	parser.add_argument(
+		'-c', '--casesensitive',
+		default=False,
+		action='store_true',
+		help='search case-sensitively (default is case-insensitive)'
+	)
+	parser.add_argument(
+		'-f', '--from',
+		help='email address of sender'
+	)
+	parser.add_argument(
+		'-t', '--to',
+		help='email address of recipient (searches to:, cc:, bcc: fields)'
+	)
+	parser.add_argument(
+		'-s', '--subject',
+		help='searches subject field'
+	)
+	parser.add_argument(
+		'-a', '--after',
+		metavar='YYYY-MM-DD',
+		help='date to search on or after'
+	)
+	parser.add_argument(
+		'-b', '--before',
+		metavar='YYYY-MM-DD',
+		help='date to search on or before'
+	)
+	parser.add_argument(
+		'-r', '--raw',
+		default=False,
+		action='store_true',
+		help='searches raw email content (see below)'
+	)
+	args = parser.parse_args()
+
+	if args.before is not None:
+		m = re.match('^([0-9]{4})-([0-9]{2})-([0-9]{2})$', args.before)
+		if m is None:
+			parser.error('before date must be in YYYY-MM-DD format (e.g. 2015-03-28)')
+		args.before = [ int(m.group(1)), int(m.group(2)), int(m.group(3)) ]
+	if args.after is not None:
+		m = re.match('^([0-9]{4})-([0-9]{2})-([0-9]{2})$', args.after)
+		if m is None:
+			parser.error('after date must be in YYYY-MM-DD format (e.g. 2015-03-28)')
+		args.after = [ int(m.group(1)), int(m.group(2)), int(m.group(3)) ]
+	if args.raw:
+		if getattr(args, 'from') is not None or \
+			getattr(args, 'to') is not None or \
+			args.subject is not None or \
+			args.before is not None or \
+			args.after is not None:
+			print('Warning: Cannot search header fields in raw mode. Ignoring.', file=sys.stderr)
+	if args.dir is None:
+		args.dir = [ '.' ]
+	else:
+		for d in args.dir:
+			if not os.path.exists(d) or not os.path.isdir(d):
+				parser.error(f'search path \'{d}\' does not exist or is not a directory')
+	if args.output is None:
+		args.output = TemporaryDirectory(prefix='Email search results (id ', suffix=')').name
+	else:
+		if not os.path.exists(args.output) or not os.path.isdir(args.output):
+			parser.error(f'output path \'{args.output}\' does not exist or is not a directory')
 
-def validate_arguments():
-	if filter is None:
-		print('No filter specified')
-		sys.exit(3)
-	pass
+def construct_filter():
+	global filter
+	criteria: List[Filter] = []
+	keyword_filters = map(lambda k : BodyKeywordFilter(k, case_sensitive=args.casesensitive), args.keywords)
+	criteria.append(BooleanFilter(BooleanOperator.or_op if args.any else BooleanOperator.and_op, keyword_filters))
+	if getattr(args, 'from') is not None:
+		criteria.append(HeaderFilter('from', getattr(args, 'from')))
+	if getattr(args, 'to') is not None:
+		criteria.append(HeaderFilter(['to', 'cc', 'bcc'], getattr(args, 'to')))
+	if args.subject is not None:
+		criteria.append(HeaderFilter('subject', args.subject))
+	# TODO: Dates
+	filter = BooleanFilter(BooleanOperator.and_op, criteria)
 
 def handle_results():
 	"""Final logic after all searching is completed."""
@@ -215,6 +315,6 @@ def handle_results():
 
 # Main logic
 parse_arguments()
-validate_arguments()
+construct_filter()
 walk_directory(start_path)
 handle_results()