chalkbot/main.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
#  main.py
#
#  Copyright 2014 Fanir <projects@mail.fanir.de>
#
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software
#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
#  MA 02110-1301, USA.
#
#


# SETTINGS CAN BE FOUND AT THE BOTTOM OF THE FILE


import sys, time, string, socket, re, signal
from select import poll, POLLIN, POLLPRI, POLLOUT,\
                        POLLERR, POLLHUP, POLLNVAL
from threading import Thread#, Event

bot = None

class log():
    DEBUG, INFO, WARNING, ERROR, FATAL, SILENT =\
        0,    1,       2,     3,     4,      5
    
    def show(level, msg):
        if level in range(log.DEBUG, log.SILENT):
            if LOGLEVEL <= level: print(msg)
        else:
            raise ValueError("That's not a loglevel!")


class pircbot():
    def encode(self, textstring):
        for codec in self.encodings:
            try: return textstring.encode(codec)
            except UnicodeDecodeError: continue
        return textstring.encode(self.encodings[0], 'ignore')
    
    def decode(self, textstring):
        for codec in self.encodings:
            try: return textstring.decode(codec)
            except UnicodeDecodeError: continue
        return textstring.decode(self.encodings[0], 'ignore')
    
    def __init__(self, nicknames, server, port=6667, ident="pircbot",
                realname="pircbot", encodings=("utf-8", "latin-1"),
                command_prefix=".", admins=""):
        self.nicknames = nicknames
        self.server = server
        self.port = port
        self.ident = ident
        self.realname = realname
        self.encodings = encodings
        self.admins = admins
        self.cmdprefix = command_prefix
        self.socket = None
        self.recvloop = None
        self.recvbuffer = bytearray(1024)
        self.run = False
        self.ready = False
        self.is_alive = False
        self.need_nick = True
    
    def recv(self):
        """
        Loop for reciving data
        """
        parser = Thread(target=self.parse, name="parser")
        p = poll()
        p.register(self.socket.fileno(), POLLIN)
        while self.run:
            ap = p.poll(1000)
            if (self.socket.fileno(), POLLIN) in ap:
                self.recvbuffer.extend(self.socket.recv(1024))
                if not parser.is_alive():
                    parser = Thread(target=self.parse, name="parser")
                    parser.start()
        if parser.is_alive(): parser.join()
    
    def connect(self):
        # connect
        self.socket = socket.socket()
        try:
            self.socket.connect((self.server, self.port))
        except socket.error as e:
            log.show(log.FATAL, "Fehler: %s" % e)
            return
        
        self.run = True
        
        # start getting data
        self.recvloop = Thread(target=self.recv, name="recvloop")
        self.recvloop.start()
        
        # get a nick
        self.socket.send(self.encode("NICK %s\r\n" % self.nicknames.pop(0)))
        
        self.socket.send(self.encode("USER %s %s bla :%s\r\n" % (self.ident, self.server, self.realname)))
        
        self.is_alive = True
    
    def parse(self):
        """
        Loop for parsing incoming data
        """
        #line = ""
        origin = {}
        while self.run and self.recvbuffer != b"":
                # get and decode line from buffer
                rawline, _, self.recvbuffer = self.recvbuffer.partition(b"\r\n")
                rawline = self.decode(rawline)
                # prepare line
                line = rawline.split(" :", 1)
                head = line[0].lstrip("\x00").split(" ")
                larg = line[1] if len(line)>1 else ""
                # parse prefix
                if head[0].startswith(":"):
                    prefix = head.pop(0)[1:]
                    origin.clear()
                    if "@" in prefix:
                        origin["nick"], origin["ident"], origin["host"] = re.match(r"(\S+)!(\S+)@(\S+)", prefix).groups()
                    else:
                        origin["server"] = prefix
                else:
                    prefix = ""
                # parse command
                command = head.pop(0)
                # parse params
                params = head
                params.append(larg)
                
                log.show(log.DEBUG, ">> %s" % rawline)
                
                if command == "PING":
                    self.socket.send(self.encode("PONG %s\r\n" % params[0]))
                elif command == "PRIVMSG":
                    if params[1].startswith(self.cmdprefix):
                        args = params[1].lstrip(self.cmdprefix).split(" ")
                        sndline = self.on_command(prefix, origin, params[0], args[0], args[1:])
                        if sndline != None:
                            log.show(log.DEBUG, "<< %s" % sndline)
                            self.socket.send(self.encode(sndline))
                elif command == "001":
                    self.ready = True
                elif command == "433":
                    self.socket.send(self.encode("NICK %s\r\n" % self.nicknames.pop(0)))
                elif command == "KILL":
                    print("Got killed: %s", rawline)
                    self.die()
    
    def is_admin(self, user):
        for admin in self.admins:
            if re.search(admin, user) != None:
                return True
        return False
    
    def die(self):
        self.run = False
        self.recvloop.join()
        self.is_alive = False
        sys.exit()
    
    def quit(self, reason=""):
        self.socket.send(self.encode("QUIT :%s\n" % reason))
        self.run = False
        self.recvloop.join()
        self.is_alive = False
    
    def join(self, channel):
        print(channel)
        self.socket.send(self.encode("JOIN %s\n" % channel))
    
    def part(self, channel):
        self.socket.send(self.encode("PART %s\n" % channel))
    
    def on_command(self, prefix, origin, source, command, params):
        """
        Executed when getting a PRIVMSG starting with self.cmdprefix
        Prefix contains the optional prefix of the raw line.
        Origin is a map holding the parsed prefix, containing
            "nick", "ident" and "host" for users and "server" for servers.
        Source is the first parameter after the irc-command, specifying
            the channel or user, from where the line comes.
        Command is the command for the bot.
        Params contains a list of originally space separated parameters.
        """
        print(params)
        if command.startswith("hello"):
            greeting = "".join(("Hi " + origin["nick"] +"!"))
            if source[0] in {"#", "+", "!", "&"}:
                return "PRIVMSG %s :%s\r\n" % (source, greeting)
            else:
                return "PRIVMSG %s :%s\r\n" % (origin["nick"], greeting)
        elif command.startswith("join") and len(params)>0:
            if self.is_admin(origin["nick"]): self.join(params[0])
            else: return "PRIVMSG %s :You can't do that!\r\n" % origin["nick"]
        elif command.startswith("part") and len(params)>0:
            if self.is_admin(origin["nick"]): self.part(params[0])
            else: return "PRIVMSG %s :You can't do that!\r\n" % origin["nick"]


