initial split of bot with a command queue

bot.py

@@ -1,327 +1,38 @@
-#!/usr/bin/env python3
-
-import random
-import threading
-import socket
-import subprocess
+import queue
+import shlex
 import sys
-import time
-from unzipper import unar
-
-
-class IRCClient():
-    HOST = "irc.irchighway.net"
-    PORT = 6667
-    NICK = ""
-    IDENT = ""
-    REALNAME = ""
-    IGNORE = [
-        "NOTICE",
-        "PART",
-        "QUIT",
-        "332",
-        "333",
-        "372",
-        "353",
-        "366",
-        "251",
-        "252",
-        "254",
-        "255",
-        "265",
-        "266",
-        "396"]
-    CHANNEL = "#ebooks"
-    STATUS = ""
-    PROGRESS = 0
-    TYPE = "SEARCH"
-    OUTPUT = ""
-    buffer = []
-    readbuffer = b''
-    NOT_EXITED = True
-    EXTRA_OUTPUT = ""
-    PATH = ""
-    joined = False
-    connected = False
-
-    def __init__(self, book, extension, t, cb, logging=False, path="/tmp/"):
-
-        self.TYPE = t
-        self.EXTENSION = extension
-        self.LOOKING_FOR = book
-        self.LOGGING = logging
-        self.PATH = path
-
-        name = "bookbot" + random_hash()
-        self.NICK = name
-        self.IDENT = name
-        self.REALNAME = name
-        self.callback = cb
-
-        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        self.socket.connect((self.HOST, self.PORT))
-        self.handle_connect()
-
-        nickstr = "NICK %s" % self.NICK
-        userstr = "USER %s %s bla :%s" % (self.IDENT, self.HOST, self.REALNAME)
-
-        self.t = threading.Thread(target=self.handle_read)
-        self.t.start()
-        self.send_queue(nickstr)
-        self.send_queue(userstr)
-
-    def handle_connect(self):
-        self.log("connected")
-        self.set_status("CONNECTED")
-        self.connected = True
-
-    def handle_close(self):
-        self.NOT_EXITED = False
-        self.log("closed")
-        self.set_status("DISCONNECTED")
-        self.connected = False
-        self.socket.close()
-
-    def do_timeout(self):
-        if not self.NOT_EXITED:
-            return
-        self.NOT_EXITED = False
-        self.log("Abort mission")
-        self.set_status("TIMED OUT")
-        self.connected = False
-        self.socket.close()
-
-    def handle_read(self):
-        while self.NOT_EXITED:
-            try:
-                data = self.socket.recv(1024)
-            except socket.error:
-                break
-            if len(data) < 2:
-                continue
-            if not (data[-1] == 10 and data[-2] == 13):
-                self.readbuffer += data
-                continue
-
-            data = self.readbuffer + data
-            print("DATA", data)
-
-            # purge mojibake
-            data = (data.replace(b'\x95', b'')
-            .replace(b'0xc2', b'')
-            .decode('utf-8', 'ignore'))
-            self.readbuffer = b''
-
-            lines = data.splitlines()
-
-            for line in lines:
-                words = line.split(' ')
-                if len(words) < 2:
-                    continue
-                msg_from = words[0]
-                comm = words[1]
-
-                if comm in self.IGNORE:
-                    continue
-                if comm == "JOIN":
-                    if self.NICK not in msg_from:  # msg "NICK joined the channel" not about me
-                        continue
-                    self.set_status("JOINED")
-                    self.log("Joined channel %s" % self.CHANNEL)
-                    self.log("Waiting 30s")
-                    time.sleep(31)
-                    self.run_query()
-                    self.joined = True
-                    continue
-
-                if comm == "PRIVMSG":
-                    # private message not addressed to me
-                    if words[2] != self.NICK:
-                        continue
-                    self.parse_msg(line)
-
-                if comm == "PING" or msg_from == "PING":  # respond ping to avoid getting kicked
-                    self.pong(line)
-                    continue
-                if comm == "376":  # END MOTD
-                    # MOTD complete, lets join the channel
-                    self.join_channel(self.CHANNEL)
-                    continue
-
-                if not self.joined:  # waiting
-                    continue
-
-    def run_query(self):
-        self.set_status("WAITING")
-        if self.TYPE == "SEARCH":
-            self.search(self.LOOKING_FOR)
-            return
-
-        self.book(self.LOOKING_FOR)
-
-    def parse_msg(self, msg):
-        msg = msg.split(':')[2]
-        msg = msg.replace("\x01", "")
-        if not msg.startswith("DCC"):
-            self.EXTRA_OUTPUT = msg
-            self.log(msg)
-            return
-
-        args = msg.replace("DCC SEND ", "").split(" ")
-        size = int(args.pop())
-        port = int(args.pop())
-        ip = ip_from_decimal(int(args.pop()))
-        filename = "_".join(args).replace('"', '')
-        self.set_status("RECEIVING")
-
-        n = self.netcat(ip, port, size, filename)
-        files = unar(n, self.PATH)
-        print("Unarchived files", files)
 
