#!/usr/bin/env python
"""Data structures and functions used to support work with mime objects."""

import sys
import unittest
import zipfile
from io import BytesIO

try:
	from StringIO import StringIO					#Python 2
except ImportError:
	from io import StringIO						#Python 3



mime_support_version = '0.7'

#Default extension to use when a particular mime type is found.  Not intended to be used directly, use ext_of_mime instead.
default_ext_of_mime = {
	'Composite Document File V2 Document': 'doc',
	'application/CDFV2': 'wps',
	'application/epub+zip': 'epub',
	'application/java-archive': 'jar',
	'application/javascript': 'js',
	'application/mac-binhex40': 'hqx',
	'application/msword': 'doc',
	'application/ogg': 'ogg',
	'application/pdf': 'pdf',
	'application/pgp': 'gpg',
	'application/pgp-keys': 'aexpk',
	'application/pkcs7-signature': 'p7',
	'application/postscript': 'ps',
	'application/vnd.ms-cab-compressed': 'cab',
	'application/vnd.ms-excel': 'xls',
	'application/vnd.ms-excel.addin.macroEnabled.12': 'xlam',
	'application/vnd.ms-excel.sheet.binary.macroEnabled.12': 'xlsb',
	'application/vnd.ms-excel.sheet.macroEnabled.12': 'xlsm',
	'application/vnd.ms-excel.template.macroEnabled.12': 'xltm',
	'application/vnd.ms-fontobject': 'eot',
	'application/vnd.ms-opentype': 'otf',
	'application/vnd.ms-powerpoint': 'ppt',
	'application/vnd.ms-powerpoint.addin.macroEnabled.12': 'ppam',
	'application/vnd.ms-powerpoint.presentation.macroEnabled.12': 'pptm',
	'application/vnd.ms-powerpoint.slideshow.macroEnabled.12': 'ppsm',
	'application/vnd.ms-powerpoint.template.macroEnabled.12': 'potm',
	'application/vnd.ms-tnef': 'dat',
	'application/vnd.ms-word.document.macroEnabled.12': 'docm',
	'application/vnd.ms-word.template.macroEnabled.12': 'dotm',
	'application/vnd.oasis.opendocument.presentation': 'odp',
	'application/vnd.oasis.opendocument.spreadsheet': 'ods',
	'application/vnd.oasis.opendocument.text': 'odt',
	'application/vnd.openxmlformats-officedocument.presentationml.presentation': 'pptx',
	'application/vnd.openxmlformats-officedocument.presentationml.slideshow': 'ppsx',
	'application/vnd.openxmlformats-officedocument.presentationml.template': 'potx',
	'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet': 'xlsx',
	'application/vnd.openxmlformats-officedocument.spreadsheetml.template': 'xltx',
	'application/vnd.openxmlformats-officedocument.wordprocessingml.document': 'docx',
	'application/vnd.openxmlformats-officedocument.wordprocessingml.template': 'dotx',
	'application/vnd.tcpdump.pcap': 'pcap',
	'application/x-7z-compressed': '7z',
	'application/x-apple-diskimage': 'dmg',
	'application/x-arc': 'arc',
	'application/x-arj': 'arj',
	'application/x-bittorrent': 'torrent',
	'application/x-bzip2': 'bz2',
	'application/x-chrome-extension': 'crx',
	'application/x-coredump': 'core',
	'application/x-cpio': 'cpio',
	'application/x-dmg': 'dmg',
	'application/x-dosexec': 'exe',
	'application/x-elc': 'elc',
	'application/x-empty': 'empty',
	'application/x-fortran': 'f',
	'application/x-gzip': 'gz',
	'application/x-executable': 'exe',
	'application/x-font-ttf': 'ttf',
	'application/x-hwp': 'hwp',
	'application/x-ima': 'ima',
	'application/x-iso9660-image': 'iso',
	'application/x-java-applet': 'jar',
	'application/x-javascript': 'js',
	'application/x-lha': 'lha',
	'application/x-lzh-compressed': 'lzh',
	'application/x-mach-binary': 'macbin',
	'application/x-ms-wmz': 'wmz',
	'application/x-msdownload': 'exe',
	'application/x-msi': 'msi',
	'application/x-pem': 'pem',
	'application/x-rar': 'rar',
	'application/x-rpm': 'rpm',
	'application/x-sharedlib': 'lib',
	'application/x-shockwave-flash': 'swf',
	'application/x-stuffit': 'sit',
	'application/x-tar': 'tar',
	'application/x-xar': 'pkg',
	'application/x-xz': 'xz',
	'application/zip': 'zip',
	'application/zlib': 'dmg',
	'audio/midi': 'mid',
	'audio/mp4': 'm4a',
	'audio/mpeg': 'mp3',
	'audio/x-m4a': 'm4a',
	'audio/x-wav': 'wav',
	'image/gif': 'gif',
	'image/jpeg': 'jpg',
	'image/jxr': 'jxr',
	'image/pjpeg': 'jpg',
	'image/tiff': 'tif',
	'image/png': 'png',
	'image/svg+xml': 'svg',
	'image/vnd.adobe.photoshop': 'psd',
	'image/webp': 'webp',
	'image/x-icns': 'icns',
	'image/x-icon': 'ico',
	'image/x-ms-bmp': 'bmp',
	'image/x-pcx': 'pcx',
	'image/x-tga': 'tga',
	'inode/x-empty': 'empty',
	'message/rfc822': 'eml',
	'text/calendar': 'ics',
	'text/javascript': 'js',
	'text/json': 'json',
	'text/PGP': 'pgp',
	'text/rtf': 'rtf',
	'text/rss': 'rss',
	'text/troff': 'troff',
	'text/v-card': 'vcf',
	'text/vcard': 'vcf',
	'text/x-asm': 'asm',
	'text/x-c': 'c',
	'text/x-c++': 'cpp',
	'text/x-diff': 'diff',
	'text/x-java': 'java',
	'text/x-lua': 'lua',
	'text/x-m4': 'm4',
	'text/x-makefile': 'make',
	'text/x-msdos-batch': 'bat',
	'text/x-pascal': 'pas',
	'text/x-perl': 'pl',
	'text/x-php': 'php',
	'text/x-python': 'py',
	'text/x-ruby': 'rb',
	'text/x-shellscript': 'sh',
	'text/x-tex': 'tex',
	'video/3gpp': '3g2',
	'video/mp4': 'mp4',
	'video/mpeg': 'mpg',
	'video/quicktime': 'mov',
	'video/x-flv': 'flv',
	'video/x-ms-asf': 'asf',
	'video/x-msvideo': 'avi'
	}
