# Unity Mail, the main application class
# Author: Dmitry Shachnev <mitya57@gmail.com>
# License: GNU GPL 3 or higher; http://www.gnu.org/licenses/gpl.html

app_name = 'unity-mail'
app_version = '1.4.1'

import email
import email.errors
import email.header
import email.utils
import imaplib
import logging
import os
import secretstorage
import shutil
import signal
import subprocess
import sys
import time

from socket import error as socketerror
from gi.repository import GLib, Gtk, Notify

from UnityMail import gettext, Unity, Dbusmenu, MessagingMenu, \
 with_unity, with_messagingmenu
from UnityMail.config import ConfigFile, config_dir, open_url_or_command
from UnityMail.dialog import PreferencesDialog

_ = gettext.gettext

logger = logging.getLogger('Unity-Mail')
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter('%(name)s: %(levelname)s: %(message)s'))
logger.addHandler(handler)
logger.propagate = False

def print_error(error, text=None):
	msg = text + ': exception occured: ' if text else 'Exception occured: '
	logger.error(msg + str(error))

def decode_wrapper(header):
	"""Decodes an Email header, returns a string"""
	# A hack for headers without a space between decoded name and email
	try:
		dec = email.header.decode_header(header.replace('=?=<', '=?= <'))
	except email.errors.HeaderParseError:
		logger.warning('Exception in decode, skipping')
		return header
	parts = []
	for dec_part in dec:
		if dec_part[1]:
			try:
				parts.append(dec_part[0].decode(dec_part[1]))
			except (AttributeError, LookupError, UnicodeDecodeError):
				logger.warning('Exception in decode, skipping')
		elif isinstance(dec_part[0], bytes):
			parts.append(dec_part[0].decode())
		else:
			parts.append(dec_part[0])
	return str.join(' ', parts)

def get_header_wrapper(message, header_name, decode=False):
	header = message[header_name]
	if isinstance(header, str):
		header = header.replace(' \r\n', '').replace('\r\n', '')
		return (decode_wrapper(header) if decode else header)
	return ''

def get_sender_name(sender):
	"""Strips address, and returns only name"""
	sname = email.utils.parseaddr(sender)[0]
	return sname if sname else sender

fix_format = lambda string: string.replace('%(t0)s', '{t0}').replace('%(t1)s', '{t1}')

class TimeoutException(Exception):
	"""Timeout exception for internal use"""
	pass

class Message(object):
	"""Message object"""
	def __init__(self, account_id, account_position, title, folder, message_id,
		         timestamp):
		self.is_active = True
		self.account_id = account_id
		self.account_position = account_position
		self.title = title
		self.folder = folder
		self.message_id = message_id
		self.timestamp = timestamp

