Can now search various fields. Better command line parsing.

.pylintrc

@@ -0,0 +1,6 @@
+[MESSAGES CONTROL]
+disable=bad-indentation, invalid-name
+
+[FORMAT]
+
+indent-string=' '

search.py

@@ -3,16 +3,17 @@ import re
 import platform
 import subprocess
 from tempfile import TemporaryDirectory
-from typing import List, Union
+from typing import List, Optional, Union
 from enum import Enum
 from email.utils import parsedate
 from email.parser import BytesParser
 from email.message import EmailMessage
-from zipfile import ZipFile, ZipInfo
+from zipfile import ZipFile
 import sys
 import os
 
 class BooleanOperator(Enum):
+	"""Boolean combinatory operator enum."""
 	and_op = 1
 	or_op = 2
 
@@ -20,10 +21,13 @@ 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"
+		raise AssertionError("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"
+		raise AssertionError("Not implemented")
+	def supports_raw(self) -> bool:
+		"""Whether this filter supports raw messages at least partially."""
+		return False
 
 class BodyKeywordFilter(Filter):
 	"""Simple substring search filter."""
@@ -46,11 +50,14 @@ class BodyKeywordFilter(Filter):
 		if self.case_sensitive:
 			if self.keyword in raw_content:
 				return True
-			else:
-				if self.keyword.lower() in raw_content.lower():
-					return True
+		else:
+			if self.keyword.lower() in raw_content.lower():
+				return True
 		return False
 
+	def supports_raw(self) -> bool:
+		return True
+
 class HeaderFilter(Filter):
 	"""Matches a value in an email header. Can search one filter or multiple.
 	Header names case-insensitive; value is case-insensitive."""
@@ -95,11 +102,62 @@ class BooleanFilter(Filter):
 			return True
 		return False
 
-args: argparse.Namespace = None
-filter = None
+	def supports_raw(self) -> bool:
+		for subfilter in self.subfilters:
+			if subfilter.supports_raw():
+				return True
+		return False
+
+class DateFilter(Filter):
+	"""Filters messages based on the date field. For each message with a parseable
+	date field, the given comparator is called with a `maketime` list representation
+	of the date and time. The comparator must return a bool of whether to match
+	the given date or not."""
+	def __init__(self, comparator):
+		self.comparator = comparator
+
+	def matches(self, message: EmailMessage) -> bool:
+		date_str = message.get('date', None)
+		if date_str is None:
+			return False
+		date_elems = parsedate(date_str)
+		if date_elems is None:
+			return False
+		return self.comparator(date_elems)
+
+class Options:
+	"""Parsed command-line options."""
+	def __init__(self):
+		self.keywords: List[str] = []
+		self.any: bool = False
+		self.dir: List[str] = []
+		self.output: Optional[str] = None
+		self.casesensitive: bool = False
+		setattr(self, 'from', None)
+		setattr(self, 'to', None)
+		self.subject: Optional[str] = None
+		self.before: Optional[List[int]] = None
+		self.after: Optional[List[int]] = None
+		self.raw: bool = False
+
+args: Options = Options()
+message_filter: Filter = None
 result_count = 0
 zip_result_count = 0
 zip_count = 0
+parser: argparse.ArgumentParser = None
+
+def compare_dates(a: List[int], b: List[int]) -> int:
+	"""Compares two list representations of `maketime` date-times. Returns -1 if a < b,
+	1 if a > b, and 0 if they are equal."""
+	for i in range(6):
+		a_elem = a[i] if i < len(a) else -1
+		b_elem = b[i] if i < len(b) else -1
+		if a_elem < b_elem:
+			return -1
+		if a_elem > b_elem:
+			return 1
+	return 0
 
 def clean_filename(original: str) -> str:
 	"""Returns a scrubbed string with safe filename characters."""
@@ -108,7 +166,7 @@ def clean_filename(original: str) -> str:
 def filename_from_email(email: EmailMessage) -> str:
 	"""Creates a safe filename to save the given email to."""
 	filename = ''
-	date_str = email.get('Date', None)
+	date_str = email.get('date', None)
 	if date_str is not None:
 		parsed_date = parsedate(date_str)
 		if parsed_date is not None:
@@ -119,7 +177,8 @@ def filename_from_email(email: EmailMessage) -> str:
 			filename += '0000-00-00T00.00.00 - '
 	else:
 		filename += '0000-00-00T00.00.00 - '
-	subject = email.get('Subject')
+	filename += f'{result_count:04} - '
+	subject = email.get('subject')
 	if subject is not None:
 		filename += clean_filename(subject)[0:50].strip()
 	else:
@@ -127,15 +186,15 @@ def filename_from_email(email: EmailMessage) -> str:
 	filename += '.eml'
 	return filename
 
-def walk_directory(dir: str) -> None:
+def walk_directory(path: str) -> None:
 	"""Spiders a directory looking for subdirectories and email zip archives."""
 	global zip_count
-	for f in os.listdir(dir):
-		full_path = dir + os.sep + f
+	for f in os.listdir(path):
+		full_path = path + os.sep + f
 		if f.lower().endswith('.zip'):
 			zip_count += 1
 			process_zip_file(full_path)
-		if os.path.isdir(f):
+		if os.path.isdir(full_path):
 			walk_directory(full_path)
 
 def process_zip_file(zip_path: str) -> None:
@@ -151,64 +210,83 @@ def process_zip_file(zip_path: str) -> None:
 			parser = BytesParser()
 			try:
 				email = parser.parsebytes(data)
-				search_content(email, zip_path, entry)
-			except UnicodeError:
-				print('Unicode error in message. Searching raw content.', file=sys.stderr)
+				search_content(email)
 			except:
-				print('Error parsing message. Searching raw content.', file=sys.stderr)
+				if message_filter.supports_raw():
+					search_raw_content(data)
+				else:
+					print('Message cannot be parsed. Skipping.')
 	if zip_result_count > 0:
 		print(f"\t{zip_result_count} results in zip")
 
-def search_content(email: EmailMessage, zip_path: str, entry: ZipInfo) -> None:
+def search_content(email: EmailMessage) -> None:
 	"""Processes an email message in a zip file."""
 	global result_count, zip_result_count
-	if filter.matches(email):
-		if not os.path.exists(output_path):
-			os.makedirs(output_path)
-		with open(output_path + os.sep + filename_from_email(email), 'wb') as f:
-			result_count += 1
-			zip_result_count += 1
-			f.write(email.as_bytes())
+	if message_filter.matches(email):
+		save_message(email)
 
-def search_raw_content(raw_bytes: bytes, zip_path: str, entry: ZipInfo) -> None:
+def search_raw_content(raw_bytes: bytes) -> None:
 	global result_count, zip_result_count
-	try:
-		content = raw_bytes.decode('iso-8859-1', errors='ignore')
-	except:
+	encodings = [ 'ascii', 'iso-8859-1', 'utf-8' ]
+	content = None
+	for encoding in encodings:
 		try:
-			content = raw_bytes.decode('utf-8', errors='ignore')
+			content = raw_bytes.decode(encoding)
+			break
 		except:
-			print('Cannot decode email bytes. Skipping.', file=sys.stderr)
+			pass
+	if content is None:
+		print('Cannot decode email bytes. Skipping message.', file=sys.stderr)
+		return
+	print('Could not parse message. Searching raw content.', file=sys.stderr)
+	if message_filter.matches_raw(content):
+		save_raw_message(raw_bytes)
+
+def save_message(email: EmailMessage) -> None:
+	"""Saves a matching message to the results directory."""
+	global result_count, zip_result_count
+	if not os.path.exists(args.output):
+		os.makedirs(args.output)
+	with open(args.output + os.sep + filename_from_email(email), 'wb') as f:
+		result_count += 1
+		zip_result_count += 1
+		f.write(email.as_bytes())
+
+def save_raw_message(content: bytes) -> None:
+	"""Saves an unparseable matching message to the results directory."""
+	global result_count, zip_result_count
+	if not os.path.exists(args.output):
+		os.makedirs(args.output)
+	filename = f'unparseable-match-{result_count:04}.eml'
+	with open(args.output + os.sep + filename, 'wb') as f:
+		result_count += 1
+		zip_result_count += 1
+		f.write(content)
 
 def parse_arguments():
 	"""Parses command-line arguments to `args`."""
-	global args
+	global args, parser
+	# TODO: Revisit raw mode and how unparseable emails should be handled
 	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.'
+		epilog='If raw mode is enabled, any messages that cannot be decoded ' + \
+			'will be searched as raw text.'
 	)
 	parser.add_argument(
 		'keywords',
 		action='append',
-		nargs='+',
+		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)'
+		help='matches messages containing any of the given search phrases (default requires ' + \
+			'all phrases appear in a message)'
 	)
 	parser.add_argument(
 		'-d', '--dir',
@@ -227,10 +305,12 @@ def parse_arguments():
 	)
 	parser.add_argument(
 		'-f', '--from',
+		metavar='sender-email',
 		help='email address of sender'
 	)
 	parser.add_argument(
 		'-t', '--to',
+		metavar='recipient-email',
 		help='email address of recipient (searches to:, cc:, bcc: fields)'
 	)
 	parser.add_argument(
@@ -251,10 +331,14 @@ def parse_arguments():
 		'-r', '--raw',
 		default=False,
 		action='store_true',
-		help='searches raw email content (see below)'
+		help='allows searching unparseable messages as raw text'
 	)
 	args = parser.parse_args()
 
