path: root/modules/tracking.py

info = {
	'author': 'Jon Bergli Heier',
	'title': 'Posten.no tracking',
	'description': 'Fetches tracking info from posten.no.',
}

cfg_section = 'module/tracking'

import urllib2, datetime, re
from lxml import etree
from twisted.internet import reactor
from twisted.internet.task import LoopingCall
from sqlalchemy import create_engine, Column, Integer, String, DateTime, ForeignKey
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relation, backref
from sqlalchemy.orm.exc import NoResultFound
from sqlalchemy.exc import IntegrityError
from sqlalchemy.sql import or_, and_

engine = Session = None
Base = declarative_base()

class Consignment(Base):
	__tablename__ = 'consignment'

	id = Column(Integer, primary_key = True)
	nick = Column(String, nullable = False)
	channel = Column(String) # NULL for private
	code = Column(String, nullable = False, unique = True)
	label = Column(String)
	added = Column(DateTime, nullable = False)
	last = Column(DateTime)
	type = Column(String, nullable = False) # Tracking module

	packages = relation('Package', backref = 'consignment', primaryjoin = 'Package.consignment_id == Consignment.id')

	def __init__(self, nick, channel, type, code, label, added):
		self.nick = nick
		self.channel = channel
		self.type = type
		self.code = code
		self.label = label
		self.added = added

	def __str__(self):
		s = '\002%s\002' % self.code.encode('utf-8')
		if self.label:
			s += ' (%s)' % self.label.encode('utf-8')
		return s

	def __unicode__(self):
		return str(self).decode('utf-8')

class Package(Base):
	__tablename__ = 'package'

	id = Column(Integer, primary_key = True)
	consignment_id = Column(Integer, ForeignKey('consignment.id'))
	code = Column(String, nullable = False, unique = True)
	last = Column(DateTime) # Last update
	status = Column(String) # Last status

	def __init__(self, consignment_id, code):
		self.consignment_id = consignment_id
		self.code = code

	def __str__(self):
		s = '\002%s\002' % self.code.encode('utf-8')
		if self.consignment.label:
			s += ' (%s)' % self.consignment.label.encode('utf-8')
		return s

	def __unicode__(self):
		return str(self).decode('utf-8')

class PackageError(Exception): pass
class NoPackageFound(PackageError): pass

class TrackingResult:
	def __init__(self, code, date, desc, delivered = False):
		self.code, self.date, self.desc, self.delivered = code, date, desc, delivered

	def __str__(self):
		return '\002{code}\002 {date} - {desc}'.format(code = self.code, date = self.date, desc = self.desc)

tracking_modules = {}

class TrackingModuleMeta(type):
	def __init__(cls, name, bases, dict):
		if name != 'TrackingModule':
			tracking_modules[cls.name] = cls

class TrackingModule(object):
	__metaclass__ = TrackingModuleMeta

class PostenModule(TrackingModule):
	name = 'posten'
	url = 'http://sporing.posten.no/sporing.xml?q=%s'

	def get_xml(self, url):
		try:
			u = urllib2.urlopen(url, timeout = config.getfloat('module/tracking', 'timeout'))
		except Exception as e:
			raise PackageError(str(e))
		if u.headers['content-type'].startswith('application/xml'):
			xml = etree.parse(u)
		else:
			xml = None
		u.close()
		return xml

	def get_url(self, code = None):
		url = 'http://sporing.posten.no/'
		if code:
			url += 'sporing.html?q=' + code
		return url

	def track(self, code):
		xml = self.get_xml(self.url % code)
		if not xml:
			return

		ns = 'http://www.bring.no/sporing/1.0'
		packages = xml.find('//{%s}PackageSet' % (ns,))
		if packages is None:
			raise NoPackageFound('No packages found for \002%s\002.' % code)

		results = []
		for package in packages:
			code = package.attrib['packageId']
			eventset = package.find('{%s}EventSet' % ns)
			if not len(eventset):
				continue
			last = eventset[0]
			desc = last.find('{%s}Description' % ns).text.replace('&lt;', '<').replace('<br/>', ' ')
			desc = re.sub(r'<[^>]*?>', '', desc).encode('utf8')
			isodate = last.find('{%s}OccuredAtIsoDateTime' % ns).text[:-10]
			isodate = datetime.datetime.strptime(isodate, '%Y-%m-%dT%H:%M:%S')
			city = last.find('{%s}City' % ns).text
			status = last.find('{%s}Status' % ns).text.encode('utf-8')
			if city:
				desc = '%s (%s)' % (desc, city.encode('utf8'))
			date = last.find('{%s}OccuredAtDisplayDate' % ns).text + ' ' + last.find('{%s}OccuredAtDisplayTime' % ns).text
			results.append(TrackingResult(code, isodate, desc, status == 'DELIVERED'))
		return results

