locationbot/locationbot.py

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

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

try:
	import simplejson as json
except ImportError:
	import json

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 = set(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.connection = bot.connection
			self.__locations = bot.__locations
			self.__logins = bot.__logins
			self.__reloading = True

		self.__variables = frozenset(['nick', 'secret', 'masks', 'channels', 'timezone', 'location', 'latitude'])
		self.__lists = self.__variables.intersection(['masks', 'channels'])
		self.__unsetable = self.__variables.difference(['nick', 'secret'])

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

	def __db(self):
		db = psycopg2.connect(self.__dsn)

		return db, db.cursor()

	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 = irclib.irc_lower(arguments.split(None, 1)[0].lstrip('!')) if arguments else None
		commands = {
			'help': ('[command]', 'show this help message'),
			'status': ('[nick]', 'show where everybody or a nick is'),
		}

		if not login:
			commands.update({
				'login': ('[nick] [secret]', 'log in as nick with secret or using masks'),
				'register': ('[nick] secret', 'register as nick with secret'),
			})
		else:
			commands.update({
				'logout': ('', 'log out as nick'),
				'set': ('[variable [value]]', 'display or set variables'),
				'unset': ('variable', '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, '\x02command     arguments               description\x0f')

		def help(command, arguments, description):
			connection.privmsg(nick, '%-11s %-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)

		try:
			db, cursor = self.__db()

			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 = json.load(response)
					except (TypeError, ValueError), error:
						print error
						continue

					try:
						properties = geo['features'][0]['properties']
						new_location = properties['reverseGeocode']
						locations[latitude] = new_location
						updated = datetime.fromtimestamp(properties['timeStamp'], pytz.utc)
					except (IndexError, KeyError), error:
						print error
						continue

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

				self.__location(nick, channels, old_location, new_location)
		except psycopg2.Error, error:
			print error

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

			self.__latitude_timer.start()

	def __location(self, nick, channels, old_location, new_location):
		if channels and new_location and new_location != old_location:
			with self.__locations_lock:
				for channel in self.__channels.intersection(channels):
					self.__locations.append((nick, channel, new_location))

	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, cursor = self.__db()

		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)

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

	def __register(self, connection, nick, arguments):
		arguments = arguments.split(None, 1)

		if len(arguments) == 2:
			login, secret = arguments
		elif len(arguments) == 1:
			login = nick
			secret = arguments[0]
		else:
			return self.__help(connection, nick, False, False, 'register')

		db, cursor = self.__db()

		try:
			cursor.execute('insert into locationbot.nick (nick, secret) values (%s, md5(%s))', (login, secret))
			db.commit()
		except psycopg2.IntegrityError:
			return connection.privmsg(nick, 'nick ("%s") is already registered' % login)

		connection.privmsg(nick, 'nick ("%s") sucessfully registered' % login)

	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 __set(self, connection, nickmask, nick, login, arguments):
		arguments = arguments.split(None, 1)

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

		if variable is not None and variable not in self.__variables:
			return self.__unknown_variable(connection, nick, variable)

		db, cursor = self.__db()

		if value is None:
			variables = sorted(self.__variables) if variable is None else [variable]

			cursor.execute('select ' + ', '.join(map(lambda variable: "'%s'" % ('*' * 8) if variable == 'secret' else variable, variables)) + ' from locationbot.nick where nick = %s', (login,))

			connection.privmsg(nick, '\x02variable    value\x0f')

			for variable, value in zip(variables, cursor.fetchone()):
				connection.privmsg(nick, '%-8s    %s' % (variable, ' '.join(value) if isinstance(value, list) else value))
		else:
			def invalid(value, variable = variable):
				connection.privmsg(nick, 'invalid %s ("%s")' % (variable, value))

			if variable in self.__lists:
				value = value.split()

				if variable == 'channels':
					for channel in value:
						if not irclib.is_channel(channel) or channel not in self.__channels:
							return invalid(channel, 'channel')
				elif variable == 'masks':
					_mask = re.compile('^.+!.+@.+$')

					for mask in value:
						if not _mask.match(mask):
							return invalid(mask, 'mask')
			elif variable == 'latitude':
				try:
					value = int(re.sub(r'^(?:http://(?:www\.)?google\.com/latitude/apps/badge/api\?user=)?([0-9]+)(?:&type=.*)?$', r'\1', value))
				except ValueError:
					return invalid(value)
			elif variable == 'location':
				cursor.execute('select channels, location from locationbot.nick where nick = %s', (login,))

				channels, location = cursor.fetchone()
			elif variable == 'nick':
				_nick = value.split(None, 1)

				if len(_nick) != 1:
					return invalid(value)
			elif variable == 'timezone':
				if value not in pytz.all_timezones_set:
					return invalid(value)

			try:
				cursor.execute('update locationbot.nick set ' + variable + ' = ' + ('md5(%s)' if variable == 'secret' else '%s') + (', updated = now()' if variable == 'location' else '') + ' where nick = %s', (value, login))
				db.commit()
			except psycopg2.IntegrityError:
				if variable == 'nick':
					return connection.privmsg(nick, 'nick ("%s") is already registered' % value)

				raise

			if variable == 'nick':
				self.__logins[nick] = (value, nickmask)
			elif variable == 'location':
				self.__location(nick, channels, location, value)

			connection.privmsg(nick, 'variable ("%s") successfully set to value ("%s")' % (variable, ' '.join(value) if isinstance(value, list) else value))

	def __status(self, connection, nick, login, arguments):
		_nick = arguments.split(None, 1)[0] if arguments else None
		db, cursor = self.__db()

		cursor.execute('select nick, location, updated from locationbot.nick where ' + ('nick = %s and ' if _nick is not None else '') + 'location is not null order by updated desc', (_nick,))

		if cursor.rowcount == 0:
			return connection.privmsg(nick, 'no location information for ' + ('"%s"' % _nick if _nick is not None else 'anybody'))

		locations = cursor.fetchall()

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

			timezone = pytz.timezone(cursor.fetchone()[0]) if cursor.rowcount == 1 else pytz.utc
		else:
			timezone = pytz.utc

		connection.privmsg(nick, '\x02nick                    location                             when\x0f')

		for _nick, location, updated in locations:
			connection.privmsg(nick, '%-23s %-36s %s' % (_nick, location, timezone.normalize(updated.astimezone(timezone)).strftime('%Y-%m-%d %H:%M %Z')))

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

	def __unknown_variable(self, connection, nick, variable):
		connection.privmsg(nick, 'unknown variable ("%s")' % variable)

	def __unset(self, connection, nick, login, arguments):
		try:
			variable = irclib.irc_lower(arguments.split(None, 1)[0])
		except IndexError:
			return self.__help(connection, nick, False, login, 'unset')

		if variable not in self.__unsetable:
			if variable in self.__variables:
				return connection.privmsg(nick, 'variable ("%s") is not unsetable' % variable)

			return self.__unknown_variable(connection, nick, variable)

		db, cursor = self.__db()

		cursor.execute('update locationbot.nick set ' + variable + ' = null' + (', updated = null' if variable == 'location' else '') + ' where nick = %s and ' + variable + ' is not null', (login,))
		db.commit()
		connection.privmsg(nick, 'variable ("%s") %s unset' % (variable, 'successfuly' if cursor.rowcount == 1 else 'already'))

	def __who(self, connection, nick):
		if self.__logins:
			connection.privmsg(nick, '\x02login nick              nick mask\x0f')

			for login in sorted(self.__logins.values()):
				connection.privmsg(nick, '%-23s %s' % login)
		else:
			connection.privmsg(nick, 'nobody logged in')

	def _connect(self):
		if len(self.server_list[0]) != 2:
			ssl = self.server_list[0][2]
		else:
			ssl = False

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

		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

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

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

			command = irclib.irc_lower(command.lstrip('!'))

			if command == 'help':
				self.__help(connection, nick, admin, login, arguments)
			elif command == 'login':
				self.__login(connection, nickmask, nick, arguments)
			elif command == 'status':
				self.__status(connection, nick, login, arguments)
			elif not login and command == 'register':
				self.__register(connection, nick, arguments)
			elif login and command == 'logout':
				self.__logout(connection, nick)
			elif login and command == 'set':
				self.__set(connection, nickmask, nick, login, arguments)
			elif login and command == 'unset':
				self.__unset(connection, nick, login, arguments)
			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)
		except psycopg2.Error, error:
			print error

			connection.privmsg(nick, 'an error occurred')

	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

			for channel in self.__channels.difference(self.channels.keys()):
				self.connection.join(channel)

		while not self.__reloading:
			if self.__locations_lock.acquire(False):
				if self.__locations and self.__channels.issubset(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!')