From df1ca3ad2df0b80771c93ed9ce2a207e2c55099d Mon Sep 17 00:00:00 2001
From: Fanir <projects@mail.fanir.de>
Date: Mon, 3 Mar 2014 18:06:42 +0100
Subject: [PATCH] cleanups, commenting, moar code and other nice things.. (i'll
 start writing better commitmessages soon, i promise!)

---
 main.py | 322 ++++++++++++++++++++++++++++++++++++++------------------
 1 file changed, 222 insertions(+), 100 deletions(-)

diff --git a/main.py b/main.py
index 788e8b0..5c8fdbb 100755
--- a/main.py
+++ b/main.py
@@ -27,6 +27,7 @@
 
 
 import sys, string, socket, re, signal
+from random import choice
 from time import sleep
 from select import poll, POLLIN, POLLPRI, POLLOUT,\
                         POLLERR, POLLHUP, POLLNVAL
@@ -39,30 +40,25 @@ class log():
     DEBUG, INFO, WARNING, ERROR, FATAL, SILENT =\
         0,    1,       2,     3,     4,      5
     
+    levels = {
+        0: "[DEBUG]   ", 1: "[INFO]    ", 2: "[WARNING] ",
+        3: "[ERROR]   ", 4: "[FATAL]   ", 5: "[SILENT]  "
+    }
+    
     def show(level, msg):
         if level in range(log.DEBUG, log.SILENT):
-            if LOGLEVEL <= level: print(msg)
+            if LOGLEVEL <= level: print("".join((log.levels[level], 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="", serverpasswd="", parser_wait_time=0.1):
+    def __init__(self, nicknames, server, port=6667,
+                 ident="pircbot", realname="pircbot", serverpasswd="",
+                 encodings=("utf-8", "latin-1"), admins="",
+                 query_type="", command_prefix=".",
+                 parser_wait_time=0.1):
         self.server = server
         self.port = port
         self.serverpasswd = serverpasswd
@@ -74,39 +70,89 @@ class pircbot():
         
         self.admins = admins
         self.cmdprefix = command_prefix
+        self.query_type = query_type.upper()
         
         self.parser_wait_time=parser_wait_time
         
         self.user = {}
         
         self.socket = None
-        self.recvloop = None
         self.recvbuffer = bytearray(1024)
         
-        self.ready = False
+        self.recvloop = None
+        self.parseloop = None
         
         self.die_event = Event()
+        self.ready = False
+        
+        self.mode_reply_to = None
     
-    def recv_loop(self):
+    def connect(self):
+        # connect
+        self.socket = socket.socket()
+        log.show(log.DEBUG, "--- SOCKET OPENING ---")
+        try:
+            self.socket.connect((self.server, self.port))
+        except socket.error as e:
+            log.show(log.FATAL, "Fehler: %s" % e)
+            return
+        
+        # start getting data
+        self.recvloop = Thread(target=self.recv, name="recvloop")
+        self.recvloop.start()
+        
+        # optionally send a server password
+        if self.serverpasswd != "": self.send("PASS %s" % self.serverpasswd)
+        # get a nick
+        self.send("NICK %s" % self.nicknames.pop(0))
+        # set user data
+        self.send("USER %s 0 * :%s" % (self.ident, self.realname))
+    
+    # implements irc command QUIT and (more or less) clean exiting
+    def disconnect(self, reason="", send_quit=True):
+        if send_quit:
+            try:
+                self.send("QUIT :%s" % reason)
+                sleep(1)
+            except: pass
+        self.die_event.set()
+        ctr = 0
+        while self.recvloop.is_alive() and self.parseloop.is_alive() and ctr < 15:
+            ctr += 1
+            sleep(1)
+        log.show(log.DEBUG, "--- SOCKET CLOSING ---")
+        try:
+            self.socket.shutdown(socket.SHUT_RDWR)
+            self.socket.close()
+        except: pass
+    
+    
+    ### threaded functions ###
+    
+    # loop for recieving data from irc
+    def recv(self):
         """
         Loop for reciving data
         """
-        parser = Thread(target=self.parse_loop, name="parser")
+        log.show(log.DEBUG, "--- RECVLOOP STARTING ---")
+        self.parseloop = Thread(target=self.parser, name="parser")
         p = poll()
         p.register(self.socket.fileno(), POLLIN)
         while not self.die_event.is_set():
             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_loop, name="parser")
-                    parser.start()
-        print("--- RECVLOOP EXITING ---")
+                if not self.parseloop.is_alive():
+                    self.parseloop = Thread(target=self.parser, name="parser")
+                    self.parseloop.start()
+        log.show(log.DEBUG, "--- RECVLOOP EXITING ---")
     
-    def parse_loop(self):
+    # loop for parsing incoming data
+    def parser(self):
         """
         Loop for parsing incoming data
         """
+        log.show(log.DEBUG, "--- PARSELOOP STARTING ---")
         while not self.die_event.is_set():# and self.recvbuffer.endswith(b"\r\n"):# != b"":
             if self.recvbuffer.endswith(b"\r\n"):
                 # get and decode line from buffer
@@ -122,6 +168,7 @@ class pircbot():
                     prefix = head.pop(0)[1:]
                     if "@" in prefix:
                         origin["nick"], origin["ident"], origin["host"] = re.match(r"(\S+)!(\S+)@(\S+)", prefix).groups()
+                        origin["mask"] = prefix
                     else:
                         origin["server"] = prefix
                 else:
@@ -130,20 +177,32 @@ class pircbot():
                 command = head.pop(0)
                 # parse params
                 params = head
-                params.append(larg)
+                params.append(larg.strip())
+                print(params)
                 
                 log.show(log.DEBUG, " > %s" % rawline)
                 
                 # PING
                 if command == "PING":
                     self.send("PONG %s" % params[0])
-                # PRIVMSG
-                elif command == "PRIVMSG":
+                # PRIVMSG and NOTICE
+                elif command == "PRIVMSG" or command == "NOTICE":
                     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:])