def code_split(code):
	if ':' in code:
		return code.split(':')
	else:
		return 'posten', code

def get_tracking_module(arg):
	type, code = code_split(arg)

	return tracking_modules[type]()

class Module:
	def __init__(self, bot):
		self.setup_db()
		self.irc = bot
		self.tracking = []
		self.lc = LoopingCall(self.lc_callback)
		if bot:
			self.lc.start(config.getfloat(cfg_section, 'interval'), False)
			self.irc.register_keyword('!track', self)

	def stop(self):
		try:
			self.lc.stop()
		except:
			pass

	def setup_db(self):
		global engine, Session

		if engine and Session:
			return

		engine = create_engine(config.get(cfg_section, 'db_path'))
		Base.metadata.create_all(engine)
		Session = sessionmaker(bind = engine, autoflush = True, autocommit = False)

	def track_start(self, code, label, nick, channel):
		type, code = code_split(code)
		if not type in tracking_modules:
			return '\002%s\002 is an invalid type.' % type
		if not len(code):
			return 'Missing tracking number.'

		msg = None
		try:
			session = Session()
			consignment = Consignment(nick, channel, type, code, label.decode('utf8'), datetime.datetime.utcnow())
			session.add(consignment)
			session.commit()
			msg = 'Now tracking \002%s\002.' % code
			reactor.callLater(1, self.track_update, code, True)
		except IntegrityError as e:
			msg = 'Already tracking \002%s\002.' % code
		finally:
			session.close()
		return msg

	def track_stop(self, code, nick, channel):
		type, code = code_split(code)
		if not len(code):
			return 'Missing tracking number or label.'
		msg = None
		try:
			session = Session()
			consignments = session.query(Consignment).filter(or_(Consignment.code == code, Consignment.label == code))
			if nick:
				consignments = consignments.filter(Consignment.nick == nick)
			deleted = []
			# we should have exactly one consignment, but need to handle cases when there are more
			for consignment in consignments:
				code = str(consignment)
				for p in consignment.packages:
					session.delete(p)
				deleted.append(code)
				session.delete(consignment)
			session.commit()
			msg = 'No longer tracking %s' % ', '.join(deleted)
		except NoResultFound:
			msg = '\002%s\002 is not being tracked.' % code.decode('utf8')
		finally:
			session.close()
		return msg

	def track_status(self, code, nick, channel):
		# don't code_split() here; this is done later if no results are returned
		msg = []
		try:
			session = Session()
			consignments = session.query(Consignment).filter(or_(Consignment.code.like(code + '%'), Consignment.label.like((code + '%'))))
			if nick:
				consignments = consignments.filter(Consignment.nick == nick)
			results = []
			for row in consignments:
				try:
					self.track_update(row.code, propagate_error = True)
				except PackageError as e:
					results.append('Failed to fetch data for %s.' % unicode(row))
					continue
				i = 0

				for package in row.packages:
					i += 1
					date = package.last
					desc = package.status

					if date and desc:
						s = '%s %s - %s' % (unicode(package), date, desc)
					else:
						s = 'No tracking info found for %s' % unicode(package)

					results.append(s)

				if i == 0:
					results.append('No packages found for %s' % unicode(row))

			if len(results):
				msg = results
			elif len(code):
				try:
					# type:code is handled by get_tracking_module()
					tm = get_tracking_module(code)
					data = tm.track(code)
					if data:
						for package in data:
							msg.append('\002%s\002 %s - %s' % (package.code, package.date, package.desc.decode('utf8')))
					elif code:
						msg = 'Failed to fetch tracking data for \002%s\002.' % code.decode('utf8')
				except PackageError as e:
					msg = str(e)
			else:
				msg = 'No tracking number given or registered.'
		finally:
			session.close()
		return msg

	def track_label(self, code, label, nick, channel):
		type, code = code_split(code)
		if not len(code):
			return 'Missing tracking number or label.'

		msg = None
		try:
			session = Session()
			consignment = session.query(Consignment).filter_by(code = code).filter_by(nick = nick).one()
			consignment.label = label.decode('utf8')
			session.add(consignment)
			session.commit()
			msg = '\002%s\002 now labelled \002%s\002' % (code, label.decode('utf8'))
		finally:
			session.close()
		return msg

	def track_list(self, nick):
		msg = None
		try:
			session = Session()
			trackings = []
			consignments = session.query(Consignment).filter_by(nick = nick)
			for row in consignments:
				trackings.append(str(row))
			if len(trackings):
				msg = 'Trackings for %s: %s' % (nick, ', '.join(trackings))
			else:
				msg = 'No trackings for %s' % nick
		except NoResultFound:
			msg = 'No trackings registered for %s.'
		finally:
			session.close()
		return msg

	def track_url(self, code, nick, channel):
		msg = []
		try:
			session = Session()
			consignments = session.query(Consignment).filter(or_(Consignment.code.like(code + '%'), Consignment.label.like(code+  '%')))
			if nick:
				consignments = consignments.filter(Consignment.nick == nick)
			results = []
			for consignment in consignments:
				for package in consignment.packages:
					m = get_tracking_module(package.code)
					msg.append(m.get_url(package.code))
			if not len(msg):
				# Rreturn an error only if we got a code to filter with.
				if code:
					msg.append('No packages found.')
				else:
					m = get_tracking_module('')
					msg = [m.get_url()]
		finally:
			session.close()
		return msg

	def keyword(self, nick, channel, kw, msg):
		usage = 'Usage: !track (start|label TRACKINGNO [LABEL])|(stop|status TRACKINGNO|LABEL)|list|(url TRACKINGNO|LABEL)'
		args = msg.split()
		if len(args) < 1:
			self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0], usage)
			return

		mode = args[0]
		if mode.lower() in ('start', 'add', 'stop', 'status'):
			if len(args) < 2 and mode.lower() != 'status':
				self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0], usage)
				return
			if mode.lower() in ('start', 'add'):
				code = args[1]
				label = ' '.join(args[2:]) if len(args) > 2 else ''
			else:
				code = ' '.join(args[1:])
		elif mode.lower() == 'label':
			if len(args) < 3:
				self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0], usage)
				return
			code = args[1]
			label = ' '.join(args[2:])
		elif mode.lower() == 'url':
			code = ' '.join(args[1:]) if len(args) > 1 else ''

		msg = None
		if mode.lower() in ('start', 'add'):
			msg = self.track_start(code, label, nick.split('!')[0], channel if not channel == self.irc.nickname else None)
		elif mode.lower() == 'stop':
			msg = self.track_stop(code, nick.split('!')[0], channel if not channel == self.irc.nickname else None)
		elif mode.lower() == 'status':
			msg = self.track_status(code, nick.split('!')[0], channel if not channel == self.irc.nickname else None)
		elif mode.lower() == 'list':
			msg = self.track_list(nick.split('!')[0])
		elif mode.lower() == 'label':
			msg = self.track_label(code, label, nick.split('!')[0], channel if not channel == self.irc.nickname else None)
		elif mode.lower() == 'url':
			msg = self.track_url(code, nick.split('!')[0], channel if not channel == self.irc.nickname else None)
		else:
			msg = 'Invalid mode "%s".' % mode

		if msg:
			if type(msg) == list:
				for i in msg:
					self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0], i.encode('utf-8'))
			else:
				self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0], msg.encode('utf-8'))
		else:
			self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0], 'No data returned (this is a bug).')

	# called by start and status
	def track_update(self, code, announce = False, propagate_error = False):
		try:
			session = Session()
			consignment = session.query(Consignment).filter(Consignment.code == code).one()
			self.update_consignment(session, consignment, announce, propagate_error)
			session.commit()
		finally:
			session.close()

	def update_consignment(self, session, consignment, announce = True, propagate_error = False):
		now = datetime.datetime.utcnow()
		td = datetime.timedelta(seconds = config.getint(cfg_section, 'cache_time'))
		# return if consignment was cached within cache_time seconds
		if consignment.last and consignment.last > now - td:
			return
		consignment.last = now
		target = consignment.channel or consignment.nick
		target = target.encode('utf-8')
		removed = False
		tm = tracking_modules[consignment.type]()
		try:
			package_results = tm.track(consignment.code) or []
		except PackageError as e:
			if propagate_error:
				raise e
			# ignore, set results to an empty list so that we can reach the stale check at the end
			package_results = []
		for data in package_results:
			try:
				package = session.query(Package).filter(and_(Package.consignment_id == consignment.id, Package.code == data.code)).one()
			except NoResultFound:
				# don't add delivered packages
				if data.delivered:
					continue
				package = Package(consignment.id, data.code)
				session.add(package)
				try:
					session.commit()
				except IntegrityError:
					# assume several packets within the same consignment was added as a consignment
					session.rollback()
					removed = True
					self.irc.msg(target, '%s: \002%s\002 conflicts with another consignment' % (consignment.nick.encode('utf-8'), data.code))
					break
				package = session.query(Package).filter(and_(Package.consignment_id == consignment.id, Package.code == data.code)).one()

			if package.last == None or data.date > package.last:
				code = data.code
				last = data.date
				desc = data.desc
				msg = '%s: %s %s - %s' % (consignment.nick.encode('utf-8'), package, last, desc)
				if data.delivered:
					session.delete(package)
					msg += ' (Package delivered - tracking stopped)'
					removed = True
				else:
					package.last = last
					package.status = desc
					session.add(package)
				if announce:
					self.irc.msg(target, msg)
			elif package.last != None and package.last < datetime.datetime.utcnow() - datetime.timedelta(hours = 24*30):
				msg = '%s: Removing stale package %s' % (consignment.nick.encode('utf-8'), consignment)
				self.irc.msg(target, msg)
				session.delete(package)
				removed = True
		# empty consignment and 30 days old
		if len(consignment.packages) == 0 and consignment.added < datetime.datetime.utcnow() - datetime.timedelta(hours = 24*30):
			msg = '%s: Removing stale consignment %s' % (consignment.nick.encode('utf-8'), consignment)
			self.irc.msg(target, msg)
			session.delete(consignment)
		if removed and len(consignment.packages) == 0:
			msg = '%s: %s is no longer being tracked' % (consignment.nick.encode('utf-8'), consignment)
			self.irc.msg(target, msg)
			session.delete(consignment)

	def lc_callback(self):
		try:
			session = Session()
			consignments = session.query(Consignment).filter(Consignment.channel.in_(config.get(self.irc.factory.server, 'channels').split()))
			for row in consignments:
				self.update_consignment(session, row)
			session.commit()
		finally:
			session.close()

if __name__ == '__main__':
	import sys, ConfigParser, os
	config = ConfigParser.ConfigParser()
	config.read([os.path.expanduser('~/.fot')])
	m = Module(None)
	for code in sys.argv[1:]:
		tm = get_tracking_module(code)
		tracking = tm.track(code)
		if tracking:
			for track in tracking:
				if type(track) == list:
					print '\n'.join(track)
				else:
					print track