add sqlite database support

[erebus.git] / erebus.py

#!/usr/bin/python
# vim: fileencoding=utf-8

# Erebus IRC bot - Author: John Runyon
# main startup code

from __future__ import print_function

import os, sys, select, time, traceback, random, gc
import bot, config, ctlmod

class Erebus(object): #singleton to pass around
	APIVERSION = 0
	RELEASE = 0

	bots = {}
	fds = {}
	numhandlers = {}
	msghandlers = {}
	chanhandlers = {}
	exceptionhandlers = [] # list of (Exception_class, handler_function) tuples
	users = {}
	chans = {}

	class User(object):
		def __init__(self, nick, auth=None):
			self.nick = nick
			if auth is None:
				self.auth = None
			else:
				self.auth = auth.lower()
			self.checklevel()

			self.chans = []

		def bind_bot(self, bot):
			return main._BoundUser(self, bot)

		def msg(self, *args, **kwargs):
			main.randbot().msg(self, *args, **kwargs)
		def slowmsg(self, *args, **kwargs):
			main.randbot().slowmsg(self, *args, **kwargs)
		def fastmsg(self, *args, **kwargs):
			main.randbot().fastmsg(self, *args, **kwargs)

		def isauthed(self):
			return self.auth is not None

		def authed(self, auth):
			if auth == '0': self.auth = None
			else: self.auth = auth.lower()
			self.checklevel()

		def checklevel(self):
			if self.auth is None:
				self.glevel = -1
			else:
				c = main.query("SELECT level FROM users WHERE auth = %s", (self.auth,))
				if c:
					row = c.fetchone()
					if row is not None:
						self.glevel = row['level']
					else:
						self.glevel = 0
				else:
					self.glevel = 0
			return self.glevel

		def setlevel(self, level, savetodb=True):
			if savetodb:
				if level != 0:
					c = main.query("REPLACE INTO users (auth, level) VALUES (%s, %s)", (self.auth, level))
				else:
					c = main.query("DELETE FROM users WHERE auth = %s", (self.auth,))
					if c == 0: # no rows affected
						c = True # is fine
				if c:
					self.glevel = level
					return True
				else:
					return False
			else:
				self.glevel = level
				return True

		def join(self, chan):
			if chan not in self.chans: self.chans.append(chan)
		def part(self, chan):
			try:
				self.chans.remove(chan)
			except: pass
			return len(self.chans) == 0
		def quit(self):
			pass
		def nickchange(self, newnick):
			self.nick = newnick

		def __str__(self): return self.nick
		def __repr__(self): return "<User %r (%d)>" % (self.nick, self.glevel)

	class _BoundUser(object):
		def __init__(self, user, bot):
			self.__dict__['_bound_user'] = user
			self.__dict__['_bound_bot'] = bot
		def __getattr__(self, name):
			return getattr(self._bound_user, name)
		def __setattr__(self, name, value):
			setattr(self._bound_user, name, value)
		def msg(self, *args, **kwargs):
			self._bound_bot.msg(self._bound_user, *args, **kwargs)
		def slowmsg(self, *args, **kwargs):
			self._bound_bot.slowmsg(self._bound_user, *args, **kwargs)
		def fastmsg(self, *args, **kwargs):
			self._bound_bot.fastmsg(self._bound_user, *args, **kwargs)
		def __repr__(self): return "<_BoundUser %r %r>" % (self._bound_user, self._bound_bot)

	class Channel(object):
		def __init__(self, name, bot):
			self.name = name
			self.bot = bot
			self.levels = {}

			self.users = []
			self.voices = []
			self.ops = []

			c = main.query("SELECT user, level FROM chusers WHERE chan = %s", (self.name,))
			if c:
				row = c.fetchone()
				while row is not None:
					self.levels[row['user']] = row['level']
					row = c.fetchone()


		def msg(self, *args, **kwargs):
			self.bot.msg(self, *args, **kwargs)
		def slowmsg(self, *args, **kwargs):
			self.bot.slowmsg(self, *args, **kwargs)
		def fastmsg(self, *args, **kwargs):
			self.bot.fastmsg(self, *args, **kwargs)

		def levelof(self, auth):
			if auth is None:
				return 0
			auth = auth.lower()
			if auth in self.levels:
				return self.levels[auth]
			else:
				return 0

		def setlevel(self, auth, level, savetodb=True):
			auth = auth.lower()
			if savetodb:
				c = main.query("REPLACE INTO chusers (chan, user, level) VALUES (%s, %s, %s)", (self.name, auth, level))
				if c:
					self.levels[auth] = level
					return True
				else:
					return False
			else:
				self.levels[auth] = level
				return True

		def userjoin(self, user, level=None):
			if user not in self.users: self.users.append(user)
			if level == 'op' and user not in self.ops: self.ops.append(user)
			if level == 'voice' and user not in self.voices: self.voices.append(user)
		def userpart(self, user):
			if user in self.ops: self.ops.remove(user)
			if user in self.voices: self.voices.remove(user)
			if user in self.users: self.users.remove(user)

		def userop(self, user):
			if user in self.users and user not in self.ops: self.ops.append(user)
		def uservoice(self, user):
			if user in self.users and user not in self.voices: self.voices.append(user)
		def userdeop(self, user):
			if user in self.ops: self.ops.remove(user)
		def userdevoice(self, user):
			if user in self.voices: self.voices.remove(user)

		def __str__(self): return self.name
		def __repr__(self): return "<Channel %r>" % (self.name)

	def __init__(self, cfg):
		self.mustquit = None
		self.starttime = time.time()
		self.cfg = cfg
		self.trigger = cfg.trigger
		if os.name == "posix":
			self.potype = "poll"
			self.po = select.poll()
		else: # f.e. os.name == "nt" (Windows)
			self.potype = "select"
			self.fdlist = []

	def query(self, sql, parameters=[], noretry=False):
		# Callers use %s-style (paramstyle='format') placeholders in queries.
		# There's no provision for a literal '%s' present inside the query; stuff it in a parameter instead.
		if db_api.paramstyle == 'format' or db_api.paramstyle == 'pyformat': # mysql, postgresql
			# psycopg actually asks for a mapping with %(name)s style (pyformat) but it will accept %s style.
			pass
		elif db_api.paramstyle == 'qmark': # sqlite doesn't like %s style.
			parameters = [str(p) for p in parameters]
			sql = sql.replace('%s', '?') # hope that wasn't literal, oopsie

		log_noretry = ''
		if noretry:
			log_noretry = ', noretry=True'
		self.log("[SQL]", "?", "query(%r, %r%s)" % (sql, parameters, log_noretry))

		try:
			curs = self.db.cursor()
			res = curs.execute(sql, parameters)
			if res:
				return curs
			else:
				return res
		except db_api.DataError as e:
			self.log("[SQL]", ".", "DB DataError: %r" % (e))
			return False
		except db_api.Error as e:
			self.log("[SQL]", "!", "DB error! %r" % (e))
			if not noretry:
				dbsetup()
				return self.query(sql, parameters, noretry=True)
			else:
				raise e

	def querycb(self, cb, *args, **kwargs):
		# TODO this should either get thrown out with getdb()/returndb(), or else be adjusted to make use of it.
		def run_query():
			cb(self.query(*args, **kwargs))
		threading.Thread(target=run_query).start()

	def newbot(self, nick, user, bind, authname, authpass, server, port, realname):
		if bind is None: bind = ''
		obj = bot.Bot(self, nick, user, bind, authname, authpass, server, port, realname)
		self.bots[nick.lower()] = obj

	def newfd(self, obj, fileno):
		self.fds[fileno] = obj
		if self.potype == "poll":
			self.po.register(fileno, select.POLLIN)
		elif self.potype == "select":
			self.fdlist.append(fileno)
	def delfd(self, fileno):
		del self.fds[fileno]
		if self.potype == "poll":
			self.po.unregister(fileno)
		elif self.potype == "select":
			self.fdlist.remove(fileno)

	def bot(self, name): #get Bot() by name (nick)
		return self.bots[name.lower()]
	def fd(self, fileno): #get Bot() by fd/fileno
		return self.fds[fileno]
	def randbot(self): #get Bot() randomly
		return self.bots[random.choice(list(self.bots.keys()))]

	def user(self, _nick, send_who=False, create=True):
		nick = _nick.lower()

		if send_who and (nick not in self.users or not self.users[nick].isauthed()):
			self.randbot().conn.send("WHO %s n%%ant,1" % (nick))

		if nick in self.users:
			return self.users[nick]
		elif create:
			user = self.User(_nick)
			self.users[nick] = user
			return user
		else:
			return None
	def channel(self, name): #get Channel() by name
		if name.lower() in self.chans:
			return self.chans[name.lower()]
		else:
			return None

	def newchannel(self, bot, name):
		chan = self.Channel(name.lower(), bot)
		self.chans[name.lower()] = chan
		return chan

	def poll(self):
		timeout_seconds = 30
		if self.potype == "poll":
			pollres = self.po.poll(timeout_seconds * 1000)
			return [fd for (fd, ev) in pollres]
		elif self.potype == "select":
			return select.select(self.fdlist, [], [], timeout_seconds)[0]

	def connectall(self):
		for bot in self.bots.values():
			if bot.conn.state == 0:
				bot.connect()

	def module(self, name):
		return ctlmod.modules[name]

	def log(self, source, level, message):
		print("%09.3f %s [%s] %s" % (time.time() % 100000, source, level, message))

	def getuserbyauth(self, auth):
		return [u for u in self.users.values() if u.auth == auth.lower()]

	def getdb(self):
		"""Get a DB object. The object must be returned to the pool after us, using returndb(). This is intended for use from child threads.
			It should probably be treated as deprecated though. Where possible new modules should avoid using threads.
			In the future, timers will be provided (manipulating the timeout_seconds of the poll() method), and that should mostly be used in place of threading."""
		return self.dbs.pop()

	def returndb(self, db):
		self.dbs.append(db)

	#bind functions
	def hook(self, word, handler):
		try:
			self.msghandlers[word].append(handler)
		except:
			self.msghandlers[word] = [handler]
	def unhook(self, word, handler):
		if word in self.msghandlers and handler in self.msghandlers[word]:
			self.msghandlers[word].remove(handler)
	def hashook(self, word):
		return word in self.msghandlers and len(self.msghandlers[word]) != 0
	def gethook(self, word):
		return self.msghandlers[word]

	def hooknum(self, word, handler):
		try:
			self.numhandlers[word].append(handler)
		except:
			self.numhandlers[word] = [handler]
	def unhooknum(self, word, handler):
		if word in self.numhandlers and handler in self.numhandlers[word]:
			self.numhandlers[word].remove(handler)
	def hasnumhook(self, word):
		return word in self.numhandlers and len(self.numhandlers[word]) != 0
	def getnumhook(self, word):
		return self.numhandlers[word]

	def hookchan(self, chan, handler):
		try:
			self.chanhandlers[chan].append(handler)
		except:
			self.chanhandlers[chan] = [handler]
	def unhookchan(self, chan, handler):
		if chan in self.chanhandlers and handler in self.chanhandlers[chan]:
			self.chanhandlers[chan].remove(handler)
	def haschanhook(self, chan):
		return chan in self.chanhandlers and len(self.chanhandlers[chan]) != 0
	def getchanhook(self, chan):
		return self.chanhandlers[chan]

	def hookexception(self, exc, handler):
		self.exceptionhandlers.append((exc, handler))
	def unhookexception(self, exc, handler):
		self.exceptionhandlers.remove((exc, handler))
	def hasexceptionhook(self, exc):
		return any((True for x,h in self.exceptionhandlers if isinstance(exc, x)))
	def getexceptionhook(self, exc):
		return (h for x,h in self.exceptionhandlers if isinstance(exc, x))


