#!/usr/bin/env python

import StringIO
import __main__
import cStringIO
import email
import re
import sys
import zipfile

MAX_ENTRIES = 100
reject_extensions = ('.bat', '.com', '.exe', '.lnk', '.pif', '.scr',)
reject_patterns = map (lambda x: re.compile (r'\%s$' % x), reject_extensions)

def analyse_filename (name):
	for i in reject_patterns:
		if i.search (name):
			return "Filename matches forbidden pattern: %s" \
			       % i.pattern
	return None

def inspect_zip_attachment (msg):
	zip_str = email.base64MIME.decode (msg.get_payload ())
	zip_file = cStringIO.StringIO (zip_str)
	try:
		zip = zipfile.ZipFile (zip_file)
	except zipfile.BadZipfile:
		return "Not a zip archive"

	names = zip.namelist ()
	entry_count = len (names)
	if entry_count > MAX_ENTRIES:
		return "Zip archive has too many entries: %d" % entry_count

	name = zip.testzip ()
	if name:
		return "Zip archive contains corrupt file: %s" % name

	for i in names:
		r = analyse_filename (i)
		if r:
			return r


recurse_patterns = {
	'.zip$' : inspect_zip_attachment,
	}


def scan_multipart_message (msg):
	for i in msg.get_payload ():
		r = scan_message (i)
		if r:
			return r

def scan_message_part (msg):
	name = msg.get_filename ()
	if name:
		r = analyse_filename (name)
		if r:
			return r
		for i in recurse_patterns.keys ():
			if re.search (i, name):
				r = recurse_patterns[i] (msg)
				if r:
					return r

def scan_message (msg):
	if msg.is_multipart ():
		return scan_multipart_message (msg)
	else:
		return scan_message_part (msg)

def main ():
	str = sys.stdin.read ()
	msg = email.message_from_string (str)
	r = scan_message (msg)
	if r:
		v = 'suspect; reason: %s' % r
	else:
		v = 'clean'
	msg.add_header ('X-Filter-Zip-Scanned', v)
	sys.stdout.write (msg.as_string (unixfrom=1))

if __name__ == '__main__':
	main ()