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.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_

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)

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

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

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

class Module:
	url = 'http://sporing.posten.no/sporing.xml?q=%s'

	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 get_xml(self, url):
		try:
			u = urllib2.urlopen(url)
		except:
			return
		if u.headers['content-type'].startswith('text/xml'):
			xml = etree.parse(u)
		else:
			xml = None
		u.close()
		return xml

	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 Exception('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;', '<')
			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
			if city:
				desc = '%s (%s)' % (desc, city.encode('utf8'))
			date = last.find('{%s}OccuredAtDisplayDate' % ns).text + ' ' + last.find('{%s}OccuredAtDisplayTime' % ns).text
			results.append((code, isodate, desc))
		return results

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

	def track_stop(self, code, nick, channel):
		msg = None
		try:
			session = Session()
			consignment = session.query(Consignment).filter(or_(Consignment.code == code, Consignment.label == code)).one()
			code = consignment.code
			for p in consignment.packages:
				session.delete(p)
			label = (' (%s)' % consignment.label) if consignment.label else ''
			session.delete(consignment)
			session.commit()
			msg = 'No longer tracking \002%s\002%s' % (code.decode('utf8'), label)
		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):
		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:
				i = 0
				if row.label:
					label = ' (%s)' % row.label
				else:
					label = ''

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

					if date and desc:
						s = '\002%s\002%s %s - %s' % (package.code, label, date, desc)
					else:
						s = 'No tracking info found for \002%s\002%s' % (package.code, label)

					results.append(s)

				if i == 0:
					results.append('No packages found for \002%s\002%s' % (row.code, label))

			if len(results):
				msg = results
			else:
				try:
					data = self.track(code)
					if data:
						for package in data:
							msg.append('\002%s\002 %s - %s' % (package[0], package[1], package[2].decode('utf8')))
					elif code:
						msg = 'Failed to fetch tracking data for \002%s\002.' % code.decode('utf8')
					else:
						msg = 'No tracking number given or registered.'
				except Exception as e:
					msg = str(e)
		finally:
			session.close()
		return msg

	def track_label(self, code, label, nick, channel):
		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:
				s = row.code
				if row.label:
					s += ' (%s)' % row.label
				trackings.append(s)
			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 keyword(self, nick, channel, kw, msg):
		usage = 'Usage: !track (start|label TRACKINGNO [LABEL])|(stop|status TRACKINGNO|LABEL)|list'
		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', '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() == 'start':
				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:])

		msg = None
		if mode.lower() == 'start':
			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)
		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).')

	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:
				label = (' (%s)' % row.label if row.label else '')
				target = row.channel or row.nick
				target = target.encode('utf-8')
				removed = False
				data_ = self.track(row.code)
				if not data_:
					continue
				for data in data_:
					try:
						package = session.query(Package).filter_by(code = data[0]).one()
					except:
						package = Package(row.id, data[0])
						session.add(package)
						package = session.query(Package).filter_by(code = data[0]).one()

					if package.last == None or data[1] > package.last:
						code = data[0]
						last = data[1]
						desc = data[2]
						msg = '%s: \002%s\002%s %s - %s' % (row.nick.encode('utf-8'), package.code.encode('utf-8'), label.encode('utf-8'), last, desc)
						if desc.startswith('Sendingen er utlevert'):
							session.delete(package)
							msg += ' (Package delivered - tracking stopped)'
							removed = True
						else:
							package.last = last
							package.status = desc
							session.add(package)
						self.irc.msg(target, msg)
				if removed and len(row.packages) == 0:
					msg = '%s: \002%s\002%s is no longer being tracked' % (row.nick, row.code, label)
					self.irc.msg(target, msg.encode('utf-8'))
					session.delete(row)
			session.commit()
		finally:
			session.close()

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