def dbsetup():
	main.db = None
	main.dbs = []
	dbtype = cfg.get('erebus', 'dbtype', 'mysql')
	if dbtype == 'mysql':
		_dbsetup_mysql()
	elif dbtype == 'sqlite':
		_dbsetup_sqlite()
	else:
		main.log('*', '!', 'Unknown dbtype in config: %s' % (dbtype))

def _dbsetup_mysql():
	global db_api
	import MySQLdb as db_api, MySQLdb.cursors
	for i in range(cfg.get('erebus', 'num_db_connections', 2)-1):
		main.dbs.append(db_api.connect(host=cfg.dbhost, user=cfg.dbuser, passwd=cfg.dbpass, db=cfg.dbname, cursorclass=MySQLdb.cursors.DictCursor))
	main.db = db_api.connect(host=cfg.dbhost, user=cfg.dbuser, passwd=cfg.dbpass, db=cfg.dbname, cursorclass=MySQLdb.cursors.DictCursor)

def _dbsetup_sqlite():
	global db_api
	import sqlite3 as db_api
	for i in range(cfg.get('erebus', 'num_db_connections', 2)):
		main.db = db_api.connect(cfg.dbhost)
		main.db.row_factory = db_api.Row
		main.db.isolation_level = None
		main.dbs.append(main.db)

