dr.botzo/dr.botzo.py

#!/usr/bin/env python2.6

from ConfigParser import ConfigParser, NoSectionError, NoOptionError
from datetime import datetime
import os
import random
import re
import sys

from dateutil.parser import *
from dateutil.relativedelta import *
from dateutil.tz import *
import irclib

#####
# sub_join_channel
# join a channel when told to by an admin
#####

def sub_join_channel(connection, event, nick, userhost, replypath, what, admin_unlocked):
    whats = what.split(' ')
    if whats[0] == 'join' and admin_unlocked and len(whats) >= 2:
        channel = whats[1]
        if irclib.is_channel(channel):
            connection.join(channel)
            connection.privmsg(replypath, 'joined ' + channel)

#####
# sub_part_channel
# leave a channel when told to by an admin. optionally provide a message
#####

def sub_part_channel(connection, event, nick, userhost, replypath, what, admin_unlocked):
    whats = what.split(' ')
    if whats[0] == 'part' and admin_unlocked and len(whats) >= 2:
        channel = whats[1]
        if irclib.is_channel(channel):
            connection.part(channel, ' '.join(whats[2:]))
            connection.privmsg(replypath, 'parted ' + channel)

#####
# sub_quit_channel
# quit irc server. optionally provide a message
#####

def sub_quit_channel(connection, event, nick, userhost, replypath, what, admin_unlocked):
    whats = what.split(' ')
    if whats[0] == 'quit' and admin_unlocked:
        connection.privmsg(replypath, 'quitting')
        connection.quit(' '.join(whats[1:]))
        with open('dr.botzo.cfg', 'w') as cfg:
            config.write(cfg)

#####
# sub_handle_autojoin
# add/remove on channel autojoin list.
#####

def sub_autojoin_manipulate(connection, event, nick, userhost, replypath, what, admin_unlocked):
    whats = what.split(' ')
    if whats[0] == 'autojoin' and admin_unlocked and len(whats) >= 3:
        if whats[1] == 'add':
            try:
                # get existing list
                channel = whats[2]
                if irclib.is_channel(channel):
                    channelset = set(config.get('channels', 'autojoin').split(','))
                    channelset.add(channel)
                    config.set('channels', 'autojoin', ','.join(channelset))
                    connection.privmsg(replypath, 'added ' + channel + ' to autojoin')
            except NoOptionError: pass
        elif whats[1] == 'remove':
            try:
                # get existing list
                channel = whats[2]
                if irclib.is_channel(channel):
                    channelset = set(config.get('channels', 'autojoin').split(','))
                    channelset.discard(channel)
                    config.set('channels', 'autojoin', ','.join(channelset))
                    connection.privmsg(replypath, 'removed ' + channel + ' from autojoin')
            except NoOptionError: pass

#####
# sub_add_to_seen
# when someone says a pubmsg, keep it in the config
#####

def sub_add_to_seen(connection, event, nick, userhost, what):
    if not config.has_section('seen'):
        config.add_section('seen')

    config.set('seen', nick, userhost + '|:|' + datetime.utcnow().isoformat() + '|:|' + what)

#####
# sub_report_seen
# report when a person has been seen, based on config
#####

def sub_report_seen(connection, event, nick, userhost, replypath, what, admin_unlocked):
    whats = what.split(' ')
    if whats[0] == 'seen' and len(whats) >= 2:
        query = whats[1]
        if query != 'debug':
            try:
                seendata = config.get('seen', query).split('|:|')
                converted = datetime.strptime(seendata[1], "%Y-%m-%dT%H:%M:%S.%f").replace(tzinfo=tzutc())
                connection.privmsg(replypath, 'last saw ' + query + ' at ' + converted.astimezone(tzlocal()).strftime("%Y/%m/%d %H:%M:%S %Z") + ' saying \'' + seendata[2] + '\'')
            except NoOptionError: pass

#####
# sub_save_config
# save the config file
#####

def sub_save_config(connection, event, nick, userhost, replypath, what, admin_unlocked):
    whats = what.split(' ')
    if whats[0] == 'save' and admin_unlocked:
        with open('dr.botzo.cfg', 'w') as cfg:
            config.write(cfg)
        connection.privmsg(replypath, 'saved config file')

#####
# sub_change_nick
# change the bot's nickname
#####

