#!/usr/local/bin/python

from socket import *
from string import strip, split
from sys import argv, exit
import re, signal, md5

Timeout = 'Timeout'
def timeout(sig, frame) :
    raise Timeout

OK_LIST = re.compile(r"\+OK ([0-9]+) ([0-9]+).*")

POP_ERROR = 'POP_ERROR'
POP_TIMEOUT = 'POP_TIMEOUT'
POP_USER = 'POP_USER'


class POP :
    "The POP class is a clientside version of the POP3 protocol"
    def __init__(self, HOST = '', PORT = 110, DEBUG = 0) :
	"Initialize instance of Pop using (optional) HOST and PORT"
	# Setup a handler for timeouts
	signal.signal(signal.SIGALRM, timeout)
	# Initate variables for socket use
	self.socket = 0			# The socket, not yet
	self.username = ''
	self.password = ''
	self.DEBUG = DEBUG
	if HOST == '' :
	    return
	self.connect(HOST, PORT)
	
    def __del__(self) :
	"Make sure that we close down before entering oblivion"
	if self.DEBUG :
	    print "Entering __del__"
	if self.socket :
	    if self.username :
		self.socket.send("RSET\r\n")
		data = self.socket.recv(1024)
		self.socket.send("QUIT\r\n")
		data = self.socket.recv(1024)
		self.username = ''
	    self.socket.close()

    def __getitem__(self, id) :
	"Same as self.retrieve(ID)"
	return self.retrieve(id)

    def __delitem__(self, id) :
	"Same as self.delete(ID)"
	self.delete(id)

    def __len__(self) :
	"Return the number of messages at POP-server"
	return self.status()[0]

    def triad(self) :
	"Return hostname, user and password on a connection"
	if not self.username :
	    raise POP_ERROR, "User is not logged in"
	return (self.host, self.username, self.password)

    def connect(self, HOST, PORT = 110) :
	"Connect to HOST using (optional) PORT"
	if self.DEBUG :
	    print "Entering connect to host %s on port %d" % ( HOST, PORT)
	if self.socket :
	    raise POP_ERROR, "Already connected"
	signal.alarm(30)
	try :
	    self.socket = socket(AF_INET, SOCK_STREAM)
	    self.socket.connect(HOST, PORT)
	    self.host = HOST
	    self.port = PORT
	except Timeout :
	    raise POP_TIMEOUT, 'Timed out when trying to connect to %s' % HOST
	signal.alarm(120)
	try :
	    data = self.socket.recv(1024)
	except Timeout :
	    raise POP_TIMEOUT, 'Timed out when getting initial message'
	signal.alarm(0)
	if data[0:3] != '+OK' :
	    self.socket.shutdown(2)
	    raise POP_ERROR, \
		  'Failed to get an OK initial message from %s' % HOST
	answer = strip(data[4:])
	if self.DEBUG :
	    print "Initial message was: '%s'" % answer
	m = re.match(r".*(<[^>]*>).*\r\n", answer)
	if m :
	    self.timestamp = m.group(1)
	    if self.DEBUG :
		print "Timestamp was: %s" % self.timestamp
	else :
	    self.timestamp = ''
	    if self.DEBUG :
		print "No timestamp was given"

    def close(self, rset = 1) :
	"Close a connection, and maybe (optional) RESET first"
	if self.DEBUG :
	    if rset :
		print "Entering close, with reset"
	    else :
		print "Entering close, without reset"
	if self.username :
	    try :
		if rset :
		    data = self.command("RSET")
		data = self.command("QUIT", 300)
	    except POP_TIMEOUT :
		raise POP_TIMEOUT, "QUIT timed out, indeterminded behavior"
	    self.username = ''
	self.socket.close()
	self.socket = 0

    def reconnect(self) :
	"Reset the current connection by closing and reopening it"
	if self.socket :
	    self.close(1)
	self.connect(self.host, self.port)

    def multiline(self, ALARM = 60) :
	"""Get a multiline answer. Ending with a line only a dot '.' which
	   will not be included. Allowing lines to be separated by only
	   <LF> instead of the required <CRLF>. Returning a list of lines"""
	if self.DEBUG :
	    print "Entering multiline"
	try :
	    result = []
	    data = filter(lambda x : x != '\r', self.socket.recv(1024))
	    while data[-3:] != '\n.\n' :
		data = data + \
		       filter(lambda x : x != '\r', self.socket.recv(1024))
	    lines = split(data, '\n')
	    for line in lines :
		if line and line[0] == '.' :
		    if len(line) == 1 :
			signal.alarm(0)
			return result
		    else :
			result.append(line[1:])
		else :
		    result.append(line)
	    raise POP_ERROR, "Unexpected end of a multiline answer"
	except Timeout :
	    raise POP_TIMEOUT, "multiline timed out when getting data"

    def command(self, COMMAND, ALARM = 60) :
	"""Send a COMMAND and return the answer"""
	if self.DEBUG :
	    print "Entering command with command %s" % COMMAND
	if COMMAND[-1] != '\n' :
	    COMMAND = COMMAND + '\r\n'
	try :
	    signal.alarm(ALARM)
	    if not self.socket :
		if self.host and self.port :
		    self.connect(self.host, self.port)
 		    signal.alarm(ALARM)
		else :
		    signal.alarm(0)
		    raise POP_ERROR, \
			  "Unable to invoke command '%s', no connection" \
			  % COMMAND
	    if self.DEBUG :
		print "Sending %s" % COMMAND
	    self.socket.send(COMMAND)
	    signal.alarm(ALARM)
	    if self.DEBUG :
		print "Reciving data. . .", 
	    data = self.socket.recv(1024)
	    if self.DEBUG :
		m = re.match(r"(.*)\r\n", data)
		if m :
		    print "'%s'" % m.group(1)
		else :
		    print data
	    signal.alarm(0)
	    return data
	except Timeout :
	    raise POP_TIMEOUT, "Timed out on command '%s'" % COMMAND
	
    def user(self, NAME, PASSWORD) :
	"Try log in USER with PASSWORD, first trying with apop command"
	if self.DEBUG :
	    print "Entering user with %s and %s" % (NAME, PASSWORD)
	if self.username :
	    raise POP_ERROR, "User %s is already in" % self.username
	# Try APOP first to avoid sendig uncrypted password, if we can
	if self.timestamp :
	    m = md5.new(self.timestamp + PASSWORD)
	    l = map(lambda x : hex(ord(x))[2:], m.digest())
	    digest = reduce(lambda x, y : x+y,
			    map(lambda x : hex(ord(x))[2:], m.digest()),
			    "")
	    data = self.command("APOP %s %s" % (NAME, digest))
	    if data[0:3] == '+OK' :
		self.username = NAME
		self.password = PASSWORD
		return
	    # This is to handle a protocol-error in (at least) qpop
	    # which disconnects after a failed APOP instead of staying
	    # in authorization state
	    self.reconnect()
	    # End of protocol-error handling
	data = self.command("USER %s" % NAME)
	if data[0:3] != '+OK' :
	    raise POP_USER, "Username %s was not accepted, %s" % (NAME, data)
	data = self.command("PASS %s" % PASSWORD, 600) # An extended timeout
	if data[0:3] != '+OK' :
	    raise POP_USER, "Password was not accepted, %s" % data
	self.username = NAME		# We are in
	self.password = PASSWORD
	
    def quit(self) :
	"""Close shop in a nice way"""
	if self.DEBUG :
	    print "Entering quit"
	self.close(0)

    def list(self, msg = 0) :
	"""List all messages, or just one MESSAGE. Returns an id, size pair
	   or a list of them"""
	if msg == 0 :
	    return self.list_all()
	if not self.username :
	    raise POP_ERROR, 'May only invoke list when user is active'
	data = self.command("LIST %d" % msg)
	m = OK_LIST.match(data)
	if m :
	    return (int(m.group(1)), int(m.group(2)))
	raise POP_ERROR, "LIST %d failed, %s" % (msg, data)
	
    def list_all(self) :
	"Auxiliary function for list"
	if not self.username :
	    raise POP_ERROR, 'May only invoke list when user is active'
	data = self.command("LIST")
	if data[0:3] == '+OK' :
	    lines = self.multiline()	# Get a multiline answer
	    result = []
	    for line in lines :
		m = re.match(r"([0-9]+) ([0-9]+).*", line)
		if m :
		    result.append((int(m.group(1)), int(m.group(2))))
		elif self.DEBUG :
		    print "List: ignoring '%s'" % line
	    return result
	raise POP_ERROR, "LIST failed, %s" % data

    def uidl(self, msg = 0) :
	"""List unique-id for all messages, or just one MESSAGE. Returns an
	   id, uid par or a list of them."""
	if msg == 0 :
	    return self.uidl_all()
	if not self.username :
	    raise POP_ERROR, 'May only invoke uidl when user is active'
	data = self.command("UIDL %d" % msg)
	m = re.match(r"\+OK ([0-9]+) ([!-~]+).*", data)
	if m :
	    return (int(m.group(1)), m.group(2))
	raise POP_ERROR, "UIDL %d failed, %s" % (msg, data)
	
    def uidl_all(self) :
	"Auxiliary function for uidl"
	if not self.username :
	    raise POP_ERROR, 'May only invoke uidl when user is active'
	data = self.command("UIDL")
	if data[0:3] == '+OK' :
	    lines = self.multiline()	# Get a multiline answer
	    result = []
	    for line in lines :
		m = re.match(r"([0-9]+) ([!-~]+).*", line)
		if m :
		    result.append((int(m.group(1)), m.group(2)))
		elif self.DEBUG :
		    print "Uidl: ignoring '%s'" % line
	    return result
	raise POP_ERROR, "UIDL failed, %s" % data

    def noop(self) :
	"Pass the time (maybe to keep a connection alive)"
	if not self.username :
	    raise POP_ERROR, 'May only invoke noop when user is active'
	data = self.command("NOOP")
	if data[0:3] != '+OK' :
	    raise POP_ERROR, "NOOP failed, %s" % data
	
    def retrieve(self, msg, delete = 0) :
	"Get a MESSAGE, and optionally DELETE it afterwards."
	if not self.username :
	    raise POP_ERROR, 'May only invoke retrieve when user is active'
	data = self.command("RETR %d" % msg)
	if data[0:3] != '+OK' :
	    raise POP_ERROR, "RETR %d failed, %s" % (msg, data)
	result = self.multiline()
	if delete :
	    self.delete(msg)
	return result

    def reset(self) :
	"Reset the state of this connection, does not log out the user"
	if not self.username :
	    raise POP_ERROR, 'May only invoke reset when user is active'
	data = self.command("RSET")
	if data[0:3] != '+OK' :
	    raise POP_ERROR, "RSET failed, %s" % data

    def delete(self, msg) :
	"Delete a MESSAGE"
	if not self.username :
	    raise POP_ERROR, 'May only invoke delete when user is active'
	data = self.command("DELE %d" % msg)
	if data[0:3] != '+OK' :
	    raise POP_ERROR, "DELE failed, %s" % data

    def top(self, msg, lines = 0) :
	"From MESSAGE get headers and optionally some LINES of the letter" 
	if not self.username :
	    raise POP_ERROR, 'May only invoke top when user is active'
	data = self.command("TOP %d %d" % (msg, lines))
	if data[0:3] != '+OK' :
	    raise POP_ERROR, "TOP %d %d failed, %s" % (msg, lines, data)
	return self.multiline()

    def status(self) :
	"Returns the number of messages, and the total size of them"
	if not self.username :
	    raise POP_ERROR, 'May only invoke status when user is active'
	data = self.command("STAT")
	m = re.match(r"\+OK ([0-9]+) ([0-9]+).*", data)
	if m :
	    return (int(m.group(1)), int(m.group(2)))
	else :
	    raise POP_ERROR, "STAT failed, %s" % data