def setup():
	global cfg, main

	cfg = config.Config('bot.config')

	if cfg.getboolean('debug', 'gc'):
		gc.set_debug(gc.DEBUG_LEAK)

	pidfile = open(cfg.pidfile, 'w')
	pidfile.write(str(os.getpid()))
	pidfile.close()

	main = Erebus(cfg)
	dbsetup()

	autoloads = [mod for mod, yes in cfg.items('autoloads') if int(yes) == 1]
	for mod in autoloads:
		ctlmod.load(main, mod)

	c = main.query("SELECT nick, user, bind, authname, authpass FROM bots WHERE active = 1")
	if c:
		rows = c.fetchall()
		c.close()
		for row in rows:
			main.newbot(row['nick'], row['user'], row['bind'], row['authname'], row['authpass'], cfg.host, cfg.port, cfg.realname)
	main.connectall()

def loop():
	poready = main.poll()
	for fileno in poready:
		try:
			data = main.fd(fileno).getdata()
		except:
			main.log('*', '!', 'Super-mega-emergency: getdata raised exception for socket %d' % (fileno))
			traceback.print_exc()
			data = None
		if data is None:
			main.fd(fileno).close()
		else:
			for line in data:
				if cfg.getboolean('debug', 'io'):
					main.log(str(main.fd(fileno)), 'I', line)
				try:
					main.fd(fileno).parse(line)
				except:
					main.log('*', '!', 'Super-mega-emergency: parse raised exception for socket %d data %r' % (fileno, line))
					traceback.print_exc()
	if main.mustquit is not None:
		main.log('*', '!', 'Core exiting due to: %s' % (main.mustquit))
		raise main.mustquit

if __name__ == '__main__':
	try: os.rename('logfile', 'oldlogs/%s' % (time.time()))
	except: pass
	sys.stdout = open('logfile', 'w', 1)
	sys.stderr = sys.stdout
	setup()
	while True: loop()