+                        args = []
+                        for v in params[1].lstrip(self.cmdprefix).split(" "):
+                            if v!="": args.append(v)
+                        sndline = self.on_command(command, prefix, origin, params[0], args[0].lower(), args[1:])
                         if sndline != None:
                             self.send(sndline)
+                # 221 (RPL_UMODEIS)
+                elif command == "221" and self.mode_reply_to is not None:
+                        self.send("PRIVMSG %s :Modes are %s" % (self.mode_reply_to, params[1]))
+                # INVITE
+                elif command == "INVITE":
+                    if self.is_admin(origin["mask"]):
+                        self.join(params[1])
+                    else:
+                        self.send("NOTICE %s :You can not force me to do that!" % origin["nick"])
                 # 001 (RPL_WELCOME)
                 elif command == "001":
                     self.user["mask"] = re.search(r" (\S+!\S+@\S+)$", rawline.split(" :", 1)[1]).groups()[0]
@@ -155,69 +214,110 @@ class pircbot():
                 # KILL
                 elif command == "KILL":
                     log.show(log.WARNING, "Got killed by %s: %s" % (params[0], params[1:]))
-                    self.quit(send=False)
+                    self.disconnect(send_quit=False)
             else:
                 sleep(self.parser_wait_time)
-        print("--- PARSELOOP EXITING ---")
+        log.show(log.DEBUG, "--- PARSELOOP EXITING ---")
     
-    def send(self, data):
-        log.show(log.DEBUG, "<  %s" % data)
-        try: self.socket.send(self.encode("".join((data, "\r\n"))))
-        except BrokenPipeError as e:
-            log.show(log.FATAL, e)
-            self.quit(send=False)
     
-    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
-        
-        # start getting data
-        self.recvloop = Thread(target=self.recv_loop, name="recvloop")
-        self.recvloop.start()
-        
-        # optionally send a server password
-        if self.serverpasswd != "": self.send("PASS %s" % self.serverpasswd)
-        # get a nick
-        self.send("NICK %s" % self.nicknames.pop(0))
-        # set user data
-        self.send("USER %s 0 * :%s" % (self.ident, self.realname))
+    ### helper functions ###
     
+    # encodes data for irc
+    def encode(self, textstring):
+        for codec in self.encodings:
+            try: return textstring.encode(codec)
+            except UnicodeDecodeError: continue
+        return textstring.encode(self.encodings[0], 'ignore')
+    
+    # decodes data from irc
+    def decode(self, textstring):
+        for codec in self.encodings:
+            try: return textstring.decode(codec)
+            except UnicodeDecodeError: continue
+        return textstring.decode(self.encodings[0], 'ignore')
+    
+    
+    # checks if a given user is listed in admins
     def is_admin(self, user):
         for admin in self.admins:
             if re.search(admin, user) != None:
                 return True
         return False
     
