#! /usr/bin/python

"""Manage a mail archive.

The archive is a directory containing a (g)dbm database, some index
files, and some "overflow" files (for messages too big to store
comfortably in the database).  Messages are stored in the archive
exactly as they come in.

The database keys messages on their message-id.  An input message
without a message-id field is refused.  For overflow files, the
database stores the headers with an initial header of the format
"Overflow-file: <filename>"; the overflow file stores the entire
original message.  Overflow filenames are generated from the date+time
they came in.

The index files can be regenerated from the database, but are normally
updated as each new message is archived.  At the moment, it is assumed
that the index files are small enough to be searched sequentially, so
updating simply means adding a line of the form "<message-id>: text"
to the index file.  Initially there are two index files: by subject
and by author.

XXX To do:
- locking
- overflow messages

"""


import rfc822
import gdbm
import string
import os
import sys
import regsub
import StringIO


class Archive:

	indices = ['from', 'subject', 'date']

	def __init__(self, mode = 'c', dir = ".", name = "database"):
		self.mode = mode
		self.dir = dir
		self.name = name
		self.db = None
		self.lock()
		self.opendbm()

	def __del__(self):
		self.close()

	def close(self):
		self.closedbm()
		self.unlock()

	def opendbm(self):
		name = self.makename(self.name)
		self.db = gdbm.open(name, self.mode, 0666)

	def closedbm(self):
		self.db = None

	def lock(self):
		pass

	def unlock(self):
		pass

	def get(self, id):
		return self.db[id]

	def select(self, key, value):
		import regex
		pattern = regex.compile(value, regex.casefold)
		hits = {}
		try:
			fp = open(self.makename("index." + key))
		except IOError:
			return hits
		while 1:
			line = fp.readline()
			if not line: break
			i = string.find(line, ':')
			if i >= 0:
				text = line[i+1:-1]
				if pattern.search(text) >= 0:
					id = line[:i]
					hits[id] = text
		return hits

	def append(self, fp):
		first = 1
		more = 1
		id = None
		while more:
			msg = rfc822.Message(fp, 0)
			if first:
				if not hasattr(msg, 'unixfrom'):
					more = 0
				else:
					more = msg.unixfrom
				first = 0
			try:
				id = msg['message-id']
			except KeyError:
				id = None
			if id and id[0] == '<' and id[-1] == '>':
				id = id[1:-1]
			if not id:
				raise RuntimeError, "No Message-ID header"
			if self.db.has_key(id):
				raise RuntimeError, "Duplicate Message-ID"
			htext = string.joinfields(msg.headers, '')
			if more:
				btext = ''
				while 1:
					line = fp.readline()
					if not line:
						more = 0
						break
					if line[:5] == "From ":
						break
					btext = btext + line
			else:
				btext = fp.read()
			# XXX To do: overflow files
			text = htext + '\n' + btext
			print '-'*20
			print text
			print '-'*20
			self.db[id] = text
			self.indexing(id, msg, btext)
		return id		# Of last appended message

	def regenindices(self):
		print "Regenerating indices ... 000\r",
		sys.stdout.flush()
		for name in self.indices:
			self.remove("index." + name)
		n = 0
		for id in self.db.keys():
			text = self.db[id]
			fp = StringIO.StringIO(text)
			msg = rfc822.Message(fp)
			self.indexing(id, msg, fp.read())
			n = n+1
			if n%25 == 0:
				print "Regenerating indices ... %03d\r" % n,
				sys.stdout.flush()
		print

	def indexing(self, id, msg, btext):
		for key in self.indices:
			try:
				value = msg[key]
			except KeyError:
				value = ''
			self.addline("index." + key, id, value)

	def addline(self, name, id, text):
		text = regsub.gsub('\n', ' ', text)
		name = self.makename(name)
		fp = open(name, "a")
		fp.write("%s: %s\n" % (id, text))
		fp.close()

	def remove(self, name):
		try:
			os.unlink(self.makename(name))
		except os.error:
			pass
 
	def makename(self, name):
		return os.path.join(self.dir, name)


def test(archiver=Archive):
	"""Simple test program -- command line interface.

	To chdir to a directory before doing anything else:
		mailarch -d directory [other options and arguments...]

	To write messages to a logfile:
		mailarch -l logfile [other options and arguments...]

	To archive messages:
		mailarch (one message from stdin)
		mailarch file ... (one message per file)
		mailarch +folder ... (archives all messages in folder)

	To select messages by from and/or subject header (prints message ids):
		mailarch [-s subject] [-f from]
	
	To select messages by body text (requires textindex interface):
		textindex [-t 'word ...']

	To retrieve messages by message-id (prints whole messages):
		mailarch -m id ...

	To regenerate the index files:
		mailarch -r

	"""

	import getopt
	import os
	try:
		opts, args = getopt.getopt(sys.argv[1:], 'd:l:rs:f:mt:')
	except getopt.error, msg:
		sys.stdout = sys.stderr
		print msg
		print test.__doc__
		sys.exit(2)
	regen = 0
	subject = None
	fromm = None
	text = None
	msgid = 0
	for o, a in opts:
		if o == '-d':
			os.chdir(a)
		if o == '-l':
			sys.stdout = sys.stderr = open(a, 'a')
		if o == '-m': msgid = 1
		if o == '-r': regen = 1
		if o == '-s': subject = a
		if o == '-f': fromm = a
		if o == '-t': text = a
	if msgid:
		a = archiver('r')
		for id in args:
			text = a.get(id)
			sys.stdout.write(text)
		return
	if text:
		a = archiver('r')
		hits = a.select('text', text)
		if not hits:
			sys.stderr.write("No hits\n")
		else:
			for id in hits.keys():
				print id
		return
	if subject or fromm:
		a = archiver('r')
		if subject:
			hit_s = a.select('subject', subject)
		else:
			hit_s = None
		if fromm:
			hit_f = a.select('from', fromm)
		else:
			hit_f = None
		if subject and fromm:
			hits = intersection(hit_s, hit_f)
		elif subject:
			hits = hit_s
		else:
			hits = hit_f
		if not hits:
			sys.stderr.write("No hits\n")
		else:
			for id in hits.keys():
				print id
		return
	a = archiver()
	if regen:
		a.regenindices()
	elif args:
		for name in args:
			if name[0] == "+":
				print "Appending MH folder %s ..." % name
				import mhlib
				mh = mhlib.MH()
				folder = mh.openfolder(name[1:])
				messages = folder.listmessages()
				for n in messages:
					file = folder.getmessagefilename(n)
					fp = open(file)
					print "\tAppend message %d ..." % n
					try:
						a.append(fp)
					except RuntimeError, msg:
						print msg
			else:
				# XXX Should split it if it is a Unix mailbox
				print "Append %s ..." % name
				fp = open(name)
				a.append(fp)
	else:
		a.append(sys.stdin)
	a.close()


def intersection(a, b):
	d = {}
	for key in a.keys():
		if b.has_key(key):
			d[key] = key
	return d

if __name__ == '__main__':
	test()