class UnityMail(object):
	"""Main Unity Mail Application"""
	
	def __init__(self):
		self.dicts = []
		self.host = []
		self.port = []
		self.login = []
		self.passwd = []
		self.accdlg = None
		self.init_argv_help()
		self.init_keyring()
		self.init_autostart()
		if not self.dialogshown:
			self.init_argv_change()
		self.init_config()
		Notify.init('Unity Mail')
		self.l = len(self.host)
		self.mail_client = [None]*self.l
		self.notinit = [True]*self.l
		self.unread_messages = []
		if with_unity:
			desktop_id = 'unity-mail.desktop'
			favorites = Unity.LauncherFavorites.get_default().enumerate_ids()
			for i in favorites:
				if 'unity-mail' in i: desktop_id = i
			self.launcher = Unity.LauncherEntry.get_for_desktop_id(desktop_id)
			root = Dbusmenu.Menuitem()
			item_markread = Dbusmenu.Menuitem()
			item_markread.property_set('label', _('Mark all as read'))
			item_markread.connect('item-activated', self.on_mark_all_read)
			separator = Dbusmenu.Menuitem()
			separator.property_set('type', 'separator')
			root.child_append(item_markread)
			root.child_append(separator)
			self.launcher.set_property('quicklist', root)
		self.first_run = True
		if with_messagingmenu:
			self.mmapp = MessagingMenu.App(desktop_id='unity-mail.desktop')
			self.mmapp.register()
			self.mmapp.connect('activate-source', self.on_mm_item_clicked)
		
		self.loop = GLib.MainLoop()
		GLib.set_application_name('unity-mail')
		GLib.timeout_add_seconds(self.interval, self.update)
		self.update()
		self.start_loop()
	
	def start_loop(self):
		try:
			self.loop.run()
		except KeyboardInterrupt:
			self.handle_keyboard_interrupt()
	
	def handle_keyboard_interrupt(self):
		signal.alarm(0)
		try:
			if input('\nDo you really want to exit (y/n)? ') in ('y', 'yes'):
				sys.exit()
		except EOFError:
			print()
		except KeyboardInterrupt:
			print()
			sys.exit()
		self.start_loop()
	
	def on_mm_item_clicked(self, mmapp, sourceid):
		for message in self.unread_messages:
			if message.message_id == sourceid:
				urlid = self.get_urlid(message.account_id)
				if urlid:
					open_url_or_command(urlid)
				else:
					self.mark_message_as_read(message)
				if message.is_active and with_unity:
					count = self.launcher.get_property('count')
					self.set_launcher_count(count-1)
	
	def mark_message_as_read(self, message):
		client = self.mail_client[message.account_id]
		try:
			if message.folder:
				client.select(message.folder)
			else:
				client.select()
			client.store(message.account_position, '+FLAGS', '\Seen')
			client.close()
		except (imaplib.IMAP4.error, socketerror) as e:
			print_error(e)
	
	def on_mark_all_read(self, menuitem, timestamp):
		for message in self.unread_messages:
			if message.is_active:
				self.mark_message_as_read(message)
			if with_messagingmenu and self.mmapp.has_source(message.message_id):
				self.mmapp.remove_source(message.message_id)
		self.set_launcher_count(0)
	
	def init_config(self):
		config = ConfigFile(config_dir+'/unity-mail.conf')
		self.interval = config.get_int_with_default('Interval', 30)
		self.max_count = config.get_int_with_default('MaxItemCount', 20)
		self.enable_notifications = config.get_bool_with_default('EnableNotifications', True)
		self.play_sound = config.get_bool_with_default('EnableSound', False)
		self.display_icons = config.get_bool_with_default('NotificationsHaveIcons', False)
		self.draw_attention = config.get_bool_with_default('DrawAttention', True)
		self.hide_count = config.get_bool_with_default('HideMessagesCount', False)
		
		self.command = config.get_str_with_default('ExecOnReceive')
		self.custom_sound = config.get_str_with_default('CustomSound')
		self.account_display = config.get_str_with_default('AccountDisplay', None)
		value = config.get_str_with_default('Blacklist')
		self.black_list = value.split(', ') if value else []
		value = config.get_str_with_default('ExtraFolders')
		self.extra_folders = value.split(', ') if value else []
		
		self.on_click_urls = []
		if config.has_section('URLs'):
			self.on_click_urls = config.items('URLs')
	
	def init_autostart(self):
		as_dir = config_dir+'/autostart/'
		if os.path.exists('/usr/share/unity-mail/unity-mail-autostart.desktop') and \
		not os.path.exists(as_dir+'unity-mail-autostart.desktop'):
			logger.info('Adding to Autostart')
			if not os.path.exists(as_dir):
				os.makedirs(as_dir)
			shutil.copy('/usr/share/unity-mail/unity-mail-autostart.desktop',
			as_dir+'unity-mail-autostart.desktop')
	
	def init_argv_help(self):
		if len(sys.argv) > 1:
			if sys.argv[1] == '--help' or sys.argv[1] == '-h':
				print('Unity Mail, version', app_version)
				print('Usage:')
				print('  unity-mail [options]')
				print('Options:')
				print('  -c, --change  : Change accounts data')
				print('  -h, --help    : Display this help message and exit')
				print('  -v, --version : Display version number and exit')
				sys.exit(0)
			if sys.argv[1] == '--version' or sys.argv[1] == '-v':
				print('Unity Mail, version', app_version)
				sys.exit(0)
	
	def init_argv_change(self):
		if len(sys.argv) > 1:
			if sys.argv[1] == '--change' or sys.argv[1] == '-c':
				self.open_dialog(set_dicts=True)
	
	def init_keyring(self):
		bus = secretstorage.dbus_init()
		try:
			self.collection = secretstorage.get_default_collection(bus)
			self.collection.is_locked()
		except secretstorage.SecretStorageException as e:
			sys.exit(e)
		if self.collection.is_locked():
			self.collection.unlock()
		if self.collection.is_locked():
			sys.exit('unity-mail: Failed to unlock the collection; exiting')
		search_attrs = {'application': 'unity-mail'}
		self.mail_keys = list(self.collection.search_items(search_attrs))
		self.dialogshown = False
		if not self.mail_keys:
			self.open_dialog()
			self.dialogshown = True
		for key in sorted(self.mail_keys, key=secretstorage.Item._item_id):
			attributes = key.get_attributes()
			self.host.append(attributes['server'])
			self.port.append(int(attributes['port']))
			self.login.append(attributes['username'])
			self.passwd.append(key.get_secret().decode('utf-8'))
			self.dicts.append({'Host': self.host[-1], 'Port': self.port[-1],
			'Login': self.login[-1], 'Passwd': self.passwd[-1]})
	
	def create_keyring_item(self, ind, update=False):
		attrs = {
			'application': 'unity-mail',
			'service': 'imap',
			'server': self.host[ind],
			'port': str(self.port[ind]),
			'username': self.login[ind]
		}
		label = 'unity-mail: '+self.login[ind]+' at '+self.host[ind]
		if update:
			self.mail_keys[ind].set_attributes(attrs)
			self.mail_keys[ind].set_secret(self.passwd[ind])
			self.mail_keys[ind].set_label(label)
		else:
			self.collection.unlock()
			self.collection.create_item(label, attrs, self.passwd[ind], True)
	
	def open_dialog(self, set_dicts=False):
		self.accdlg = PreferencesDialog()
		self.accdlg.connect('response', self.on_dialog_response)
		if set_dicts:
			self.accdlg.set_dicts(self.dicts)
		self.accdlg.run()
	
	def on_dialog_response(self, dlg, response):
		if response == Gtk.ResponseType.APPLY:
			dlg.update_datadicts()
			dlg.save_all_settings()
			self.dicts = dlg.datadicts
			self.load_data_from_dicts()
			for index in range(len(self.host), len(self.mail_keys)):
				# Remove old keys
				self.mail_keys[index].delete()
			for index in range(len(self.host)):
				# Create new keys or update existing
				self.create_keyring_item(index,
				update=(index < len(self.mail_keys)))
		dlg.destroy()
		if (not self.dicts) or (not self.host):
			sys.exit('unity-mail: No accounts registered; exiting')
	
	def load_data_from_dicts(self):
		self.host = []
		self.port = []
		self.login = []
		self.passwd = []
		for dict_ in self.dicts:
			if not dict_['Login']:
				continue
			self.host.append(dict_['Host'])
			try:
				self.port.append(int(dict_['Port']))
			except ValueError:
				# Port empty or non-integer
				self.port.append(993)
			self.login.append(dict_['Login'])
			self.passwd.append(dict_['Passwd'])
	
	def establish_connection(self, cn):
		try:
			self.mail_client[cn] = imaplib.IMAP4_SSL(self.host[cn], int(self.port[cn]))
		except TimeoutException:
			raise
		except Exception:
			self.mail_client[cn] = imaplib.IMAP4(self.host[cn], int(self.port[cn]))
			self.mail_client[cn].starttls()
	
	def try_establish_connection(self, cn):
		try:
			self.establish_connection(cn)
		except Exception:
			return False
		try:
			self.mail_client[cn].login(self.login[cn], self.passwd[cn])
		except (imaplib.IMAP4.error, UnicodeEncodeError) as e:
			print_error(e)
			logger.critical('Invalid account data provided for connection #{0}; exiting'.format(cn))
			sys.exit(1)
		else:
			logger.info('Connection #{0} established'.format(cn))
			self.notinit[cn]=False
			return True
	
	def raise_error(self, index, frame):
		raise TimeoutException('Timeout error')
	
	def update_single(self, cn, folder_name=None):
		if self.notinit[cn]:
			signal.alarm(15)
			if not self.try_establish_connection(cn):
				return
		if folder_name:
			status = self.mail_client[cn].select(mailbox=folder_name, readonly=True)
			if status[0] != 'OK':
				return
		else:
			self.mail_client[cn].select(readonly=True)
		search = self.mail_client[cn].search(None, '(UNSEEN)')
		if search[1][0] == None:
			ui = []
		else:
			ui = search[1][0].split()
		self.unread_count[cn] = len(ui)
		for m in ui:
			signal.alarm(3)
			typ, msg_data = self.mail_client[cn].fetch(m, '(BODY.PEEK[HEADER])')
			msg = None
			for response_part in msg_data:
				if isinstance(response_part, tuple):
					msg = email.message_from_bytes(response_part[1])
			if msg is None:
				continue
			message_id = msg['Message-Id']
			if not isinstance(message_id, str):
				message_id = str(cn)+':'+str(m)
			existing_message = None
			for message in self.unread_messages:
				if message.account_id == cn and message.message_id == message_id:
					existing_message = message
			if existing_message == None:
				sender = get_header_wrapper(msg, 'From', decode=True)
				subj = get_header_wrapper(msg, 'Subject', decode=True)
				date = get_header_wrapper(msg, 'Date')
				try:
					tuple_time = email.utils.parsedate_tz(date)
					timestamp = email.utils.mktime_tz(tuple_time)
					if timestamp > time.time():
						# Message time is larger than the current one
						timestamp = time.time()
				except TypeError:
					# Failed to get time from message
					timestamp = time.time()
				# Number of seconds to number of microseconds
				timestamp *= (10**6)
				while subj.lower().startswith('re:'):
					subj = subj[3:]
				while subj.lower().startswith('fwd:'):
					subj = subj[4:]
				subj = subj.strip()
				if sender.startswith('"'):
					pos = sender[1:].find('"')
					if pos >= 0:
						sender = sender[1:pos+1]+sender[pos+2:]
				ilabel = subj if subj else _('No subject')
				message = Message(account_id=cn, account_position=m, title=ilabel,
					folder=folder_name, message_id=message_id, timestamp=timestamp)
				self.unread_messages.append(message)
				if self.enable_notifications:
					self.notifications_queue.append((sender, subj, cn))
				if self.command and not self.first_run:
					try:
						subprocess.call((self.command, sender, ilabel))
					except OSError as e:
						# File doesn't exist or is not executable
						print_error(e, 'Cannot execute command')
			else:
				existing_message.is_active = True
		self.mail_client[cn].close()
	
	def check_email(self, cn, email):
		if self.login[cn] == email:
			return True
		if '@' not in email:
			return False
		login, host = email.rsplit('@', 1)
		return (self.login[cn] == login) and self.host[cn].endswith(host)
	
	def get_urlid(self, cn):
		for urlid, url in self.on_click_urls:
			if urlid.startswith('Inbox[') and urlid.endswith(']'):
				if self.check_email(cn, urlid[6:-1]):
					return urlid
	
	def msg_indicator(self, message):
		if self.mmapp.has_source(message.message_id):
			return
		title = message.title
		if len(title) > 50:
			title = title[:50] + '...'
		self.mmapp.append_source_with_time(message.message_id, None, title,
			message.timestamp)
		if self.draw_attention:
			allowed = True
			for item in self.black_list:
				if self.check_email(message.account_id, item):
					allowed = False
			if allowed:
				self.mmapp.draw_attention(message.message_id)
	
	def get_folders_list(self, cn):
		result = []
		for folder in self.extra_folders:
			email, folder_name = folder.split(':')
			if self.check_email(self.cn, email):
				folder_name_utf7 = bytes(folder_name, 'utf-7').replace(b'+',
					b'&').replace(b' &', b'- &')
				if b' ' in folder_name_utf7:
					folder_name_utf7 = b'"'+folder_name_utf7+b'"'
				result.append(folder_name_utf7)
		return result
	
	def update(self):
		self.unread_count = [0]*self.l
		self.notifications_queue = []
		for msg in self.unread_messages:
			msg.is_active = False
		for self.cn in range(self.l):
			signal.signal(signal.SIGALRM, self.raise_error)
			folders = self.get_folders_list(self.cn)
			for folder in [None]+folders:
				signal.alarm(7)
				try:
					self.update_single(self.cn, folder_name=folder)
				except TimeoutException:
					logger.warning('Timeout error occured in connection #{0}'.format(self.cn))
					self.notinit[self.cn]=True
				except (imaplib.IMAP4.error, socketerror) as e:
					print_error(e, 'Error in connection #{0}'.format(self.cn))
					self.notinit[self.cn]=True
				except KeyboardInterrupt:
					self.notinit[self.cn]=True
					self.handle_keyboard_interrupt()
				signal.alarm(0) # Cancels alarm
		if with_messagingmenu:
			enum = list(enumerate(sorted(self.unread_messages, key=lambda m: m.timestamp)))
			shown_indices = [ind for ind, msg in enum if msg.is_active][-self.max_count:]
			for ind, msg in enum:
				# Show or hide indicator depending on it's activeness
				if ind in shown_indices:
					self.msg_indicator(msg)
				else:
					self.mmapp.remove_source(msg.message_id)
		try:
			self.show_notifications()
		except GLib.GError as e:
			print_error(e)
		self.first_run = False
		if with_unity:
			number_of_mails = 0
			for i in self.unread_count:
				number_of_mails += int(i)
			self.set_launcher_count(number_of_mails)
		return True
	
	def set_launcher_count(self, count):
		self.launcher.set_property('count', count)
		if all(self.notinit):
			# None of the connections is working
			count_visible = False
		else:
			count_visible = (count > 0) or not self.hide_count
		self.launcher.set_property('count_visible', count_visible)
	
	def show_notifications(self):
		icon = 'unity-mail' if self.display_icons else None
		basemessage = gettext.ngettext('You have %d unread mail', 'You have %d unread mails',
			len(self.notifications_queue))
		basemessage = basemessage.replace('%d', '{0}')
		if self.notifications_queue and self.play_sound:
			try:
				if self.custom_sound:
					subprocess.call(('canberra-gtk-play', '-f', self.custom_sound))
				else:
					subprocess.call(('canberra-gtk-play', '-i', 'message-new-email'))
			except OSError as e:
				print_error(e)
		if len(self.notifications_queue) > (1 if self.first_run else 2):
			senders = set(get_sender_name(notification[0]) for notification in self.notifications_queue)
			unknown_sender = ('' in senders)
			if unknown_sender:
				senders.remove('')
			ts = tuple(senders)
			if len(ts) > 2 or (len(ts) == 2 and unknown_sender):
				message = fix_format(_('from %(t0)s, %(t1)s and others')).format(t0=ts[0], t1=ts[1])
			elif len(ts) == 2 and not unknown_sender:
				message = fix_format(_('from %(t0)s and %(t1)s')).format(t0=ts[0], t1=ts[1])
			elif len(ts) == 1 and not unknown_sender:
				message = _('from %s').replace('%s', '{0}').format(self.notifications_queue[0][0])
			else:
				message = None
			notification = Notify.Notification.new(basemessage.format(len(self.notifications_queue)),
				message, icon)
			notification.set_hint('desktop-entry', GLib.Variant.new_string('unity-mail'))
			notification.show()
		else:
			for notification in self.notifications_queue:
				if self.account_display is None:
					account_string = ''
				else:
					cn = notification[2]
					account_string = '\n' + self.account_display
					account_string = account_string.replace('%HOST%', self.host[cn])
					account_string = account_string.replace('%LOGIN%', self.login[cn])
				if notification[0]:
					message = _('New mail from %s').replace('%s', '{0}').format(notification[0])
				else:
					message = basemessage.format(1)
				notification = Notify.Notification.new(message, notification[1]+account_string, icon)
				notification.set_hint('desktop-entry', GLib.Variant.new_string('unity-mail'))
				notification.show()
