#!/usr/bin/env python

"""
poptoimap
(c) 2007 Tony Houghton
Home page: <http://www.realh.co.uk>
Distributed under the GPL: <http://www.gnu.org/licenses/gpl.html>
"""

version = (0, 2, 0)

"""
Changelog:

0.2.0 (2008-04-01)
	* replace() function added

0.1.3 (2007-05-27)
	* Control over how errors from to/cc pipes are handled

0.1.2 (2007-05-02)
	* Prints '[killed]' for killed messages

0.1.1 (2007-03-12)
	* Logs date/time in verbose mode
"""

import email.Utils
import exceptions
import getpass
import imaplib
import optparse
import os
import poplib
import re
import socket
import subprocess
import sys

#####################################################################

" Some private stuff. "

pop_descriptions = []

DEFAULT_FOLDER = 'INBOX'

create_imap = None

class ConfigError(exceptions.Exception):
	pass

#####################################################################

"All available to .poptoimaprc at top-level but only needed for advanced use."

verbose = False

delivery = None

def _(s):
	return s

class LoginFailed(exceptions.Exception):
	pass

class FilterFailed(exceptions.Exception):
	pass

class DeliveryFailed(exceptions.Exception):
	pass

class DeliverySucceeded(exceptions.Exception):
	""" This isn't an error, but is raised by to() to prevent further
	processing after a successful delivery. Arg should indicate where it
	was delivered to. """
	pass

class KilledOK(exceptions.Exception):
	""" This is also a "success" exception. """
	pass

#####################################################################

def imap_parse(full_desc):
	""" Parses an imap URI as passed on the command line, returning an array:
	[proto, host, port, username, password, folder] where proto is "imap" or
	"imaps". Prompts for password if not given in URI. """
	global DEFAULT_FOLDER
	proto, desc = full_desc.split('://', 1)
	if proto == 'imap':
		port = 143
	elif proto == 'imaps':
		port = 993
	else:
		raise ConfigError(_("Unrecognised IMAP protocol in '%s'") % full_desc)
	desc = desc.split('/', 1)
	if len(desc) > 1 and desc[1]:
		folder = desc[1]
	else:
		folder = None
	user = None
	passwd = None
	desc = desc[0].split('@')
	if len(desc) > 1:
		user = desc[0].split(':')
		if len(user) > 1:
			passwd = user[1]
		user = user[0]
		desc = desc[1]
	else:
		desc = desc[0]
	host = desc.split(':')
	if len(host) > 1:
		port = host[1]
	host = host[0]
	if not passwd:
		if user:
			prompt = _("Enter password for '%s': ") % full_desc
		else:
			user = getpass.getuser()
			prompt = _("Enter password for user '%s' on '%s': ") \
					% (user, full_desc)
		passwd = getpass.getpass(prompt)
	if not user:
		user = getpass.getuser()
	if not folder:
		folder = DEFAULT_FOLDER
	return [proto, host, port, user, passwd, folder]

def imap_connect(full_desc):
	global create_imap
	id = imap_parse(full_desc)
	if id[0] == 'imap':
		imo = imaplib.IMAP4(id[1], id[2])
	elif id[0] == 'imaps':
		imo = imaplib.IMAP4(id[1], id[2])
	response = imo.login(id[3], id[4])
	if response[0] != 'OK':
		raise LoginFailed(_("IMAP login to '%s' failed: %s %s") \
				% (full_desc, response[0], response[1][0]))
	if create_imap:
		response = imo.list()
		if response[0] == 'OK':
			imo.folder_names = []
			for f in response[1]:
				imo.folder_names.append( \
						f.split(') ', 1)[1].split(' ', 1)[1].strip('"'))
		else:
			raise LoginError(_("Unable to list IMAP folders on %s") % full_desc)
	return imo