#Notes:

#	'application/msword': 'doc'			#Could also be .dot
#	'application/pgp': 'gpg'			#Could also be .asc or .pgp
#	'application/vnd.ms-excel': 'xls'		#Could also be .xlt or .xla
#	'application/vnd.ms-powerpoint'): 'ppt'		#Could also be .pot, .pps, or .ppa
#	'application/x-dosexec'): 'exe'			#Should one of these be .com?
#	'application/x-gzip'): 'gz'			#Note - this catches .tgz as well
#	'application/x-executable'): 'exe'		#Should one of these be .com?
#	'text/PGP'): 'pgp'				#Could also be .gpg
#	'text/x-pascal'): 'pas'				#Lots of these come back for text files.


#Default extension to use, but if found in this list check against the pri and sec overrides list for overrides
mime_check_overrides = {
                       'application/octet-stream': 'bin',
                       'application/xml': 'xml',
                       'audio/x-ape': 'ape',
                       'text/html': 'htm',
                       'text/plain': 'txt',
                       'text/xml': 'xml'
                       }


#If a file's mime-type doesn't correctly identify the extension, find a string that
#better identifies this file and the extension you feel it should have and place it here.
sec_overrides = {'new ActiveXObject(': 'js',
                 '("WScript.Shell")': 'js',
                 '("ADODB.Stream")': 'js',
                 '("Scripting.FileSystemObject")': 'js',
                 'document.writeln("': 'js',
                 'WScript.CreateObject(': 'js',
                 'WScript.Echo("': 'js',
                 "<script language='JScript'>": 'js',
                 '} catch(': 'js',
                 '<script language="JScript">': 'js',
                 "<script  language='JScript'>": 'js',
                 'Math.floor(': 'js',
                 ".split('').reverse().join('')": 'js',
                 '= new Function(': 'js',
                 'Math.pow(': 'js',
                 '} catch (': 'js',
                 'eval(unescape(': 'js',
                 'document.location=': 'js',
                 'MSXML2.XMLHTTP': 'js',
                 'new Function(': 'js',
                 '<script>function': 'js',
                 'String.fromCharCode(': 'js',
                 'WScript[': 'js',
                 '#!/usr/bin/perl': 'pl',
                 '.charCodeAt(': 'js',
                 ' = eval(': 'js',
                 'begin 6': 'uue',
                 ' = CreateObject("WScript"': 'js',
                 '/*@cc_on': 'js',
                 'Non-System disk or disk error': 'img'}