##
# Functions that uses POP
##

def kill_large(pop, the_size = 209715152, verbose = 0) :
    """Connect to HOST use USER and PASSWORD to enter, and remove all
       messages larger than (optional) SIZE, and be TALKATIVE.
       Returns the number of remaing messages, and their size."""
    (msgs, total) = pop.status()
    if verbose :
	print "Total of %d kbytes in %d messages" % \
	      ((total+512) / 1024 , msgs)
    for (n, size) in pop.list() :
	if size > the_size :
	    if verbose :
		print "Deleting %d because of size (%d)" % (n, size)
	    pop.delete(n)
    (msgs, total) = pop.status()
    if verbose :
	print "Total size of %d remaining messages %d kbytes" % \
	      (msgs, (total + 512) / 1024)
    return (msgs, (total + 512) / 1024)
	
def fetch_mail(host, user, password, mailfile) :
    "Poll mail on HOST for USER with PASSWORD and write it to MAILFILE"
    pop = POP(host)
    pop.user(user, password)

    triad = pop.triad()
    (N, size) = pop.status()
    if N == 0 :
	pop.quit()
	return triad
    # Start by copying the old file
    from tempfile import mktemp
    from time import asctime, localtime, time
    import re, os
    eol = '\n'
    filename = mktemp()
    outfile = open(filename, "w")
    try :
	infile = open(mailfile, "r")
    except IOError :
	infile = None
    if infile :
	line = infile.readline()
	while line :
	    outfile.write(line)
	    line = infile.readline()
	infile.close()

    returnpath = re.compile(r"^[Rr][Ee][Tt][Uu][Rr][Nn]-[Pp][Aa][Tt][Hh]:[\t ]*(.*)")
    for id in range(1, N+1) :
       	message = pop.retrieve(id)
	sender = "<>"
	for line in message :
	    m = returnpath.match(line)
	    if m :
		sender = m.group(1)
		break
	    elif line == '' :
		break
	outfile.write("From %s %s\n" % (sender, asctime(localtime(time()))))
	for line in message :
	    outfile.write("%s%s" % (line, eol))
	if id < N :
	    outfile.write(eol)
    outfile.close()
    try :
	os.rename(filename, mailfile)
    except os.error :			# Rename failed, try to copy
	outfile = open(mailfile, "w")
	infile = open(filename, "r")
	line = infile.readline()
	while line :
	    outfile.write(line)
	    line = infile.readline()
	infile.close()
	outfile.close()
	os.remove(filename)
    # Mailbox is updated, delete mail
    for id in range(1, N+1) :
	pop.delete(id)
    pop.quit()
    return triad

def poll_mail(host, user, password, mailfile, delay = 120) :
    "Go and fetch the mail, again and again"
    from time import sleep
    while 1 :
	(host, user, password) = fetch_mail(host, user, password, mailfile)
	sleep(delay)
	
if __name__ == '__main__' :
    import re, sys, os

    host = user = password = ''
    for i in range(1, len(sys.argv)) :
	m = re.match(r"([a-zA-Z0-9]+)@([-.a-zA-Z0-9]+)", sys.argv[i])
	if m :
	    user = m.group(1)
	    host = m.group(2)
	else :
	    password = sys.argv[i]
    mailfile = os.environ['MAIL']
    if mailfile and host and user and password :
	# fetch_mail(host, user, password, '/tmp/mail')
	poll_mail(host, user, password, mailfile)
    else :
	print "usage: %s <user>@<maildrop> <password or secret phrase>" % \
	      argv[0]