-    def quit(self, reason="", send=True):
-        if send:
-            try:
-                self.send("QUIT :%s" % reason)
-            except: pass
-        self.die_event.set()
-        print("--- SOCKET CLOSING ---")
-        try:
-            self.socket.shutdown(socket.SHUT_RDWR)
-            self.socket.close()
-        except: pass
+    # checks if s is a valid channel name
+    def is_channel(self, s):
+        return (True if s[0] in ('&', '#', '+', '!') else False)
+    
+    
+    def exec_on_ready(self, command, retry_times=-1, interval=1):
+        if command.startswith("self."):
+            if retry_times == -1:
+                while not self.ready: sleep(interval)
+            else:
+                cnt = 0
+                while not self.ready and cnt<retry_times: sleep(interval)
+                if not self.ready:
+                    log.show(log.WARNING, "Connection did not get ready in time (%sx %s seconds), \"%s\" not executed" % (retry_times, interval, command))
+                    return True
+            exec("ret = %s" % command)
+            return False
+        else:
+            log.show(log.FATAL, "exec_on_ready() called with an invalid command: %s" % command)
+            return True
+    
+    def exec_retry(self, command, times=5, interval=1, wait_for_ready=True):
+        if command.startswith("self."):
+            if wait_for_ready:
+                while not self.ready: sleep(interval)
+            cnt = 0
+            exec("while %s and cnt<times: sleep(interval)" % command)
+        else: log.show(log.FATAL, "exec_retry() called with an invalid command.")
+    
+    
+    def send(self, data):
+        log.show(log.DEBUG, "<  %s" % data)
+        try: self.socket.send(self.encode("".join((data, "\r\n"))))
+        except BrokenPipeError as e:
+            log.show(log.FATAL, e)
+            self.disconnect(send_quit=False)
+    
+    # replies to channel by PRIVMSG
+    def chanmsg(self, channel, msg):
+        self.send("".join(("PRIVMSG ", channel, " :", msg)))
+    
+    # replies to user by NOTICE or PRIVMSG
+    def query(self, nick, msg, in_query_type):
+        self.send("".join((self.query_type if self.query_type!="" else in_query_type, " ", nick, " :", msg)))
+    
+    # decides whether to print to user or to channel based on the string origin
+    def reply(self, origin, source, msg, in_query_type):
+        if self.is_channel(source):
+            self.chanmsg(source, msg)
+        else:
+            self.query(origin["nick"], msg, in_query_type)
+    
+    
+    ### irc command wrapper ###
+    
+    def nick(self, nick):
+        self.exec_on_ready("".join(('self.send("JOIN %s" % "', channel, '")')))
     
     def join(self, channel):
-        while not self.ready: sleep(1)
-        self.send("JOIN %s" % channel)
+        self.exec_on_ready("".join(('self.send("JOIN %s" % "', channel, '")')))
     
     def part(self, channel):
-        while not self.ready: sleep(1)
-        self.send("PART %s" % channel)
+        self.exec_on_ready("".join(('self.send("PART %s" % "', channel, '")')))
     
     def set_mode(self, modes):
-        while not self.ready: sleep(1)
-        self.send("MODE %s :%s" % (self.user["nick"], modes))
+        self.exec_on_ready("".join(('self.send("MODE %s :%s" % (self.user["nick"], "', modes, '"))')))
+        
+    def get_modes(self):
+        self.exec_on_ready("".join(('self.send("MODE %s" % self.user["nick"])')))
     