def parseargs():
    import argparse
    p = argparse.ArgumentParser(
        description = "i think my desc is missing")
    p.add_argument("action", default="help", choices = ["start", "stop"], help="What to do?")
    #p.add_argument("--daemon", "-d", type = bool, choices = [1, 0], default=1, help="Daemonize, Default: 1")
    return p.parse_args()

def KeyboardInterruptHandler(signum, frame):
    global bot
    print("Got Ctrl-C, dying now...")
    bot.quit("")
    sys.exit()

def main():
    global bot
    signal.signal(signal.SIGINT, KeyboardInterruptHandler)
    
    args = parseargs()
    if args.action == "start":
        bot = pircbot(nicknames=NICKNAMES, ident=IDENT, realname=REALNAME,
            server=SERVER, port=PORT, encodings=ENCODINGS, command_prefix=COMMAND_PREFIX, admins=ADMINS)
        try:
            bot.connect()
            while bot.ready == False: time.sleep(1)
            for channel in CHANNELS:
                bot.join(channel)
            while 1:
                if bot.is_alive == False: print("X")
                time.sleep(10)
        except SystemExit:
            print("exiting...")
    elif args.action == "stop": print("nope!")
    return 0


if __name__ == '__main__':
    
    ################
    ### SETTINGS ###
    ################
    
    # With how much information do you want to be annoyed?
    # DEBUG spams most, FATAL least. WARNING should be a good tradeoff.
    # One of: log.DEBUG, log.INFO, log.WARNING, log.ERROR, log.FATAL
    LOGLEVEL = log.DEBUG
    
    
    # Also known as username. Some IRC-internal.
    # By default, the nickname will be used as ident.
    IDENT = "chalkbot"
    
    # The list of nicknames to try. If the first one is not aviable, it will
    # try the second, and so on...
    # You should specfy at least two nicks.
    NICKNAMES = ["chalkbot", "chalkbot_", "chalkbot__"]
    
    REALNAME = "A ChalkBot Instance"
    
    # Command for registering with nickserv, without leading slash.
    NICKSERVCMD = ""
    
    MODES = "+B"
    
    
    # The Server name to connect to. Duh!
    #SERVER = "fanir.de"
    SERVER = "localhost"
    
    # "Default" is 6667. An often used port for SSL would be 6697, if SSL would be
    # supported. Maybe in a future, far far away...
    PORT = 6667
    
    # A comma-seperated list of channels to join, enclosed by braces.
    CHANNELS = ["#bots", "#main"]
    
    # The encodings to try when getting messages from IRC. Will be tried in the given order.
    # The first one is used to encode data when sending stuff.
    # The list given shoud do just fine in most networks, I assume.
    # Also comma seperated and enclosed by braces.
    ENCODINGS = ['utf-8', 'latin-1', 'iso-8859-1', 'cp1252']
    
    
    # List of users (hostmasks, will be parsed as regex) who can do important stuff,
    # like joining and parting channels and shutting down the bot.
    ADMINS = ["Fanir.*"]
    
    COMMAND_PREFIX = "."
    
    #######################
    ### END OF SETTINGS ###
    #######################
    
    
    main()