#This is reserved for things we should look for before looking for the entries in sec_overrides
#It's especially useful when a file might have more than one signature string, and we want certain
#ones to take precedence over others; they go in here.
pri_overrides = {'<job id=': 'wsf',
                 '<job  id=': 'wsf'}


def debug(message):
	"""Write debug message to stderr"""
	sys.stderr.write(message + '\n')


def content_matches(content_string, search_file, match_strings):
	"""This looks for any of the match_strings in either content_string (a string variable) or search_file (a filename).
	If any lines in the file contain any of the search strings, send back the first search string found, otherwise return None.
	"""

	if content_string is None:					#This must be a literal check for == None, as opposed to "if not content_string:" because content_string could be empty but not None
		with open(search_file, 'r') as searchcontent:
			for line in searchcontent:
				for signature in match_strings:		#Match_strings is a dictionary; this extracts the keys, which are the strings for which we're searching.
					if signature in line:
						#debug("!!!! Found " + signature + " in " + search_file)
						return signature
	else:
		with StringIO(str(content_string)) as searchcontent:
			for line in searchcontent:
				for signature in match_strings:		#Match_strings is a dictionary; this extracts the keys, which are the strings for which we're searching.
					if signature in line:
						#debug("!!!! Found " + signature + " in " + search_file)
						return signature
	return None


def check_overrides(content_string, content_file, pri_dict, sec_dict, default_extension):
	"""Check to see if we need to override the extension suggested by the mime type."""

	first_match = content_matches(content_string, content_file, pri_dict)
	if first_match is not None:
		return pri_dict[first_match]			#We found a match, return the override extension
	else:
		second_match = content_matches(content_string, content_file, sec_dict)
		if second_match is not None:
			return sec_dict[second_match]			#We found a match, return the override extension
		else:
			return default_extension


def forced_zip_ext(zip_file_list):
	"""Figures out if the zip file in question is a compound doc (whose top level container is a zip), such as jar or opendoc format.
	Hand in the file content list of the zip rather than the zip itself.
	Returns None if a normal zip, and the appropriate extension to use if a compound doc.
	"""

	force_extension = None
	if 'META-INF/MANIFEST.MF' in zip_file_list:
		force_extension = 'jar'
	elif '[Content_Types].xml' in zip_file_list:
		if 'word/vbaProject.bin' in zip_file_list:
			force_extension = 'docm'
		elif 'word/document.xml' in zip_file_list:
			force_extension = 'doc'
		elif 'xl/vbaProject.bin' in zip_file_list:
			force_extension = 'xlsm'
		elif 'xl/workbook.xml' in zip_file_list:
			force_extension = 'xls'
		elif 'ppt/vbaProject.bin' in zip_file_list:
			force_extension = 'pptm'
		elif 'ppt/presentation.xml' in zip_file_list:
			force_extension = 'ppt'

	return force_extension