+def validate_arguments():
+	"""Validate and parse special field types"""
+	global args
+	args.keywords = args.keywords[0]  # no idea why it nests it 2D
 	if args.before is not None:
 		m = re.match('^([0-9]{4})-([0-9]{2})-([0-9]{2})$', args.before)
 		if m is None:
@@ -285,26 +369,36 @@ def parse_arguments():
 			parser.error(f'output path \'{args.output}\' does not exist or is not a directory')
 
 def construct_filter():
-	global filter
+	global message_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))
+	keyword_filters = []
+	for k in args.keywords:
+		k = k.strip()
+		if len(k) > 0:
+			keyword_filters.append(BodyKeywordFilter(k, case_sensitive=args.casesensitive))
+	if len(keyword_filters) > 0:
+		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)
+	if args.before is not None:
+		criteria.append(DateFilter(lambda d: compare_dates(d, args.before) <= 0))
+	if args.after is not None:
+		criteria.append(DateFilter(lambda d: compare_dates(d, args.after) >= 0))
+	if len(criteria) == 0:
+		parser.error('No filters specified')
+	message_filter = BooleanFilter(BooleanOperator.and_op, criteria)
 
 def handle_results():
 	"""Final logic after all searching is completed."""
 	if result_count > 0:
 		if platform.system() == 'Darwin':
-			subprocess.call(['open', output_path])
+			subprocess.call(['open', args.output])
 		elif platform.system() == 'Windows':
-			subprocess.call(['explorer.exe', output_path])
+			subprocess.call(['explorer.exe', args.output])
 		print(f'Found {result_count} result(s) total')
 	elif zip_count == 0:
 		print('No zip files found')
@@ -315,6 +409,8 @@ def handle_results():
 
 # Main logic
 parse_arguments()
+validate_arguments()
 construct_filter()
-walk_directory(start_path)
+for path in args.dir:
+	walk_directory(path)
 handle_results()