def imap_append(imap, folder, message):
	""" Convenience function to call an IMAP4 object's append() method and
	raise an exception if the result isn't OK. message may be a single string
	or an array of lines. folder may be None (defaults to DEFAULT_FOLDER). """
	global create_imap, DEFAULT_FOLDER
	if not folder:
		folder = DEFAULT_FOLDER
	if create_imap and not folder in imap.folder_names:
		result = imap.create(folder)
		if result[0] != 'OK':
			raise DeliveryFailed(_("Error creating IMAP folder '%s': %s %s") \
					% (folder, result[0], str(result[1])))
	if not isinstance(message, str):
		message = ''.join(message)
	result = imap.append(folder, None, None, message)
	if result[0] != 'OK':
		raise DeliveryFailed(_("Error appending to IMAP folder '%s': %s %s") \
				% (folder, result[0], result[1][0]))


class DeliveryBase:
	""" The main MDA class. Holds a message and provides methods for testing
	and delivering it. A message is stored as a list of lines with trailing
	newlines. """
	def __init__(self, imap):
		""" imap is the IMAP connection to which messages are delivered """
		self.imap = imap
		self.message = None
		self.boundary = None	# index into message of header-body separator
		self.folder = DEFAULT_FOLDER
		self.size = 0
	
	def set_folder(self, folder):
		self.folder = folder
	
	def set_imap(self, imap):
		self.imap = imap

	def set_message(self, message):
		self.message = message
		self.size = 0
		for l in self.message:
			self.size += len(l)
		self.boundary = None
	
	def message_size(self):
		return self.size
	
	def get_boundary(self):
		" May return -1 if message has no separator. "
		if not self.boundary:
			for i in range(len(self.message)):
				if self.message[i] == "\n":
					self.boundary = i
					break
			else:
				self.boundary = -1
		return self.boundary
	
	def match(self, expression, case_sensitive = False, match = None):
		""" expression is a string or a compiled regex. In the latter case
		case_sensitive flag is ignored.
		match = 'header' (default), 'body', or 'all'. """
		#print "Testing against re", expression
		if isinstance(expression, str):
			if case_sensitive:
				flags = None
			else:
				flags = re.I
			expression = re.compile(expression, flags)
		if match == None or match == 'header':
			rng = [0, self.get_boundary()]
			if rng[1] == -1:
				rng[1] = len(self.message)
		elif match == 'body':
			rng = [self.get_boundary() + 1, len(self.message)]
			if not rng[0]:
				return None
		elif match == 'all':
			rng = [0, len(self.message)]
		for i in range(rng[0], rng[1]):
			m = expression.search(self.message[i].rstrip("\n"))
			if m:
				return m
		return None

	def xfilter(self, args, shell = True):
		""" args and shell as in subprocess.Popen(). Raises exception if
		exit code NZ. """
		subproc = subprocess.Popen(args, shell = shell,
				stdin = subprocess.PIPE, stdout = subprocess.PIPE)
		stdout, stderr = subproc.communicate("".join(self.message))
		if subproc.returncode:
			if not isinstance(args, str):
				args = ' '.join(args)
			raise FilterFailed(_("Filtering through '%s'failed because " \
					"it returned an error code (%d)") % (args,
						subproc.returncode))
		# This is very inefficient, but the communicate method greatly
		# simplifies things and only returns a flattened string
		message = stdout.split("\n")
		for n in range(len(message)):
			message[n] = message[n].rstrip("\r") + "\n"
		self.set_message(message)

	def pipe(self, args, shell = True):
		""" args and shell as in subprocess.Popen().
		Returns program's exit code. """
		subproc = subprocess.Popen(args, shell = shell, stdin = subprocess.PIPE)
		if not self.message[0].startswith('From '):
			subproc.stdin.write("From %s@%s %s\n" % (getpass.getuser(),
					socket.getfqdn(), email.Utils.formatdate()))
		subproc.stdin.writelines(self.message)
		subproc.stdin.close()
		return subproc.wait()

	def forward(self, address):
		""" Forward a mail to another address or addresses. address arg may be
		a list/tuple of addresses or a single string delimited by whitespace.
		"""
		if isinstance(address, str):
			address = [address]
		addresses = []
		for a in address:
			addresses += a.split()
		p = self.pipe(["sendmail", "-bm"] + addresses, False)
		if p:
			raise DeliveryFailed(_("Forwarding failed because " \
				"sendmail returned an error code (%d)") % p)

	def cc(self, folder = None, error = 0):
		""" Response to NZ exit codes from pipes:
			0	exception
			1	warn
			2	ignore
		"""
		if folder:
			if folder[0] == '|':
				p = self.pipe(folder[1:])
				if p:
					if error == 1:
						print >> sys.stderr, \
								_("Warning: %s returned %d") % (folder, p)
					elif error == 0:
						raise DeliveryFailed(_("'%s' failed, returning " \
								"error code %d") % (folder, p))
				return
			elif folder[0] == '!':
				self.forward(folder[1:])
				return
		else:
			folder = self.folder
		imap_append(self.imap, folder, self.message)

	def to(self, folder = None, error = 0):
		global DEFAULT_FOLDER
		self.cc(folder, error)
		if not folder:
			folder = self.folder
			if not folder:
				folder = DEFAULT_FOLDER
		raise DeliverySucceeded(folder)

	def kill(self):
		raise DeliverySucceeded('[killed]')

