Mercurial > public > sg101

annotate tools/sg101Bot.py @ 339:b871892264f2

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Adding the sg101 IRC bot code to SVN. This code is pretty rough and needs love, but it gets the job done (one of my first Python apps). This fixes #150.
author Brian Neal <bgneal@gmail.com>
date Sat, 26 Feb 2011 21:27:49 +0000
parents
children
rev   line source
bgneal@339 1 #! /usr/bin/env python
bgneal@339 2 """sg101Bot.py
bgneal@339 3
bgneal@339 4 IRC Bot for SurfGuitar101.com's IRC server.
bgneal@339 5 This bot watches who is in the given channel and updates a MySQL
bgneal@339 6 database accordingly. The database is read by the website to display
bgneal@339 7 who is in the channel.
bgneal@339 8
bgneal@339 9 Brian Neal <brian@surfguitar101.com>
bgneal@339 10 """
bgneal@339 11
bgneal@339 12 import re
bgneal@339 13 import sys
bgneal@339 14 import random
bgneal@339 15 import logging
bgneal@339 16 from optparse import OptionParser
bgneal@339 17 import ConfigParser
bgneal@339 18
bgneal@339 19 import MySQLdb
bgneal@339 20 import irclib
bgneal@339 21 from daemon import Daemon
bgneal@339 22
bgneal@339 23 #irclib.DEBUG = True
bgneal@339 24
bgneal@339 25 class QuoteDb(object):
bgneal@339 26 """QuoteDb is a class that reads a quote file and provdes a random quote when asked"""
bgneal@339 27
bgneal@339 28 def __init__(self, filename):
bgneal@339 29 random.seed()
bgneal@339 30 f = open(filename, 'r')
bgneal@339 31 self.quotes = []
bgneal@339 32 for line in f:
bgneal@339 33 self.quotes.append(line.strip())
bgneal@339 34 f.close()
bgneal@339 35
bgneal@339 36 def get(self):
bgneal@339 37 return random.choice(self.quotes)
bgneal@339 38
bgneal@339 39
bgneal@339 40 class Sg101Table(object):
bgneal@339 41 """Sg101Table class for abstracting the database table that is used to store who is in the channel"""
bgneal@339 42
bgneal@339 43 def __init__(self, dbHost, dbName, dbUser, dbPassword, tableName):
bgneal@339 44 self.dbArgs = {'host' : dbHost, 'user' : dbUser, 'passwd' : dbPassword, 'db' : dbName}
bgneal@339 45 self.table = tableName
bgneal@339 46 self.dbRefCnt = 0
bgneal@339 47
bgneal@339 48 def empty(self):
bgneal@339 49 self.__dbOpen()
bgneal@339 50 """empties the table of channel users"""
bgneal@339 51 logging.debug('emptying the table')
bgneal@339 52 self.cursor.execute('TRUNCATE TABLE ' + self.table)
bgneal@339 53 self.__dbClose()
bgneal@339 54
bgneal@339 55 def add(self, nick):
bgneal@339 56 self.__dbOpen()
bgneal@339 57 """adds the nickname to the table"""
bgneal@339 58 logging.debug('adding to the table: ' + nick)
bgneal@339 59 self.cursor.execute("INSERT INTO " + self.table +
bgneal@339 60 " (name, last_update) VALUES (%s, NOW())", nick)
bgneal@339 61 self.__dbClose()
bgneal@339 62
bgneal@339 63 def remove(self, nick):
bgneal@339 64 self.__dbOpen()
bgneal@339 65 """removes the nickname from the table"""
bgneal@339 66 logging.debug('removing from the table: ' + nick)
bgneal@339 67 self.cursor.execute("DELETE FROM " + self.table + " WHERE `name` = %s", nick)
bgneal@339 68 self.__dbClose()
bgneal@339 69
bgneal@339 70 def set(self, nicks):
bgneal@339 71 self.__dbOpen()
bgneal@339 72 """empties and then adds all the nicknames in the nicks list to the table"""
bgneal@339 73 self.empty()
bgneal@339 74 logging.debug('setting the table: ' + ', '.join(nicks))
bgneal@339 75 for nick in nicks:
bgneal@339 76 self.add(nick)
bgneal@339 77 self.__dbClose()
bgneal@339 78
bgneal@339 79 def __dbOpen(self):
bgneal@339 80 if self.dbRefCnt <= 0:
bgneal@339 81 self.connection = MySQLdb.connect(
bgneal@339 82 host = self.dbArgs['host'],
bgneal@339 83 user = self.dbArgs['user'],
bgneal@339 84 passwd = self.dbArgs['passwd'],
bgneal@339 85 db = self.dbArgs['db'])
bgneal@339 86 self.cursor = self.connection.cursor()
bgneal@339 87 self.dbRefCnt = 1
bgneal@339 88 else:
bgneal@339 89 self.dbRefCnt = self.dbRefCnt + 1
bgneal@339 90
bgneal@339 91 def __dbClose(self):
bgneal@339 92 self.dbRefCnt = self.dbRefCnt - 1
bgneal@339 93 if self.dbRefCnt <= 0:
bgneal@339 94 self.cursor.close()
bgneal@339 95 self.connection.commit()
bgneal@339 96 self.connection.close()
bgneal@339 97 self.dbRefCnt = 0
bgneal@339 98
bgneal@339 99
bgneal@339 100 class Sg101Bot(irclib.SimpleIRCClient):
bgneal@339 101 """Sg101Bot class for monitoring a given channel and updating a database table accordingly"""
bgneal@339 102
bgneal@339 103 def __init__(self, channel, nickname, server, port, password, db, quotes, quoteMod = 5):
bgneal@339 104 """__init__ function for Sg101Bot"""
bgneal@339 105
bgneal@339 106 irclib.SimpleIRCClient.__init__(self)
bgneal@339 107 self.channel = channel
bgneal@339 108 self.nickname = nickname
bgneal@339 109 self.server = server
bgneal@339 110 self.port = port
bgneal@339 111 self.password = password
bgneal@339 112 self.users = []
bgneal@339 113 self.db = db
bgneal@339 114 self.quotes = quotes
bgneal@339 115 self.quoteMod = quoteMod
bgneal@339 116 self.pingCount = 0;
bgneal@339 117
bgneal@339 118 def start(self):
bgneal@339 119 """call this function to start the bot"""
bgneal@339 120 self.db.empty()
bgneal@339 121 self.__connect()
bgneal@339 122 irclib.SimpleIRCClient.start(self)
bgneal@339 123
bgneal@339 124 def on_nicknameinuse(self, c, e):
bgneal@339 125 """called if nick is in use; attempts to change nick"""
bgneal@339 126 c.nick(c.get_nickname() + "_")
bgneal@339 127
bgneal@339 128 def on_welcome(self, c, e):
bgneal@339 129 """called when welcomed to the IRC server; attempts to join channel"""
bgneal@339 130 c.join(self.channel)
bgneal@339 131
bgneal@339 132 def on_namreply(self, c, e):
bgneal@339 133 """called when we hear a namreply; we scan the names and if different than our
bgneal@339 134 info of who is in the channel we update the database"""
bgneal@339 135 me = c.get_nickname()
bgneal@339 136 users = [re.sub('^[@\+]', '', u) for u in e.arguments()[2].split(' ') if u != me]
bgneal@339 137 users.sort()
bgneal@339 138 if (self.users != users):
bgneal@339 139 logging.debug('on_namreply: detecting user list difference')
bgneal@339 140 self.users = users
bgneal@339 141 self.db.set(users)
bgneal@339 142 self.__printUsers()
bgneal@339 143
bgneal@339 144 def on_join(self, c, e):
bgneal@339 145 """called after we or someone else joins the channel; update our list of users and db"""
bgneal@339 146 nick = irclib.nm_to_n(e.source())
bgneal@339 147 if nick != c.get_nickname() and nick not in self.users:
bgneal@339 148 self.__addUser(nick)
bgneal@339 149 self.__printUsers()
bgneal@339 150
bgneal@339 151 def on_part(self, c, e):
bgneal@339 152 """called after we or someone else parts the channel; update our list of users and db"""
bgneal@339 153 nick = irclib.nm_to_n(e.source())
bgneal@339 154 if nick != c.get_nickname():
bgneal@339 155 self.__delUser(nick)
bgneal@339 156 self.__printUsers()
bgneal@339 157
bgneal@339 158 def on_kick(self, c, e):
bgneal@339 159 """called after someone is kicked; update our list of users and db"""
bgneal@339 160 nick = e.arguments()[0]
bgneal@339 161 if nick != c.get_nickname():
bgneal@339 162 self.__delUser(nick)
bgneal@339 163 self.__printUsers()
bgneal@339 164
bgneal@339 165 def on_quit(self, c, e):
bgneal@339 166 """called after someone quits; update our list of users and db"""
bgneal@339 167 nick = irclib.nm_to_n(e.source())
bgneal@339 168 if nick != c.get_nickname():
bgneal@339 169 self.__delUser(nick)
bgneal@339 170 self.__printUsers()
bgneal@339 171
bgneal@339 172 def on_ping(self, c, e):
bgneal@339 173 """called when we are pinged; we use this opportunity to ask who is in the channel via NAMES"""
bgneal@339 174 c.names()
bgneal@339 175 self.pingCount = self.pingCount + 1
bgneal@339 176 if (self.quoteMod > 0) and len(self.users) > 0 and (self.pingCount % self.quoteMod) == 0:
bgneal@339 177 c.privmsg(self.channel, self.quotes.get())
bgneal@339 178
bgneal@339 179 def on_privmsg(self, c, e):
bgneal@339 180 self.doCommand(e, e.arguments()[0], e.source())
bgneal@339 181
bgneal@339 182 def on_pubmsg(self, c, e):
bgneal@339 183 a = e.arguments()[0].split(':', 1)
bgneal@339 184 if len(a) > 1 and irclib.irc_lower(a[0]) == irclib.irc_lower(c.get_nickname()):
bgneal@339 185 self.doCommand(e, a[1].strip(), e.target())
bgneal@339 186
bgneal@339 187 def on_disconnect(self, c, e):
bgneal@339 188 logging.critical('received on_disconnect()')
bgneal@339 189 logging.debug('scheduling a routine in 30 secs')
bgneal@339 190 self.users = []
bgneal@339 191 self.db.empty()
bgneal@339 192 self.connection.execute_delayed(30, self.__connectedChecker)
bgneal@339 193
bgneal@339 194 def __connectedChecker():
bgneal@339 195 logging.debug('__connectedChecker()')
bgneal@339 196 if not self.connection.is_connected():
bgneal@339 197 logging.debug('rescheduling __connectedChecker in 30 secs')
bgneal@339 198 self.connection.execute_delayed(30, self._connected_checker)
bgneal@339 199 self.__connect()
bgneal@339 200
bgneal@339 201
bgneal@339 202 def doCommand(self, e, cmd, reply):
bgneal@339 203 if cmd == 'yow':
bgneal@339 204 self.connection.privmsg(reply, self.quotes.get())
bgneal@339 205
bgneal@339 206 def __connect(self):
bgneal@339 207 """our internal connect function"""
bgneal@339 208 self.connect(self.server,
bgneal@339 209 self.port,
bgneal@339 210 self.nickname,
bgneal@339 211 self.password)
bgneal@339 212
bgneal@339 213 def __printUsers(self):
bgneal@339 214 """internal print users command"""
bgneal@339 215 msg = 'My users are: (' + ', '.join(self.users) + ')'
bgneal@339 216 logging.debug(msg)
bgneal@339 217
bgneal@339 218 def __addUser(self, nick):
bgneal@339 219 """adds a user to our list and db"""
bgneal@339 220 self.users.append(nick)
bgneal@339 221 self.users.sort()
bgneal@339 222 self.db.add(nick)
bgneal@339 223
bgneal@339 224 def __delUser(self, nick):
bgneal@339 225 """removes a user from our list and db"""
bgneal@339 226 if nick in self.users:
bgneal@339 227 self.users.remove(nick)
bgneal@339 228 self.db.remove(nick)
bgneal@339 229
bgneal@339 230
bgneal@339 231 class BotDaemon(Daemon):
bgneal@339 232 def __init__(self, filename, options, db_config):
bgneal@339 233 Daemon.__init__(self, filename)
bgneal@339 234 self.options = options
bgneal@339 235 self.db_config = db_config
bgneal@339 236
bgneal@339 237 def run(self):
bgneal@339 238 logging.basicConfig(level = logging.DEBUG,
bgneal@339 239 format = '%(asctime)s %(levelname)s %(message)s',
bgneal@339 240 filename = '/home/brian/irc/sg101Bot.log',
bgneal@339 241 filemode = 'w')
bgneal@339 242 logging.info('sg101Bot starting')
bgneal@339 243
bgneal@339 244 server = self.options.server
bgneal@339 245 port = self.options.port
bgneal@339 246 channel = self.options.channel
bgneal@339 247 nickname = self.options.nick
bgneal@339 248 password = self.options.password
bgneal@339 249
bgneal@339 250 try:
bgneal@339 251 quotes = QuoteDb('/home/brian/irc/zippy.txt')
bgneal@339 252 db = Sg101Table(self.db_config['host'],
bgneal@339 253 self.db_config['database'],
bgneal@339 254 self.db_config['user'],
bgneal@339 255 self.db_config['password'],
bgneal@339 256 self.db_config['table'])
bgneal@339 257 bot = Sg101Bot(channel, nickname, server, port, password, db, quotes, 0)
bgneal@339 258 bot.start()
bgneal@339 259 except MySQLdb.OperationalError, message:
bgneal@339 260 logging.exception('MySQLdb.OperationalError: %d %s' % (message[0], message[1]))
bgneal@339 261 except MySQLdb.ProgrammingError, message:
bgneal@339 262 logging.exception('MySQLdb.ProgrammingError: %d %s' % (message[0], message[1]))
bgneal@339 263 except:
bgneal@339 264 print "Got an exception:", sys.exc_info()[0]
bgneal@339 265 logging.exception('Exception received ' + str(sys.exc_info()))
bgneal@339 266
bgneal@339 267 logging.info('sg101Bot ending')
bgneal@339 268
bgneal@339 269 if __name__ == "__main__":
bgneal@339 270 parser = OptionParser(usage='usage: %prog [options] start|stop|restart',
bgneal@339 271 description="""\
bgneal@339 272 SG101 IRC Bot. Monitors who is in the specified channel and
bgneal@339 273 updates a database table accordingly for display by the website.
bgneal@339 274 """)
bgneal@339 275
bgneal@339 276 parser.add_option("-s", "--server", default="localhost",
bgneal@339 277 help="server address")
bgneal@339 278 parser.add_option("-p", "--port", type="int", default=6667,
bgneal@339 279 help="port number")
bgneal@339 280 parser.add_option("-c", "--channel", default="#ShallowEnd",
bgneal@339 281 help="channel name")
bgneal@339 282 parser.add_option("-n", "--nick", default="bot",
bgneal@339 283 help="bot nickname")
bgneal@339 284 parser.add_option("-w", "--password", default="morereverb",
bgneal@339 285 help="channel password")
bgneal@339 286 parser.add_option("-q", "--quote", default="/home/brian/irc/zippy.txt",
bgneal@339 287 help="Quote file")
bgneal@339 288 (options, args) = parser.parse_args()
bgneal@339 289
bgneal@339 290 commands = ('start', 'stop', 'restart')
bgneal@339 291 if len(args) != 1 and args[0] not in commands:
bgneal@339 292 parser.print_help()
bgneal@339 293 sys.exit("Please provide a single command: start, stop, or restart.")
bgneal@339 294
bgneal@339 295 config = ConfigParser.ConfigParser()
bgneal@339 296 config.read('/home/brian/irc/sg101Bot.ini')
bgneal@339 297 db_config = dict(config.items('Database'))
bgneal@339 298
bgneal@339 299 daemon = BotDaemon('/tmp/sg101Bot.pid', options, db_config)
bgneal@339 300 cmd = args[0]
bgneal@339 301 if cmd == 'start':
bgneal@339 302 daemon.start()
bgneal@339 303 elif cmd == 'stop':
bgneal@339 304 daemon.stop()
bgneal@339 305 elif cmd == 'restart':
bgneal@339 306 daemon.restart()