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):
		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)

		self.__locations_lock = threading.Lock()
		self.__locations = []
		self.__logins = {}

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

	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 __do_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]))

		self.__latitude = threading.Timer((self.__latitude_next - datetime.utcnow()).seconds, self.__do_latitude)

		self.__latitude.start()

	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({
				'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 %-22s %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 __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 __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:
			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 == '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 = threading.Thread(None, self.__do_latitude)
		self.__latitude.daemon = True

		self.__latitude.start()
		self._connect()

		while True:
			self.ircobj.process_once(0.2)

			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()

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

if __name__ == '__main__':
	locationbot = LocationBot()

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