locationbot/locationbot.py

#!/usr/bin/env python
# Location Bot
#
# Douglas Thrift
#
# $Id$

from ConfigParser import NoOptionError, SafeConfigParser
from datetime import datetime, timedelta
import geojson
import getpass
import ircbot
import irclib
import os
import psycopg2
import pytz
import sys
import time
import threading
import urllib, urllib2
import warnings

class LocationBot(ircbot.SingleServerIRCBot):
	def __init__(self, bot = None):
		config = SafeConfigParser({'hostname': ''})

		config.read('locationbot.ini')

		try:
			nick = config.sections()[0]
		except IndexError:
			sys.exit('No nick configured')

		servers = []

		try:
			for server in config.get(nick, 'servers').split():
				server = server.split(':', 2)

				if len(server) == 1:
					servers.append((server[0], 6667))
				else:
					host = server[0]
					port = int(server[1])
					ssl = server[1].startswith('+')

					if len(server) == 3:
						servers.append((host, port, ssl, server[2]))
					else:
						servers.append((host, port, ssl))

			self.__admins = config.get(nick, 'admins').split()
			self.__channels = config.get(nick, 'channels').split()
			self.__dsn = config.get(nick, 'dsn')
			self.__hostname = config.get(nick, 'hostname')
		except NoOptionError, error:
			sys.exit(error)

		ircbot.SingleServerIRCBot.__init__(self, servers, nick, 'Location Bot')

		self.__latitude_timer_lock = threading.Lock()
		self.__latitude_timer = None
		self.__locations_lock = threading.Lock()

		if bot is None:
			self.__locations = []
			self.__logins = {}
			self.__reloading = False
		else:
			irclibobj = self.ircobj.connections[0].irclibobj
			self.ircobj.connections[0] = bot.ircobj.connections[0]
			self.ircobj.connections[0].irclibobj = irclibobj
			self.channels = bot.channels
			self.__locations = bot.__locations
			self.__logins = bot.__logins
			self.__reloading = True

	def __admin(self, nickmask):
		return _or(lambda admin: irclib.mask_matches(nickmask, admin), self.__admins)

	def __channel(self, nick, exclude = None):
		if exclude is not None:
			exclude = irclib.irc_lower(exclude)
			channels = map(lambda channel: channel[1], filter(lambda channel: irclib.irc_lower(channel[0]) == exclude, self.channels))
		else:
			channels = self.channels.values()

		return _or(lambda channel: channel.has_user(nick), channels)

	def __help(self, connection, nick, admin, login, arguments):
		command = arguments.split(None, 1)[0].lstrip('!') if arguments else None
		commands = {
			'help': ('[command]', 'show this help message'),
			'login': ('[nick] [secret]', 'log in as nick with secret or using masks'),
			'register': ('[nick] secret', 'register as nick with secret'),
			'status': ('[nick]', 'show where everybody or a nick is'),
		}

		if login:
			commands.update({
				'latitude': ('[id]', 'shortcut for !set latitude [id]'),
				'logout': ('', 'log out as nick'),
				'set': ('[key [value]]', 'display or set variables'),
				'unset': ('key', 'unset a variable'),
			})

		if admin:
			commands.update({
				'reload': ('', 'reload with more up to date code'),
				'restart': ('', 'quit and join running more up to date code'),
				'say': ('nick|channel message', 'say message to nick or channel'),
				'who': ('', 'show who is logged in'),
			})

		connection.privmsg(nick, 'Command     Arguments               Description')

		def help(command, arguments, description):
			connection.privmsg(nick, '!%-10s %-23s %s' % (command, arguments, description))

		if command in commands:
			help(command, *commands[command])
		else:
			for command, (arguments, description) in sorted(commands.iteritems()):
				help(command, arguments, description)

	def __latitude(self):
		now = datetime.utcnow()

		try:
			while now < self.__latitude_next:
				time.sleep((self.__latitude_next - now).seconds)

				now = datetime.utcnow()
		except AttributeError:
			pass

		self.__latitude_next = now.replace(minute = now.minute - now.minute % 5, second = 0, microsecond = 0) + timedelta(minutes = 5)
		db = psycopg2.connect(self.__dsn)
		cursor = db.cursor()

		cursor.execute('select nick, channels, location, latitude from locationbot.nick where latitude is not null order by latitude')

		locations = {}

		for nick, channels, old_location, latitude in cursor.fetchall():
			new_location = locations.get(latitude)

			if latitude not in locations:
				url = 'http://www.google.com/latitude/apps/badge/api?' + urllib.urlencode({'user': latitude, 'type': 'json'})

				try:
					response = urllib2.urlopen(url)
				except urllib2.URLError, error:
					print error
					continue

				try:
					geo = geojson.load(response, object_hook = geojson.GeoJSON.to_instance)
				except (TypeError, ValueError), error:
					print error
					continue

				if len(geo.features):
					properties = geo.features[0].properties
					new_location = properties['reverseGeocode']
					locations[latitude] = new_location
					updated = datetime.fromtimestamp(properties['timeStamp'], pytz.utc)

					cursor.execute('update locationbot.nick set location = %s, updated = %s where latitude = %s', (new_location, updated, latitude))
					db.commit()

			if channels and new_location and new_location != old_location:
				with self.__locations_lock:
					for channel in frozenset(channels).intersection(self.__channels):
						self.__locations.append((nick, channel, locations[latitude]))

		with self.__latitude_timer_lock:
			self.__latitude_timer = threading.Timer((self.__latitude_next - datetime.utcnow()).seconds, self.__latitude)

			self.__latitude_timer.start()

	def __login(self, connection, nickmask, nick, arguments = ''):
		login = nick

		if connection is not None:
			arguments = arguments.split(None, 1)

			if len(arguments) == 2:
				login, secret = arguments
			elif len(arguments) == 1:
				secret = arguments[0]
			else:
				secret = None
		else:
			secret = None

		if nick in self.__logins:
			login = self.__logins[nick][0]

			if connection is not None:
				return connection.privmsg(nick, 'already logged in as "%s"' % login)

			return login

		db = psycopg2.connect(self.__dsn)
		cursor = db.cursor()

		def success():
			connection.privmsg(nick, 'successfully logged in as "%s"' % login)

		if secret is not None:
			cursor.execute('select true from locationbot.nick where nick = %s and secret = md5(%s)', (login, secret))

			if cursor.rowcount == 1:
				self.__logins[nick] = (login, nickmask)

				return success()

		cursor.execute('select nick, masks from locationbot.nick where nick in (%s, %s)', (login, secret if len(arguments) != 2 else None))

		for login, masks in cursor.fetchall():
			if _or(lambda mask: irclib.mask_matches(nickmask, mask), masks):
				self.__logins[nick] = (login, nickmask)

				return success() if connection else login

		if connection is not None:
			return connection.privmsg(nick, 'failed to log in as "%s"' % login)

		return False

	def __logout(self, connection, nick):
		connection.privmsg(nick, 'logged out as "%s"' % self.__logins.pop(nick)[0])

	def __reload(self, connection, nick):
		self.__reloading = True

		connection.privmsg(nick, 'reloading')

	def __restart(self, connection):
		connection.disconnect('Restarting')
		os.execvp(sys.argv[0], sys.argv)

	def __say(self, connection, nick, arguments):
		try:
			nick_channel, message = arguments.split(None, 1)
		except ValueError:
			return self.__help(connection, nick, True, False, 'say')

		if irclib.is_channel(nick_channel):
			if nick_channel not in self.channels:
				return connection.privmsg(nick, 'not in channel ("%s")' % nick_channel)
		elif not self.__channel(nick_channel):
			return connection.privmsg(nick, 'nick ("%s") not in channel(s)' % nick_channel)
		elif nick_channel == connection.get_nickname():
			return connection.privmsg(nick, 'nice try')

		connection.privmsg(nick_channel, message)

	def __unknown(self, connection, nick, command):
		connection.privmsg(nick, 'unknown command ("!%s"); try "!help"' % command)

	def __who(self, connection, nick):
		connection.privmsg(nick, 'Login Nick                      Nick Mask')

		for login in sorted(self.__logins.values()):
			connection.privmsg(nick, '%-31s %s' % login)

	def _connect(self):
		password = None

		if len(self.server_list[0]) != 2:
			ssl = self.server_list[0][2]

		if len(self.server_list[0]) == 4:
			password = self.server_list[0][3]

		try:
			with warnings.catch_warnings():
				warnings.filterwarnings('ignore', r'socket\.ssl\(\) is deprecated\.  Use ssl\.wrap_socket\(\) instead\.', DeprecationWarning)
				self.connect(self.server_list[0][0], self.server_list[0][1], self._nickname, password, getpass.getuser(), self._realname, localaddress = self.__hostname, ssl = ssl)
		except irclib.ServerConnectionError:
			pass

	def get_version(self):
		return 'locationbot ' + sys.platform

	def on_kick(self, connection, event):
		nick = event.arguments()[0]

		if not self.__channel(nick, event.target()):
			self.__logins.pop(nick, None)

	def on_nick(self, connection, event):
		nickmask = event.source()
		login = self.__logins.pop(irclib.nm_to_n(nickmask), (None,))[0]

		if login is not None:
			nick = event.target()
			self.__logins[nick] = (login, nick + '!' + nm_to_uh(nickmask))

	def on_nicknameinuse(self, connection, event):
		connection.nick(connection.get_nickname() + '_')

	def on_part(self, connection, event):
		nick = irclib.nm_to_n(event.source())

		if not self.__channel(nick, event.target()):
			self.__logins.pop(nick, None)

	def on_privmsg(self, connection, event):
		nickmask = event.source()
		nick = irclib.nm_to_n(nickmask)
		admin = self.__admin(nickmask)

		if not admin and not self.__channel(nick):
			return

		login = self.__login(None, nickmask, nick)

		try:
			command, arguments = event.arguments()[0].split(None, 1)
		except ValueError:
			command = event.arguments()[0].strip()
			arguments = ''

		if command.startswith('!'):
			command = command[1:]

			if command == 'help':
				self.__help(connection, nick, admin, login, arguments)
			elif command == 'login':
				self.__login(connection, nickmask, nick, arguments)
			elif login and command == 'logout':
				self.__logout(connection, nick)
			elif admin and command == 'reload':
				self.__reload(connection, nick)
			elif admin and command == 'restart':
				self.__restart(connection)
			elif admin and command == 'say':
				self.__say(connection, nick, arguments)
			elif admin and command == 'who':
				self.__who(connection, nick)
			else:
				self.__unknown(connection, nick, command)
		elif event.eventtype() == 'privmsg':
			self.__unknown(connection, nick, command)

	def on_quit(self, connection, event):
		self.__logins.pop(irclib.nm_to_n(event.source()), None)

	def on_welcome(self, connection, event):
		for channel in self.__channels:
			connection.join(channel)

	def start(self):
		self.__latitude_thread = threading.Thread(None, self.__latitude)
		self.__latitude_thread.daemon = True

		self.__latitude_thread.start()

		if not self.__reloading:
			self._connect()
		else:
			self.__reloading = False

		while not self.__reloading:
			if self.__locations_lock.acquire(False):
				if self.__locations and sorted(self.__channels) == sorted(self.channels.keys()):
					for nick, channel, location in self.__locations:
						self.connection.notice(channel, '%s is in %s' % (nick, location))

					self.__locations = []

				self.__locations_lock.release()

			self.ircobj.process_once(0.2)

		self.__latitude_thread.join()

		with self.__latitude_timer_lock:
			if self.__latitude_timer is not None:
				self.__latitude_timer.cancel()
				self.__latitude_timer.join()

def _or(function, values):
	return reduce(lambda a, b: a or b, map(function, values) if values else [False])

if __name__ == '__main__':
	import locationbot

	bot = None

	try:
		while True:
			bot = reload(locationbot).LocationBot(bot)

			bot.start()
	except KeyboardInterrupt:
		bot.connection.disconnect('Oh no!')