-        out = []
-        for f in files:
-            if "searchbot" in f.lower() or "searchook" in f.lower():
-                out.append(self.list_books(f))
-
-        print("Current output", out)
-        self.OUTPUT = []
-        if len(out) > 0:
-            if isinstance(out[0], list):
-                print("1")
-                self.OUTPUT = [item for sublist in out for item in sublist]
-            else:
-                print("2")
-                self.OUTPUT = out
-            print("3")
-            self.OUTPUT = list(set(self.OUTPUT))
-        else:
-            print("4")
-            self.OUTPUT = files
-
-        print(self.OUTPUT)
-
-        if self.TYPE == "BOOK":
-            self.log("Files: ")
-            self.log(self.OUTPUT)
-            for f in self.OUTPUT:
-                if f.lower().endswith(".epub"):
-                    print("EPUB %s" % f)
-                    new_fname = f.replace("epub", "mobi")
-                    p = subprocess.Popen(["ebook-convert", f, new_fname])
-                    p.wait()
-                    self.OUTPUT.append(new_fname)
-                    break
-
-        self.OUTPUT = [i.replace(self.PATH, "") for i in self.OUTPUT]
-        self.handle_close()
-        self.set_status("FINISHED")
-
-    def list_books(self, f):
-        f = open(f, "r")
-        lines = f.readlines()
-        lines = [l.strip().replace('\r', '') for l in lines if self.EXTENSION in l.lower(
-        ) and l.startswith('!') and "htm" not in l.lower()]
-        for l in lines:
-            self.log(l)
-        return lines
-
-    def netcat(self, ip, port, size, filename):
-        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        self.log("Receiving file")
-        s.connect((ip, port))
-        fname = ('%s/%s' % (self.PATH, filename)).replace("//", "/")
-        f = open(fname, 'wb')
-        count = 0
-        while True:
-            data = s.recv(4096)
-            if len(data) == 0:
-                self.log("empty")
-                break
-            count += len(data)
-            f.write(data)
-            perc = int(100 * count / size)
-            self.PROGRESS = perc
-            self.set_status("DOWNLOADING")  # % perc #progress
-
-            self.log("\r%d%%" % perc, instant=True)
-            if count >= size:
-                break
-
-        s.close()
-        f.close()
-        self.log("")  # newline
-        return fname
-
-    def search(self, book):
-        self.log("searching for %s" % book)
-        self.send_queue("PRIVMSG %s :@searchook %s " % (self.CHANNEL, book))
-
-    def pong(self, data):
-        msg = data.replace("PING ", "")
-        self.send_queue("PONG %s" % msg)
-
-    def join_channel(self, channel):
-        self.send_queue("JOIN %s" % channel)
-
-    def send_queue(self, msg):
-        add = bytes(str(msg), "utf-8") + bytes([13, 10])
-        self.socket.send(add)
-
-    def book(self, book):
-        self.log("Asking for '%s'" % book)
-        self.send_queue("PRIVMSG %s :%s " % (self.CHANNEL, book))
-
-    def who(self):
-        self.send_queue("WHO %s" % self.CHANNEL)
-
-    def log(self, val, instant=False):
-        if self.LOGGING:
-            if not instant:
-                print(val)
-            else:
-                sys.stdout.write(val)
-
-    def set_status(self, value):
-        self.STATUS = value
-        self.callback(value=="FINISHED")
-
-
-def ip_from_decimal(dec):
-    return ".".join([str(int(b, 2)) for b in b_octets(bin(dec)[2:])])
-
-def b_octets(l):
-    l = l.zfill(32)
-    return [l[0:8], l[8:16], l[16:24], l[24:32]]
-
-def random_hash():
-    return ("%032x" % random.getrandbits(128))[:10]
+from ircclient import IRCClient, MODE_SEARCH, MODE_BOOK
 
 def usage():
     print("USAGE:")
-    print("\t %s SEARCH <BOOK> <FORMAT>" % sys.argv[0])
-    print("\t %s BOOK <BOT COMMAND>" % sys.argv[0])
-    sys.exit(1)
+    print("\t SEARCH <BOOK> <FORMAT>")
+    print("\t BOOK <BOT COMMAND>")
 
-if __name__ == "__main__":
-    if len(sys.argv) == 4:
-        if(sys.argv[1] != "SEARCH"):
-            usage()
-        mode = sys.argv[1]
-        query = sys.argv[2]
-        grep = sys.argv[3]
+def main():
+    q = queue.Queue()
+    client = IRCClient(q)
+    client.start()
+
+    for line in sys.stdin:
+        line = line.strip().lower()
+        split = shlex.split(line)
 
