Login/logout!

locationbot.py

@@ -13,6 +13,7 @@ import ircbot
 import irclib
 import os
 import psycopg2
+import pytz
 import sys
 import time
 import threading
@@ -57,18 +58,21 @@ class LocationBot(ircbot.SingleServerIRCBot):
 
 		self.__locations_lock = threading.Lock()
 		self.__locations = []
+		self.__logins = {}
 
 		ircbot.SingleServerIRCBot.__init__(self, servers, nick, 'Location Bot')
 
-	def __admin(self, nick):
-		if isinstance(nick, irclib.Event):
-			nick = nick.source()
+	def __admin(self, nickmask):
+		return _or(lambda admin: irclib.mask_matches(nickmask, admin), self.__admins)
 
-		for admin in self.__admins:
-			if irclib.mask_matches(nick, admin):
-				return True
+	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 False
+		return _or(lambda channel: channel.has_user(nick), channels)
 
 	def __do_latitude(self):
 		now = datetime.utcnow()
@@ -89,19 +93,34 @@ class LocationBot(ircbot.SingleServerIRCBot):
 
 		locations = {}
 
-		for nick, channels, location, latitude in cursor.fetchall():
+		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'})
-				response = urllib2.urlopen(url)
-				geo = geojson.load(response, object_hook = geojson.GeoJSON.to_instance)
+
+				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):
-					locations[latitude] = geo.features[0].properties['reverseGeocode']
+					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 where latitude = %s', (locations[latitude], latitude))
+					cursor.execute('update locationbot.nick set location = %s, updated = %s where latitude = %s', (new_location, updated, latitude))
 					db.commit()
 
-			if channels and locations.get(latitude, location) != location:
+			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]))
@@ -110,33 +129,93 @@ class LocationBot(ircbot.SingleServerIRCBot):
 
 		self.__latitude.start()
 
-	def __help(self, connection, nick, admin, arguments):
-		connection.privmsg(nick, 'Command     Arguments        Description')
+	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'),
+			})
 
-		commands = [
-			('help', '', 'show this help message'),
-			('login', '[nick] secret', 'login as nick with secret'),
-			('register', '[nick] secret', 'register as nick with secret'),
-			('status', '[nick]', 'show where everybody or a nick is'),
-		]
+		connection.privmsg(nick, 'Command     Arguments              Description')
 
-		if True:
-			commands += [
-				('latitude', '[id]', 'shortcut for !set latitude [id]'),
-				('set', '[key [value]]', 'display or set variables'),
-				('unset', 'key', 'unset a variable'),
-			]
+		def help(command, arguments, description):
+			connection.privmsg(nick, '!%-10s %-22s %s' % (command, arguments, description))
 
-		if admin:
-			commands += [
-				('restart', '', 'quit and join running more up to date code'),
-				('say', 'channel message', 'say message in channel'),
-			]
+		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)
 
-		commands.sort()
+				return success()
 
-		for command in commands:
-			connection.privmsg(nick, '!%-10s %-16s %s' % command)
+		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')
@@ -144,14 +223,28 @@ class LocationBot(ircbot.SingleServerIRCBot):
 
 	def __say(self, connection, nick, arguments):
 		try:
-			channel, message = arguments.split(None, 1)
-
-			connection.privmsg(channel, message)
+			nick_channel, message = arguments.split(None, 1)
 		except ValueError:
-			self.__help(connection, nick, True, 'say')
+			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)
+		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
@@ -172,24 +265,38 @@ class LocationBot(ircbot.SingleServerIRCBot):
 	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_privmsg(self, connection, event):
+	def on_part(self, connection, event):
 		nick = irclib.nm_to_n(event.source())
-		admin = self.__admin(event)
 
-		if not admin:
-			ok = False
+		if not self.__channel(nick, event.target()):
+			self.__logins.pop(nick, None)
 
-			for channel in self.channels.values():
-				if channel.has_user(nick):
-					ok = True
+	def on_privmsg(self, connection, event):
+		nickmask = event.source()
+		nick = irclib.nm_to_n(nickmask)
+		admin = self.__admin(nickmask)
 
-					break
+		if not admin and not self.__channel(nick):
+			return
 
-			if not ok:
-				return
+		login = self.__login(None, nickmask, nick)
 
 		try:
 			command, arguments = event.arguments()[0].split(None, 1)
@@ -201,19 +308,25 @@ class LocationBot(ircbot.SingleServerIRCBot):
 			command = command[1:]
 
 			if command == 'help':
-				self.__help(connection, nick, admin, arguments)
-			elif admin:
-				if command == 'restart':
-					self.__restart(connection)
-				elif command == 'say':
-					self.__say(connection, nick, arguments)
-				else:
-					self.__unknown(connection, nick, command)
+				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)
@@ -237,6 +350,9 @@ class LocationBot(ircbot.SingleServerIRCBot):
 
 				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()
 

locationbot.sql

@@ -13,8 +13,10 @@ create table locationbot.nick (
 	masks text[],
 	channels text[],
 	location text,
+	updated timestamp with time zone,
 	latitude bigint
 );
 
 create index nick_location_index on locationbot.nick (location);
+create index nick_updated_index on locationbot.nick (updated);
 create index nick_latitude_index on locationbot.nick (latitude);