-    def on_command(self, prefix, origin, source, command, params):
+    ### handler ###
+    
+    # bot-command handler
+    def on_command(self, in_query_type, prefix, origin, source, command, params):
         """
         Executed when getting a PRIVMSG starting with self.cmdprefix
         Prefix contains the optional prefix of the raw line.
@@ -231,22 +331,36 @@ class pircbot():
         # hello
         if command == "hello":
             greeting = "".join(("Hi " + origin["nick"] +"!"))
-            if source[0] in {"#", "+", "!", "&"}:
-                return "PRIVMSG %s :%s" % (source, greeting)
-            else:
-                return "PRIVMSG %s :%s" % (origin["nick"], greeting)
+            self.reply(origin, source, greeting, in_query_type)
         # join <channel>
         elif command == "join" and len(params)>0:
-            if self.is_admin(origin["nick"]): self.join(params[0])
-            else: return "PRIVMSG %s :You can't do that!" % origin["nick"]
+            if self.is_admin(origin["mask"]): self.join(params[0])
+            else: return "PRIVMSG %s :You cannot do that!" % origin["nick"]
         # part <channel>
         elif command == "part" and len(params)>0:
-            if self.is_admin(origin["nick"]): self.part(params[0])
-            else: return "PRIVMSG %s :You can't do that!" % origin["nick"]
+            if self.is_admin(origin["mask"]): self.part(params[0])
+            else: return "PRIVMSG %s :You cannot do that!" % origin["nick"]
+        # mode ±<modes>
+        elif command == "mode":
+            if self.is_admin(origin["mask"]):
+                if len(params) > 0:
+                    self.set_mode(params[0])
+                else:
+                    self.mode_reply_to = origin["nick"]
+                    self.get_modes()
+            else: return "PRIVMSG %s :You cannot do that!" % origin["nick"]
         # die [<quitmsg>]
         elif command == "die":
-            if self.is_admin(origin["nick"]):
-                self.quit("".join(params) if len(params)>0 else "".join((origin["nick"], " shot me, dying now... Bye...")))
+            if self.is_admin(origin["mask"]):
+                self.disconnect("".join(params) if len(params)>0 else "".join((origin["nick"], " shot me, dying now... Bye...")))
+            else: return "PRIVMSG %s :Go die yourself!" % origin["nick"]
+        else:
+            replies = [
+                ("PRIVMSG %s :What? \"%s\" is not a command!", 15),
+                ("PRIVMSG %s :%s? What's that?", 3),
+                ("PRIVMSG %s :Sorry, I don't know how to %s...", 1)
+            ]
+            return choice([val for val, cnt in replies for i in range(cnt)]) % (origin["nick"], command)
 
 
 
@@ -262,10 +376,12 @@ def main():
     global bot
     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, serverpasswd=SERVERPASSWD, parser_wait_time=PARSER_WAIT_TIME)
         try:
+            bot = pircbot(nicknames=NICKNAMES, ident=IDENT, realname=REALNAME,
+                          server=SERVER, port=PORT, serverpasswd=SERVERPASSWD,
+                          encodings=ENCODINGS, admins=ADMINS,
+                          query_type=QUERY_TYPE, command_prefix=COMMAND_PREFIX,
+                          parser_wait_time=PARSER_WAIT_TIME)
             bot.connect()
             # wait for the bot to become ready
             while bot.ready and not bot.die_event.is_set() == False:
@@ -282,10 +398,10 @@ def main():
                 sleep(1)
         except KeyboardInterrupt:
             log.show(log.INFO, "Got Ctrl-C, dying now...")
-            bot.quit()
+            bot.disconnect("Ouch! Got shot by Ctrl-C, dying now... See you!")
     
     elif args.action == "stop": print("nope!")
-    print("--- MAIN EXITING ---")
+    log.show(log.DEBUG, "--- MAIN EXITING ---")
     return 0
 
 
@@ -295,12 +411,6 @@ 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"
@@ -319,17 +429,18 @@ if __name__ == '__main__':
     
     
     # The Server name to connect to. Duh!
-    #SERVER = "fanir.de"
-    SERVER = "localhost"
+    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
     
+    # Serverpassword. Empty in most cases.
     SERVERPASSWD = ""
     
     # A comma-seperated list of channels to join, enclosed by braces.
-    CHANNELS = ["#bots", "#main"]
+    CHANNELS = ["#bots"]
     
     # 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.
@@ -340,10 +451,21 @@ if __name__ == '__main__':
     
     # 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.*"]
+    ADMINS = ["Fanir\!.*"]
     
+    # The prefix for commands for the bot.
     COMMAND_PREFIX = "."
     
+    # Which way shouldbe used to speak to users?
+    # "" means, the type of the incomming message should be used.
+    # One of: "NOTICE", "PRIVMSG", ""
+    QUERY_TYPE = ""
+    
+    
+    # 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
     
     # The time the parser for incoming data should wait between each attempt to read new data in seconds.
     # High values will certainly make the bot reply slowly while very low values increads cpu load and therefore will perform badly on slow machines.