-    elif len(sys.argv) == 3:
-        if(sys.argv[1] != "BOOK"):
+        if len(split) < 2 or len(split) > 3 or split[0].lower() not in [MODE_SEARCH, MODE_BOOK]:
             usage()
-        mode = sys.argv[1]
-        query = sys.argv[2]
-        grep = ""
-    else:
-        usage()
+            continue
 
-    if mode == "SEARCH":
-        print("Looking for '%s' in format '%s'" % (query, grep))
-    if mode == "BOOK":
-        print("Downloading %s" % query)
+        mode = split[0]
+        query = split[1]
+        grep = ""
+        if len(split) == 3:
+            grep = split[2]
+        data = {'query': query, 'mode': mode, 'grep': grep}
+        q.put(data)
+    client.join()
 
-    client = IRCClient(query, grep, mode, print, logging=True)
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nBye")

ircclient.py

@@ -0,0 +1,275 @@
+#!/usr/bin/env python3
+
+import logging
+import os
+import queue
+import socket
+import subprocess
+import time
+
+from unzipper import unar
+from threading import Thread
+
+logging.basicConfig(level=logging.DEBUG)
+log = logging.getLogger(__name__)
+log.setLevel(logging.DEBUG)
+
+MODE_SEARCH = 'search'
+MODE_BOOK = 'book'
+
+class IRCClient(Thread):
+    TIME_TO_FIRST_COMMAND = 30
+    HOST = "irc.irchighway.net"
+    PORT = 6667
+    IGNORE = [
+        "NOTICE",
+        "PART",
+        "QUIT",
+        "332",
+        "333",
+        "372",
+        "353",
+        "366",
+        "251",
+        "252",
+        "254",
+        "255",
+        "265",
+        "266",
+        "396"]
+    CHANNEL = "#ebooks"
+    PATH = "/tmp/"
+    joined_channel = False
+    connected = False
+    name = "bookbot" + utils.random_hash()
+    readbuffer = b''
+    time_joined = None
+
+    def __init__(self, q):
+        super(IRCClient, self).__init__(daemon=True)
+        self.command_queue = q
+        self.send_queue = queue.Queue()
+
+        nickstr = "NICK %s" % self.name
+        userstr = "USER %s %s bla :%s" % (self.name, self.HOST, self.name)
+
+        self.send_queue.put(nickstr)
+        self.send_queue.put(userstr)
+
+    def run(self):
+        self.handle_connect()
+        while True:
+            if self.connected and self.joined_channel:
+                self.handle_commands()
+            self.process_send_queue()
+            try:
+                self.handle_read()
+            except socket.timeout:
+                continue
+            except socket.error as e:
+                log.error('socket error')
+                log.exception(e)
+                self.handle_connect()
+            except Exception as e:
+                log.exception(e)
+                break
+        self.handle_close()
+        log.info('Exiting RUN')
+
+    def handle_commands(self):
+        if self.command_queue.empty():
+            time.sleep(0.2)
+            return
+
+        elapsed = time.time() - self.time_joined
+        if elapsed < self.TIME_TO_FIRST_COMMAND:
+            log.info("commands to process, but we have to wait %d", self.TIME_TO_FIRST_COMMAND - elapsed)
+            time.sleep(1)
+            return
+        command = self.command_queue.get()
+        log.info("command %s", command)
+        if command['mode'] == MODE_SEARCH:
+            self.send_queue.put("PRIVMSG %s :@searchook %s " % (self.CHANNEL, command['query']))
+            self.EXTENSION = command['grep'] # FIXME 
+            return
+        self.send_queue.put("PRIVMSG %s :%s " % (self.CHANNEL, command['query']))
+
+    def handle_connect(self):
+        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.socket.connect((self.HOST, self.PORT))
+        self.socket.settimeout(2)
+        log.info("connected")
+        self.set_status("CONNECTED")
+        self.connected = True
+
+    def handle_close(self):
+        log.info("closed")
+        self.set_status("DISCONNECTED")
+        self.connected = False
+        self.socket.close()
+
+    def handle_read(self):
+        try:
+            data = self.socket.recv(4096)
+        except socket.error as e:
+            raise
+        except Exception as e:
+            raise
+
+        if len(data) < 2:
+            return
+        # newline at the end?
+        if not (data[-1] == 10 and data[-2] == 13):
+            self.readbuffer += data
+            return
+
+        data = self.readbuffer + data
+        self.readbuffer = b''
+
+        # purge crap
+        data = (data.replace(b'\x95', b'').replace(b'0xc2', b'').decode('utf-8', 'ignore'))
+
+        lines = data.splitlines()
+
+        for line in lines:
+            words = line.split(' ')
+            if len(words) < 2:
+                continue
+            msg_from = words[0]
+            comm = words[1]
+
+            if self.joined_channel:
+                # log.debug(line)
+                pass
+            if comm in self.IGNORE:
+                continue
+
+            if comm == "JOIN":
+                if self.name not in msg_from:  # msg "NICK joined the channel" not about me
+                    continue
+                self.set_status("JOINED")
+                log.info("Joined channel %s" % self.CHANNEL)
+                self.time_joined = time.time()
+                self.joined_channel = True
+                continue
+
+            if comm == "PRIVMSG":
+                # private message not addressed to me
+                if words[2] != self.name:
+                    continue
+                self.parse_msg(line)
+
+            if comm == "PING" or msg_from == "PING":  # respond ping to avoid getting kicked
+                self.pong(line)
+                continue
+            if comm == "376":  # END MOTD
+                # MOTD complete, lets join the channel
+                self.join_channel(self.CHANNEL)
+                continue
+
+    def handle_dcc(self, msg):
+        msg = msg.split(':')[2]
+        msg = msg.replace("\x01", "")
+        log.info('msg %s', msg)
+        if not msg.startswith("DCC"):
+            return
+
+        args = msg.replace("DCC SEND ", "").split(" ")
+        size = int(args.pop())
+        port = int(args.pop())
+        ip = utils.ip_from_decimal(int(args.pop()))
+        filename = "_".join(args).replace('"', '')
+        self.set_status("RECEIVING")
+        return self.netcat(ip, port, size, filename)
+
+    # FIXME rename
+    def parse_msg(self, msg):
+        log.info('complete msg %s', msg)
+        downloaded_filename = self.handle_dcc(msg)
+        files = unar(downloaded_filename, self.PATH)
+        log.info("Unarchived files %s", files)
+
+        list_of_books = False
+        out = []
+        for f in files:
+            if "searchbot" in f.lower() or "searchook" in f.lower():
+                list_of_books = True
+                out.extend(self.list_books(f))
+
+        if list_of_books:
+            log.info("Final output %s", out)
+            return
+
+        log.info("Files: %s", files)
+        ret = []
+        for f in files:
+            if f.lower().endswith(".epub"):
+                log.info("EPUB %s" % f)
+                new_fname = f.replace("epub", "mobi")
+                p = subprocess.Popen(["ebook-convert", f, new_fname])
+                p.wait()
+                ret.append(new_fname)
+            ret.append(f)
+        # make paths absolute?
+        ret = [i.replace(self.PATH, "") for i in ret]
+        for filename in ret:
+            log.info(filename)
+
+    def list_books(self, f):
+        f = open(f, "r")
+        lines = f.readlines()
+        ret = []
+        for l in lines:
+            if self.EXTENSION in l.lower() and l.startswith('!') and "htm" not in l.lower():
+                ret.append(l.strip())
+                log.info("Book matches: %s", l.strip())
+        ret = list(set(ret)) # dedup
+        return ret
+
+    def netcat(self, ip, port, size, filename):
+        filename = os.path.basename(filename).replace(" ", "_")
+        log.info('netcat: %s %d %d %s', ip, port, size, filename)
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        log.info("Receiving file")
+        s.connect((ip, port))
+
+        fname = os.path.join(self.PATH, filename)
+        f = open(fname, 'wb')
+        count = 0
+        while True:
+            data = s.recv(16384)
+            if len(data) == 0:
+                log.info("No data received - finished")
+                break
+            count += len(data)
+            f.write(data)
+            perc = int(100 * count / size)
+            self.PROGRESS = perc
+            self.set_status("DOWNLOADING")  # % perc #progress
+
+            log.info("Download percentage: %d", perc)
+            if count >= size:
+                break
+        s.close()
+        f.close()
+        return fname
+
+    def pong(self, data):
+        msg = data.replace("PING ", "")
+        self.send_queue.put("PONG %s" % msg)
+
+    def join_channel(self, channel):
+        self.send_queue.put("JOIN %s" % channel)
+
+    def process_send_queue(self):
+        if self.send_queue.empty():
+            return
+        data = self.send_queue.get()
+        log.info("Sending %s", data)
+        add = bytes(str(data), "utf-8") + bytes([13, 10])
+        self.socket.send(add)
+
+    def set_status(self, value):
+        self.STATUS = value
+
+

utils.py

@@ -0,0 +1,11 @@
+import random
+
+def random_hash():
+    return ("%032x" % random.getrandbits(128))[:10]
+
+def ip_from_decimal(dec):
+    return ".".join([str(int(b, 2)) for b in b_octets(bin(dec)[2:])])
+
+def b_octets(l):
+    l = l.zfill(32)
+    return [l[0:8], l[8:16], l[16:24], l[24:32]]