dustbot/bot.py

"""
Bot module
Core functionality for a bot instance
Contains a global module map containing loaded modules and their functions
Includes some builtin administrative commands
"""

import logging
import importlib
import sys

from protocol import IrcProtocol

import modules

# global module map, maps registered functions for each loaded module
try: _MOD_MAP
except NameError:
	_MOD_MAP = {}

def command( name ):
	"""
	Command registering decorator.
	Passed name is name of command, decorated function's docstring becomes command's help message.
	Function's module name is automatically appended to docstring.
	Command function gets passed reference to bot instance, array of params, and context info
	containing the nick that sent, channel it was sent in, and reference to connection object.
	"""
	def __deco__( func ):
		# "fixed" module name. removes top-level 'module' package name
		mod_name = func.__module__[func.__module__.find( '.' ) + 1: ]
		logging.info( 'Registering \'%s\' module command: \'%s\' (%s)...',
			mod_name, name, func.__name__ )

		if not func.__doc__: func.__doc__ = 'No help available for command'
		func.__doc__ += ' (Module: {})'.format( mod_name )

		if mod_name not in _MOD_MAP.keys():
			_MOD_MAP[mod_name] = {}

		_MOD_MAP[mod_name][name.lower()] = func

		return func #__inner__ # func

	return __deco__

class Bot:
	""" Bot class """
	def __init__( self, loop, cfg ):
		self._loop = loop
		self.cfg = cfg
		# list of network connections
		self._con = []
		# create and connect all servers in config
		for s in self.cfg['server']:
			ncon = IrcProtocol( self._loop, self.cfg['user'] )
			# add message callback
			ncon.add_message_callback( self._on_irc_message )
			self._con.append( ncon )
			self._loop.create_task( ncon.do_connect( s ) )

	# called when a connected connection receives a full message
	# passes the connection that sent and raw string message
	def _on_irc_message( self, con, raw_msg ):
		words = raw_msg.split()
		if words[1] == 'PRIVMSG':
			chan = words[2]
			raw_user = words[0]
			nick = raw_user[1 : words[0].find( '!' )]
			msg = raw_msg[raw_msg.find( ':' , 1 ) + 1 :]
			# check for command-like message
			# message split into words, first word is command rest are parms list
			pfx = self.cfg['bot']['prefix']
			strp_msg = msg.strip()
			if strp_msg[0 : len( pfx )] == pfx:
				msg_words = strp_msg.split()
				cmd = msg_words[0][len( pfx ):].lower()
				parms = msg_words[1:]
				# check if command is in a registered module
				for k, m in _MOD_MAP.items():
					if cmd in m.keys():
						con.log( 'CMD \'{}\' run in ({}) by <{}> with parms: {}'.format(
							cmd, chan, nick, parms ) )
						self._loop.create_task( m[cmd](
							self, parms, {'con': con,'dst': chan,'nick': nick} ) )

# core commands

# join / part
@command( 'join' )
async def cmd_join( bot, parms, ctx ):
	""" Joins a channel. """

	nick = ctx['nick']
	#dst = ctx['dst']
	con = ctx['con']
	if nick == con.cfg['usr']['owner'] and parms \
	and not (parms[0] in con.cfg['sv']['channels']) and parms[0][0] == '#':
		con.send( 'JOIN {}'.format( parms[0] ) )
		con.cfg['sv']['channels'].append( parms[0] )

@command( 'part' )
async def cmd_part( bot, parms, ctx ):
	""" Leaves a channel. """

	nick = ctx['nick']
	con = ctx['con']
	if nick == con.cfg['usr']['owner'] and parms \
	and parms[0] in con.cfg['sv']['channels']:
		con.send( 'PART {}'.format( parms[0] ) )
		con.cfg['sv']['channels'].remove( parms[0] )

# gets the doc string for specified command
@command( 'help' )
async def cmd_help( b, p, c ):
	""" Shows help message for a specified command. """

	con = c['con']
	dst = c['dst']

	if not p:
		con.say_to( dst, 'No command specified. See list of commands with \";list\".' )

	helpstr = ''

	for k, m in _MOD_MAP.items():
		if p[0].lower() in m.keys():
			helpstr = '> {}: {}'.format( p[0].lower(), ' '.join( m[p[0].lower()].__doc__.split() ) )

	if helpstr:
		con.say_to( dst, helpstr )
	else:
		con.say_to( dst, 'Command not found.' )

# list commands
@command( 'list' )
async def cmd_list( bot, p, ctx ):
	""" Lists all available commands. TODO: module filter """

	con = ctx['con']

	cmdlist = [list( c.keys() ) for c in _MOD_MAP.values()]

	con.say_to( ctx['dst'], '>: {}'.format( cmdlist ) )

# reload a specified module
@command( 'reload' )
async def cmd_import( bot, p, ctx ):
	""" Reloads a previously loaded module. """
	con = ctx['con']
	if ctx['nick'] != con.cfg['usr']['owner']: return
	dst = ctx['dst']
	if p and p[0] in _MOD_MAP.keys():
		prevMod = None
		try:
			prevMod = _MOD_MAP.pop( p[0], None )
			# if not reloading root bot module, re-append the modules. package name
			if p[0] != 'bot': p[0] = 'modules.' + p[0]
			importlib.reload( sys.modules[p[0]] )
			con.say_to( dst, 'Success!' )
		except Exception as e:
			con.say_to( dst, str( e ) )
			_MOD_MAP[p[0]] = prevMod
	else: con.say_to( dst, 'Module not found' )

# connect to a new server
@command( 'connect' )
async def cmd_connect( b, p, c ):
	"""
	Connects to a new network.
	Parameters: 'name' 'host'
	"""

	con = c['con']
	if c['nick'] != con.cfg['usr']['owner']: return
	if not p and len( p ) < 2: return

	new_sv_cfg = {'host':p[1],'port':6697,'name':p[0]}
	ncon = IrcProtocol( b._loop, con.cfg['usr'] )
	# add message callback
	ncon.add_message_callback( b._on_irc_message )
	b._con.append( ncon )
	b._loop.create_task( ncon.do_connect( new_sv_cfg ) )

# disconnect from currently connected server
@command( 'disconnect' )
async def cmd_discon( b, p, c ):
	""" Disconnects from currently connected network """
	con = c['con']
	if c['nick'] != con.cfg['usr']['owner']: return
	con._stop = True
	con.send( 'QUIT :POOF' )
	b._con.remove( con )

# custom modules
importlib.import_module( '.qdb', 'modules' )