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 psycopg2.extensions
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):
		self.__config = SafeConfigParser()

		self.__config.read('locationbot.ini')

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

		servers = []

		try:
			for server in self.__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 = self.__config.get(nick, 'admins').split()
			self.__channels = set(self.__config.get(nick, 'channels').split())
			self.__dsn = self.__config.get(nick, 'dsn')
		except NoOptionError, error:
			sys.exit(error)

		try:
			self.__hostname = self.__config.get(nick, 'hostname')
		except NoOptionError:
			self.__hostname = ''

		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.__nick = None
			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.__nick = bot.__nick
			self.__reloading = True

		self.__quiting = False
		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 __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 __coordinates(self, coordinates):
		return ' http://maps.google.com/maps?q=%f,%f' % coordinates if coordinates else ''

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

		def point(value, cursor):
			if value is not None:
				return tuple(map(lambda a: float(a), re.match(r'^\(([^)]+),([^)]+)\)$', value).groups()))

		psycopg2.extensions.register_type(psycopg2.extensions.new_type((600,), 'point', point), db)

		return db, db.cursor()

	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({
				'join': ('channel', 'join a channel'),
				'part': ('channel [message]', 'part from a channel'),
				'quit': ('[message]', 'quit and do not come back'),
				'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 __join(self, connection, nick, arguments):
		try:
			channel = arguments.split(None, 1)[0]
		except IndexError:
			return self.__help(connection, nick, True, False, 'join')

		connection.join(channel)
		self.__channels.add(channel)
		self.__config.set(self._nickname, 'channels', ' '.join(self.__channels))
		self.__write()
		connection.privmsg(nick, 'successfully joined channel ("%s")' % channel)

	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:
						feature = geo['features'][0]
						coordinates = tuple(reversed(feature['geometry']['coordinates']))
						properties = feature['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, coordinates = point %s, updated = %s where latitude = %s', (new_location, coordinates, updated, latitude))
					db.commit()

				self.__location(nick, channels, old_location, new_location, coordinates)
		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, coordinates = None):
		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, coordinates))

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

		if len(arguments) == 2:
			channel, message = arguments
			message = ':' + message
		elif len(arguments) == 1:
			channel = arguments[0]
			message = ''
		else:
			return self.__help(connection, nick, True, False, 'part')

		if channel in self.__channels:
			connection.part(channel, message)
			self.__channels.remove(channel)
			self.__config.set(self._nickname, 'channels', ' '.join(self.__channels))
			self.__write()
			connection.privmsg(nick, 'successfully parted channel ("%s")' % channel)
		else:
			connection.privmsg(nick, 'not in channel ("%s")' % channel)

	def __quit(self, connection, nick, arguments):
		self.__reloading = True
		self.__quiting = True

		connection.privmsg(nick, 'quiting')
		self.disconnect(arguments)

	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.__nick = 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') + (', coordinates = null, 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, coordinates, 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                     map\x0f')

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

	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' + (', coordinates = 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 __write(self):
		with open('locationbot.ini', 'w') as config:
			self.__config.write(config)

	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 disconnect(self, message = 'oh no!'):
		ircbot.SingleServerIRCBot.disconnect(self, message)

	def error(self, error):
		print error

		self.connection.privmsg(self.__nick, 'an error occured')

	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 == 'join':
				self.__join(connection, nick, arguments)
			elif admin and command == 'part':
				self.__part(connection, nick, arguments)
			elif admin and command == 'quit':
				self.__quit(connection, nick, 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.symmetric_difference(self.channels.keys()):
				if channel in self.__channels:
					self.connection.join(channel)
				else:
					self.connection.part(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, coordinates in self.__locations:
						self.connection.notice(channel, '%s is in %s' % (nick, location + self.__coordinates(coordinates)))

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

		return not self.__quiting

	def success(self):
		self.connection.privmsg(self.__nick, 'successfully reloaded')

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

if __name__ == '__main__':
	os.chdir(os.path.abspath(os.path.dirname(__file__)))

	pid = os.fork()

	if pid != 0:
		with open('locationbot.pid', 'w') as _file:
			_file.write('%u\n' % pid)

		sys.exit(0)

	sys.stdin = open('/dev/null')
	sys.stdout = open('locationbot.log', 'a', 1)
	sys.stderr = sys.stdout

	import locationbot

	bot = locationbot.LocationBot()

	try:
		while bot.start():
			try:
				bot = reload(locationbot).LocationBot(bot)
			except (ImportError, SyntaxError), error:
				bot.error(error)
			else:
				bot.success()
	except KeyboardInterrupt:
		bot.disconnect()

	os.unlink('locationbot.pid')