Set!

locationbot.py

@@ -13,6 +13,7 @@ import irclib
 import os
 import psycopg2
 import pytz
+import re
 import sys
 import time
 import threading
@@ -54,7 +55,7 @@ class LocationBot(ircbot.SingleServerIRCBot):
 						servers.append((host, port, ssl))
 
 			self.__admins = config.get(nick, 'admins').split()
-			self.__channels = config.get(nick, 'channels').split()
+			self.__channels = set(config.get(nick, 'channels').split())
 			self.__dsn = config.get(nick, 'dsn')
 			self.__hostname = config.get(nick, 'hostname')
 		except NoOptionError, error:
@@ -81,6 +82,7 @@ class LocationBot(ircbot.SingleServerIRCBot):
 			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):
@@ -127,7 +129,7 @@ class LocationBot(ircbot.SingleServerIRCBot):
 				'who': ('', 'show who is logged in'),
 			})
 
-		connection.privmsg(nick, 'Command     Arguments               Description')
+		connection.privmsg(nick, '\x02command     arguments               description\x0f')
 
 		def help(command, arguments, description):
 			connection.privmsg(nick, '%-11s %-23s %s' % (command, arguments, description))
@@ -188,10 +190,7 @@ class LocationBot(ircbot.SingleServerIRCBot):
 					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.__location(nick, channels, old_location, new_location)
 		except psycopg2.Error, error:
 			print error
 
@@ -200,6 +199,12 @@ class LocationBot(ircbot.SingleServerIRCBot):
 
 			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
 
@@ -296,6 +301,83 @@ class LocationBot(ircbot.SingleServerIRCBot):
 
 		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 + ' = %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()
@@ -314,7 +396,7 @@ class LocationBot(ircbot.SingleServerIRCBot):
 		else:
 			timezone = pytz.utc
 
-		connection.privmsg(nick, 'Nick                    Location                             When')
+		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')))
@@ -341,13 +423,16 @@ class LocationBot(ircbot.SingleServerIRCBot):
 
 		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 ("' + variable + '") ' + ('' if cursor.rowcount == 1 else 'already ') + 'unset')
+		connection.privmsg(nick, 'variable ("%s") %s unset' % (variable, 'successfuly' if cursor.rowcount == 1 else 'already'))
 
 	def __who(self, connection, nick):
-		connection.privmsg(nick, 'Login Nick              Nick Mask')
+		if self.__logins:
+			connection.privmsg(nick, '\x02login nick              nick mask\x0f')
 
-		for login in sorted(self.__logins.values()):
-			connection.privmsg(nick, '%-23s %s' % login)
+			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:
@@ -422,6 +507,8 @@ class LocationBot(ircbot.SingleServerIRCBot):
 				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':
@@ -457,9 +544,12 @@ class LocationBot(ircbot.SingleServerIRCBot):
 		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 sorted(self.__channels) == sorted(self.channels.keys()):
+				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))