def sub_change_nick(connection, event, nick, userhost, replypath, what, admin_unlocked):
    whats = what.split(' ')
    if whats[0] == 'nick' and admin_unlocked and len(whats) >= 2:
        newnick = whats[1]
        connection.nick(newnick)
        config.set('IRC', 'nick', newnick)
        connection.privmsg(replypath, 'changed nickname')

#####
# sub_countdown
# add a countdown item to the bot
#####

def sub_countdown(connection, event, nick, userhost, replypath, what, admin_unlocked):
    whats = what.split(' ')
    if whats[0] == 'countdown' and len(whats) >= 2:
        if whats[1] == 'add' and len(whats) >= 4:
            item = whats[2]
            target = parse(' '.join(whats[3:]), default=datetime.now().replace(tzinfo=tzlocal()))
            if not config.has_section('countdown'):
                config.add_section('countdown')

            config.set('countdown', item, target.astimezone(tzutc()).isoformat())
            connection.privmsg(replypath, 'added countdown item ' + whats[2])
        elif whats[1] == 'remove':
            try:
                if config.remove_option('countdown', whats[2]):
                    connection.privmsg(replypath, 'removed countdown item ' + whats[2])
            except NoSectionError: pass
        elif whats[1] == 'list':
            try:
                cdlist = config.options('countdown')
                cdlist.remove('debug')
                cdlist.sort()
                liststr = ', '.join(cdlist)
                connection.privmsg(replypath, liststr)
            except NoSectionError: pass
        else:
            try:
                timestr = config.get('countdown', whats[1])
                time = parse(timestr)
                rdelta = relativedelta(time, datetime.now().replace(tzinfo=tzlocal()))
                relstr = whats[1] + ' will occur in '
                if rdelta.years != 0:
                    relstr += str(rdelta.years) + ' years '
                if rdelta.months != 0:
                    relstr += str(rdelta.months) + ' months '
                if rdelta.days != 0:
                    relstr += str(rdelta.days) + ' days '
                if rdelta.hours != 0:
                    relstr += str(rdelta.hours) + ' hours '
                if rdelta.minutes != 0:
                    relstr += str(rdelta.minutes) + ' minutes '
                if rdelta.seconds != 0:
                    relstr += str(rdelta.seconds) + ' seconds'
                #relstr += ' (' + timestr + ')'
                connection.privmsg(replypath, relstr)
            except NoOptionError: pass

#####
# sub_dice
# roll dice, primarily for roleplaying games
#####

def sub_dice(connection, event, nick, userhost, replypath, what, admin_unlocked):
    overallroll = what

    rolls = re.split(';\s*', overallroll)
    for roll in rolls:
        pattern = '^(?:(\d+)#)?(?:(\d+)/)?(\d+)?d(\d+)(?:(\+|\-)(\d+))?(?:\s+(.*))?'
        regex = re.compile(pattern)
        matches = regex.search(roll)

        if matches is not None:
            # set variables including defaults for unspecified stuff
            faces = int(matches.group(4))
            comment = matches.group(7)

            if matches.group(1) is None:
                times = 1
            else:
                times = int(matches.group(1))

            if matches.group(3) is None:
                dice = 1
            else:
                dice = int(matches.group(3))

            if matches.group(2) is None:
                top = dice
            else:
                top = int(matches.group(2))

            if matches.group(5) is None or matches.group(6) is None:
                modifier = 0
            else:
                if str(matches.group(5)) == '-':
                    modifier = -1 * int(matches.group(6))
                else:
                    modifier = int(matches.group(6))

            result = ''
            if comment is not None:
                result += comment + ': '

            for t in range(times):
                ressubstr = ""
                rolls = []
                for d in range(dice):
                    rolls.append(str(random.randint(1, faces)))
                rolls.sort()
                rolls.reverse()
                ressubstr = ','.join(rolls[0:top])
                sum = 0
                for r in rolls[0:top]:
                    sum += int(r)
                sumplus = sum + modifier
                result += str(sumplus) + ' [' + ressubstr
                if modifier != 0:
                    if modifier > 0:
                        result += ' + ' + str(modifier)
                    else:
                        result += ' - ' + str(-1 * modifier)
                result += ']'

                if t != times-1:
                    result += ', '

            connection.privmsg(replypath, result)

#####
# on_connect
# handler for when the bot has connected to IRC
#####