def zip_override(content_string, content_file, def_extension):
	"""If this is a compound document, return the correct extension to use.  Otherwise, send back def_extension (which was almost certainly handed down to us as 'zip'."""

	if content_string is None:
		#We assume we were handed a filename pointing at a zip
		if zipfile.is_zipfile(content_file):
			handle_created = True
			try:
				z_h = zipfile.ZipFile(content_file, 'r')
			except OSError:
				handle_created = False
			if handle_created:
				z_list = z_h.namelist()				#a python list of all files and dirs in the zip
				ext_to_force = forced_zip_ext(z_list)
				if ext_to_force:
					return ext_to_force
				z_h.close()
	else:
		#We've been handed a variable (content_string) that we assume is zip content
		zip_pseudo_handle = BytesIO(content_string)
		if zipfile.is_zipfile(zip_pseudo_handle):
			handle_created = True
			try:
				z_h = zipfile.ZipFile(zip_pseudo_handle, 'r')
			except OSError:
				handle_created = False
			if handle_created:
				z_list = z_h.namelist()				#a python list of all files and dirs in the zip
				ext_to_force = forced_zip_ext(z_list)
				if ext_to_force:
					return ext_to_force
				z_h.close()
		zip_pseudo_handle.close()

	return def_extension


def ext_of_mime(mime_type, content_string, content_file):
	"""When handed a mime type of the form 'category/specific', this returns the most likely extension for that mime type.
	If there's no exact match for the provided mime type, we see if that mime type _starts_ with any of the listed mime types
	(for example, ext_of_mime('application/zip;compressed') will return 'zip'.  Sample call:

	obj_mime = 'text/html'
	file_ext = ext_of_mime(obj_mime)
	print file_ext

	returns 'htm' .

	If we need to inspect the actual content (zip files/compound documents, or text files that could be .txt or .js, for example),
	we first look at content_string.  If that's None, we fall back on content_file (which is the filename, not a filehandle).
	"""

	ret_extension = 'zzz'

	if mime_type.startswith('application/zip'):
		ret_extension = 'zip'
		ret_extension = zip_override(content_string, content_file, ret_extension)		#This replaces 'zip' with the correct extension if content_string/content_file is a compound doc.
	elif mime_type.startswith(tuple(mime_check_overrides.keys())):
		#debug('mime_type In mime_check_overrides')
		for one_type in mime_check_overrides:
			if mime_type.startswith(one_type):
				ret_extension = mime_check_overrides[one_type]
				#debug('ret_extension is: ' + str(ret_extension))
				ret_extension = check_overrides(content_string, content_file, pri_overrides, sec_overrides, ret_extension)
				#debug('ret_extension now: ' + str(ret_extension))
				break
	elif mime_type.startswith(tuple(default_ext_of_mime.keys())):
		#debug('mime_type In default_ext_of_mime')
		for one_type in default_ext_of_mime:
			if mime_type.startswith(one_type):
				ret_extension = default_ext_of_mime[one_type]
				break
	else:
		debug('Unknown file extension for mime:' + str(mime_type))
	#FIXME - deleteme
	#Following no longer needed, we check startswith for both dictionaries above
	#else:							#No perfect match found, lets see if our supplied string starts with one of these types.
	#	combined = {}
	#	combined.update(default_ext_of_mime)
	#	combined.update(mime_check_overrides)
	#	for one_type in combined:
	#		if mime_type.startswith(one_type):
	#			ret_extension = combined[one_type]
	#			break

	return ret_extension


class TestMimeExtension(unittest.TestCase):
	"""Unit tests for the mimesupport library."""
	def test_simple_membership(self):
		"""Check returns for simple mime types."""
		self.assertEqual(ext_of_mime('text/html', '', None), 'htm')
		self.assertEqual(ext_of_mime('application/xml', '', None), 'xml')

	def test_mime_startswith(self):
		"""Check returns for mime types that match but have additional text."""
		self.assertEqual(ext_of_mime('text/html;and more', '', None), 'htm')
		self.assertEqual(ext_of_mime('application/zip;compressed', '', None), 'zip')

if __name__ == "__main__":
	unittest.main()