#######################################################################

" As seen by .poptoimaprc at top-level "

def match(expression, case_sensitive = False, match = None):
	return delivery.match(expression, case_sensitive, match)

def xfilter(args, shell = True):
	delivery.xfilter(args, shell)

def size():
	return delivery.message_size()

def cc(folder = None, error = 0):
	delivery.cc(folder, error)

def to(folder = None, error = 0):
	delivery.to(folder, error)

def kill():
	delivery.kill()

def pop(host = None, port = None, user = None, passwd = None,
		ssl = False, folder = None):
	""" folder may be a full imap URI. """
	if not host:
		host = 'localhost'
	if not port:
		if ssl:
			port = 995
		else:
			port = 110
	if not user:
		user = getpass.getuser()
	if not passwd:
		passwd = getpass.getpass(_("Enter password for POP user '%s' on '%s'") \
				% (user, host))
	global pop_descriptions
	pop_descriptions.append([host, port, user, passwd, ssl, folder])

def replace_message(message):
	delivery.set_message(message)


########################################################################

IMAP_SERVER = None

def loadrc(filename = None):
	global DEFAULT_FOLDER, IMAP_SERVER, delivery, deliver_function
	global pop_descriptions, verbose
	if not filename:
		def get_from_base(base):
			f = os.path.expanduser(os.path.join(base, "realh.co.uk",
				"poptoimaprc"))
			if os.path.exists(f):
				return f
			else:
				return None
		rcfile = get_from_base(os.environ.get("XDG_CONFIG_HOME",
				os.path.join("~", ".config")))
		if not rcfile:
			rcfile = os.path.expanduser('~/.poptoimaprc')
		if not os.path.exists(rcfile):
			xdg = os.environ.get("XDG_CONFIG_DIRS", "/etc/xdg")
			for x in xdg.split(':'):
				rcfile = get_from_base(x)
				if rcfile:
					break
	else:
		rcfile = os.path.expanduser(filename)
	if not rcfile or not os.path.exists(rcfile):
		if filename:
			raise ConfigError(_("%s missing") % filename)
		else:
			raise ConfigError(_("No config file"))
	globals = { \
			'_' : _,
			'LoginFailed' : LoginFailed,
			'FilterFailed' : FilterFailed,
			'DeliveryFailed' : DeliveryFailed,
			'DeliverySucceeded' : DeliverySucceeded,
			'KilledOK' : KilledOK,
			'imap_parse' : imap_parse,
			'imap_connect' : imap_connect,
			'imap_append' : imap_append,
			'DeliveryBase' : DeliveryBase,
			'match' : match,
			'xfilter' : xfilter,
			'size' : size,
			'cc' : cc,
			'to' : to,
			'kill' : kill,
			'version' : version,
			'pop' : pop,
			'replace_message' : replace_message }
	locals = { \
			'delivery' : delivery,
			'verbose' : verbose }
	execfile(rcfile, globals, locals)
	DEFAULT_FOLDER = locals.get('DEFAULT_FOLDER')
	if not DEFAULT_FOLDER:
		DEFAULT_FOLDER = 'INBOX'
	IMAP_SERVER = locals.get('IMAP_SERVER')
	delivery = locals.get('delivery')
	deliver_function = locals['deliver']
	verbose = locals['verbose']
	if not pop_descriptions:
		raise ConfigError(_("No POP server(s) specified"))
	if IMAP_SERVER:
		imap_server = imap_connect(IMAP_SERVER)
	else:
		imap_server = None
	if not delivery:
		if not imap_server:
			raise ConfigError(_("No IMAP server specified"))
		delivery = DeliveryBase(imap_server)