def on_connect(connection, event):
    # user modes
    try:
        usermode = config.get('IRC', 'usermode')
        connection.mode(botnick, usermode)
    except NoOptionError: pass

    # join the specified channels
    try:
        autojoins = config.get('channels', 'autojoin').split(',')
        for channel in autojoins:
            if irclib.is_channel(channel):
                connection.join(channel)
    except NoOptionError: pass

#####
# on_privmsg
# private messages to the bot
#####

def on_privmsg(connection, event):
    nick = irclib.nm_to_n(event.source())
    userhost = irclib.nm_to_uh(event.source())
    replypath = nick
    what = event.arguments()[0]

    admin_unlocked = False

    try:
        if userhost == config.get('admin', 'userhost'):
            admin_unlocked = True
    except NoOptionError: pass

    # admin commands
    sub_join_channel(connection, event, nick, userhost, replypath, what, admin_unlocked)
    sub_part_channel(connection, event, nick, userhost, replypath, what, admin_unlocked)
    sub_quit_channel(connection, event, nick, userhost, replypath, what, admin_unlocked)
    sub_autojoin_manipulate(connection, event, nick, userhost, replypath, what, admin_unlocked)
    sub_save_config(connection, event, nick, userhost, replypath, what, admin_unlocked)
    sub_change_nick(connection, event, nick, userhost, replypath, what, admin_unlocked)

    # standard commands
    sub_report_seen(connection, event, nick, userhost, replypath, what, admin_unlocked)
    sub_countdown(connection, event, nick, userhost, replypath, what, admin_unlocked)
    sub_dice(connection, event, nick, userhost, replypath, what, admin_unlocked)

#####
# on_pubmsg
# public messages in a channel where the bot is
#####

def on_pubmsg(connection, event):
    nick = irclib.nm_to_n(event.source())
    userhost = irclib.nm_to_uh(event.source())
    replypath = event.target()
    what = event.arguments()[0]

    admin_unlocked = False

    try:
        if userhost == config.get('admin', 'userhost'):
            admin_unlocked = True
    except NoOptionError: pass

    sub_add_to_seen(connection, event, nick, userhost, what)

    # only do commands if the bot has been addressed directly
    addressed_pattern = '^' + connection.get_nickname() + '[:,]?\s+'
    addressed_re = re.compile(addressed_pattern)

    if not addressed_re.match(what):
        return
    else:
        what = addressed_re.sub('', what)

    # admin commands
    sub_join_channel(connection, event, nick, userhost, replypath, what, admin_unlocked)
    sub_part_channel(connection, event, nick, userhost, replypath, what, admin_unlocked)
    sub_quit_channel(connection, event, nick, userhost, replypath, what, admin_unlocked)
    sub_autojoin_manipulate(connection, event, nick, userhost, replypath, what, admin_unlocked)
    sub_save_config(connection, event, nick, userhost, replypath, what, admin_unlocked)
    sub_change_nick(connection, event, nick, userhost, replypath, what, admin_unlocked)

    # standard commands
    sub_report_seen(connection, event, nick, userhost, replypath, what, admin_unlocked)
    sub_countdown(connection, event, nick, userhost, replypath, what, admin_unlocked)
    sub_dice(connection, event, nick, userhost, replypath, what, admin_unlocked)

#####
# init
#####

# read config file

config = ConfigParser({'debug': 'false'})
config.read([os.path.expanduser('~/.dr.botzo.cfg'), 'dr.botzo.cfg'])

# load necessary options
try:
    # load connection info
    botserver   = config.get('IRC', 'server')
    botport     = config.getint('IRC', 'port')
    botnick     = config.get('IRC', 'nick')
    botircname  = config.get('IRC', 'name')
except NoSectionError as e:
    sys.exit("Aborted due to error with necessary configuration: " + str(e))
except NoOptionError as e:
    sys.exit("Aborted due to error with necessary configuration: " + str(e))

# load additional options
irclib.DEBUG = config.getboolean('IRC', 'debug')

# start up the IRC bot

# create IRC and server objects and connect
irc = irclib.IRC()
server = irc.server().connect(botserver, botport, botnick, botircname)

# install handlers
server.add_global_handler("welcome", on_connect)
server.add_global_handler('privmsg', on_privmsg)
server.add_global_handler('pubmsg', on_pubmsg)

# loop forever
irc.process_forever()

# vi:tabstop=4:expandtab:autoindent