def humanise_size(size):
	exp = 0
	while size > 2048:
		size = float(size) / 1024
		exp += 1
	# Truncate decimal
	size = "%f" % size
	if len(size) > 4:
		size = size[:4]
	size = size.rstrip('0').rstrip('.')
	return size + " KMGTP"[exp]

def run_deliver(index, size, message):
	global verbose, delivery, deliver_function
	if verbose:
		from_ = ""
		to_ = ""
		subj_ = ""
		for l in message:
			if l.startswith('From: '):
				from_ = l
			elif l.startswith('To: '):
				to_ = l
			elif l.startswith('Subject: '):
				subj_ = l
		sys.stdout.write(from_ + to_ + subj_)
	print "%4d %-5s" % (index, humanise_size(size)),
	delivery.set_message(message)
	try:
		if deliver_function:
			deliver_function(message)
		delivery.to()
	except DeliverySucceeded, e:
		e = str(e)
		if (e.startswith("INBOX.") or e.startswith("INBOX/")) \
				and len(e) > 6:
			e = e[6:]
		print ">>", e
	except KilledOK:
		print "Killed"

def run_pop(desc):
	# desc = [host, port, user, passwd, ssl, folder]
	global delivery
	if desc[4] == True:
		pop = poplib.POP3_SSL(desc[0], desc[1])
	else:
		pop = poplib.POP3(desc[0], desc[1])
	delivery.set_folder(desc[5])
	try:
		if verbose:
			print pop.getwelcome().rstrip()
		response = pop.user(desc[2])
		if verbose:
			print response.rstrip()
		response = pop.pass_(desc[3])
		if verbose:
			print response.rstrip()
		response = pop.list()
		size = 0
		for m in response[1]:
			size += int(m.split()[1])
		print _("%s on %s: %d message(s), %sbytes") % (desc[2], desc[0],
				len(response[1]), humanise_size(size))
		for m in response[1]:
			index = int(m.split()[0])
			r = pop.retr(index)
			size = r[2]
			msg = r[1]
			for i in range(len(msg)):
				msg[i] = msg[i].rstrip("\r\n") + "\n"
			run_deliver(index, size, msg)
			pop.dele(index)
	except:
		pop.quit()
		raise
	pop.quit()

def main():
	global version, verbose, pop_descriptions, create_imap
	parser = optparse.OptionParser(version = "%d.%d.%d" % version,
			description = _("Fetches mail from a POP server and delivers it " \
					"to an IMAP server."), prog = 'poptoimap')
	parser.add_option('-v', '--verbose', action = 'store_true',
			help = _("Print some extra information."))
	parser.add_option('-c', '--rcfile', action = 'store',
			default = None,
			help = _("Use an alternative config file instead of the default " \
					"(see manual)."))
	parser.add_option('-r', '--create', action = 'store_true',
			help = _("Create IMAP folders if they don't exist."))
	options = parser.parse_args(sys.argv)[0]
	if options.verbose:
		verbose = True
		import time
		print "**********************"
		print time.strftime("%c")
	if options.create:
		create_imap = True
	loadrc(options.rcfile)
	for p in pop_descriptions:
		run_pop(p)

if __name__ == '__main__':
	main()