Ver código fonte

Linted a lot, documented a little

David 10 anos atrás
pai
commit
69e1601bcc
3 arquivos alterados com 313 adições e 247 exclusões
  1. 167 148
      bot.py
  2. 110 73
      qmanager.py
  3. 36 26
      unzipper.py

+ 167 - 148
bot.py

@@ -2,254 +2,273 @@
 
 import random
 import threading
-import re
 import socket
 import sys
 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=""
+    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,logging=False,path="/tmp/"):
-
-        self.TYPE=t
-        self.EXTENSION=extension
-        self.LOOKING_FOR=book
-        self.LOGGING=logging
-        self.PATH=path
-
-        name="bookbot"+self.random_hash()
-        self.NICK=name
-        self.IDENT=name
-        self.REALNAME=name
-
-        self.socket=socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        self.socket.connect( (self.HOST, self.PORT) )
+    readbuffer = b''
+    NOT_EXITED = True
+    EXTRA_OUTPUT = ""
+    PATH = ""
+    joined = False
+    connected = False
+
+    def __init__(self, book, extension, t, 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.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)
+        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 = threading.Thread(target=self.handle_read)
         self.t.start()
         self.send_queue(nickstr)
         self.send_queue(userstr)
 
-    def random_hash(self):
-        return ("%032x" % random.getrandbits(128))[:10]
-
     def handle_connect(self):
         self.log("connected")
-        self.STATUS="CONNECTED"
-        self.connected=True
-        pass
+        self.STATUS = "CONNECTED"
+        self.connected = True
 
     def handle_close(self):
-        self.NOT_EXITED=False
+        self.NOT_EXITED = False
         self.log("closed")
-        self.STATUS="DISCONNECTED"
-        self.connected=False
+        self.STATUS = "DISCONNECTED"
+        self.connected = False
         self.socket.close()
 
     def do_timeout(self):
         if not self.NOT_EXITED:
             return
-        self.NOT_EXITED=False
+        self.NOT_EXITED = False
         self.log("Abort mission")
-        self.STATUS="TIMED OUT"
-        self.connect=False
+        self.STATUS = "TIMED OUT"
+        self.connected = False
         self.socket.close()
 
-
     def handle_read(self):
         while self.NOT_EXITED:
             try:
-                data=self.socket.recv(1024)
-            except e:
+                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
+            if not (data[-1] == 10 and data[-2] == 13):
+                self.readbuffer += data
                 continue
-    
-            data=self.readbuffer+data
-            data=data.replace(b'\x95', b'').replace(b'0xc2',b'').decode('utf-8','ignore') #purge mojibake
-            self.readbuffer=b''
-    
-            lines=data.splitlines()
-    
+
+            data = self.readbuffer + 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:
+                words = line.split(' ')
+                if len(words) < 2:
                     continue
-                msg_from=words[0]
-                comm=words[1]
-    
+                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
+                if comm == "JOIN":
+                    if self.NICK not in msg_from:  # msg "NICK joined the channel" not about me
                         continue
-                    self.STATUS="JOINED"
+                    self.STATUS = "JOINED"
                     self.log("Joined channel %s" % self.CHANNEL)
                     self.run_query()
-                    self.joined=True
+                    self.joined = True
                     continue
-    
-                if comm=="PRIVMSG":
-                    if words[2] != self.NICK: #private message not addressed to me
+
+                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.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
-                    self.join_channel(self.CHANNEL) #MOTD complete, lets join the channel
-                    continue
-    
-                if not self.joined: #waiting 
+                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.STATUS="WAITING"
-        if self.TYPE=="SEARCH":
+        self.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","")
+    def parse_msg(self, msg):
+        msg = msg.split(':')[2]
+        msg = msg.replace("\x01", "")
         if not msg.startswith("DCC"):
-            self.EXTRA_OUTPUT=msg;
+            self.EXTRA_OUTPUT = msg
             self.log(msg)
             return
 
-        args=msg.replace("DCC SEND ","").split(" ")
-        size=int(args.pop())
-        port=int(args.pop())
-        ip=self.ip_from_decimal(int(args.pop()))
-        filename="_".join(args).replace('"','')
-        self.STATUS="RECEIVING"
+        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.STATUS = "RECEIVING"
 
-        n=self.netcat(ip,port,size,filename)
-        files=unar(n,self.PATH)
+        n = self.netcat(ip, port, size, filename)
+        files = unar(n, self.PATH)
 
-        out=[]
+        out = []
         for f in files:
             if "searchbot" in f.lower():
                 out.append(self.list_books(f))
 
-        self.OUTPUT=[]
-        if len(out)>0:
-            if type(out[0]) is list:
-                self.OUTPUT=[item for sublist in out for item in sublist]
+        self.OUTPUT = []
+        if len(out) > 0:
+            if isinstance(out[0], list):
+                self.OUTPUT = [item for sublist in out for item in sublist]
             else:
-                self.OUTPUT=out
-            self.OUTPUT=list(set(self.OUTPUT))
+                self.OUTPUT = out
+            self.OUTPUT = list(set(self.OUTPUT))
         else:
-            self.OUTPUT=files
+            self.OUTPUT = files
 
         if self.TYPE == "BOOK":
             self.log("Files: ")
             self.log(self.OUTPUT)
 
-        self.OUTPUT=[i.replace(self.PATH,"") for i in self.OUTPUT  ]
+        self.OUTPUT = [i.replace(self.PATH, "") for i in self.OUTPUT]
         self.handle_close()
 
-    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() ]
+    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)
+    def netcat(self, ip, port, size, filename):
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self.log("Receiving file")
-        s.connect((ip,port))
-        sizelen=len(str(size))
-        fname=('%s/%s' % (self.PATH,filename)).replace("//","/")
-        f=open(fname, 'wb')
-        count=0
+        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:
+            if len(data) == 0:
                 self.log("empty")
                 break
-            count+=len(data)
+            count += len(data)
             f.write(data)
-            perc=int(100*count/size)
-            self.STATUS="DOWNLOADING" #% perc #progress
-            self.PROGRESS=perc
+            perc = int(100 * count / size)
+            self.STATUS = "DOWNLOADING"  # % perc #progress
+            self.PROGRESS = perc
 
-            self.log("\r%d%%" % perc,instant=True)
+            self.log("\r%d%%" % perc, instant=True)
             if count >= size:
                 break
 
         s.close()
         f.close()
-        self.log("") #newline
+        self.log("")  # newline
         return fname
 
-    def search(self,book):
+    def search(self, book):
         self.log("searching for %s" % book)
-        self.send_queue("PRIVMSG %s :@search %s " % (self.CHANNEL,book))
+        self.send_queue("PRIVMSG %s :@search %s " % (self.CHANNEL, book))
 
-    def pong(self,data):
-        msg=data.replace("PING ","")
+    def pong(self, data):
+        msg = data.replace("PING ", "")
         self.send_queue("PONG %s" % msg)
 
-    def join_channel(self,channel):
+    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])
+    def send_queue(self, msg):
+        add = bytes(str(msg), "utf-8") + bytes([13, 10])
         self.socket.send(add)
 
-    def book(self,book):
+    def book(self, book):
         self.log("Asking for '%s'" % book)
-        self.send_queue("PRIVMSG %s :%s " % (self.CHANNEL,book))
+        self.send_queue("PRIVMSG %s :%s " % (self.CHANNEL, book))
 
     def who(self):
         self.send_queue("WHO %s" % self.CHANNEL)
 
-    def ip_from_decimal(self,dec):
-        return ".".join([ str(int(b,2)) for b in self.b_octets(bin(dec)[2:]) ])
-
-    def b_octets(self,l):
-        l=l.zfill(32)
-        return [l[0:8],l[8:16],l[16:24],l[24:32]]
-
-    def log(self,val,instant=False):
+    def log(self, val, instant=False):
         if self.LOGGING:
             if not instant:
                 print(val)
             else:
                 sys.stdout.write(val)
 
+
+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]
+
 def usage():
     print("USAGE:")
     print("\t %s SEARCH <BOOK> <FORMAT>" % sys.argv[0])
@@ -257,25 +276,25 @@ def usage():
     sys.exit(1)
 
 if __name__ == "__main__":
-    if len(sys.argv)==4:
+    if len(sys.argv) == 4:
         if(sys.argv[1] != "SEARCH"):
             usage()
-        mode  = sys.argv[1]
+        mode = sys.argv[1]
         query = sys.argv[2]
-        grep  = sys.argv[3]
+        grep = sys.argv[3]
 
-    elif len(sys.argv)==3:
+    elif len(sys.argv) == 3:
         if(sys.argv[1] != "BOOK"):
             usage()
-        mode  = sys.argv[1]
+        mode = sys.argv[1]
         query = sys.argv[2]
-        grep  = ""
+        grep = ""
     else:
         usage()
-    
+
     if mode == "SEARCH":
-        print("Looking for '%s' in format '%s'" % (query,grep))
+        print("Looking for '%s' in format '%s'" % (query, grep))
     if mode == "BOOK":
         print("Downloading %s" % query)
 
-    client = IRCClient(query,grep,mode,logging=True)
+    client = IRCClient(query, grep, mode, logging=True)

+ 110 - 73
qmanager.py

@@ -1,113 +1,150 @@
 #!/usr/bin/env python3
+"""Manage a list of book-jobs,
+return a description list for the website"""
 
-from  bot import IRCClient
-import queue
 import shlex
-import asyncore
-import datetime
 import re
-from time import strftime,time
+from time import time
+from bot import IRCClient
+
 
 class qManager:
+    """Main module class"""
 
     tasks = []
-    last_id=0
-    PATH=""
+    last_id = 0
+    PATH = ""
 
-    def __init__(self,path="/tmp/"):
-        self.PATH=path
+    def __init__(self, path="/tmp/"):
+        """Init the manager. Set output path"""
+        self.PATH = path
 
-    def new_dl(self,string):
-        self.new_task(string,IRCClient(string,"","BOOK",logging=False,path=self.PATH))
+    def new_dl(self, string):
+        """Download a book"""
+        irc_client = IRCClient(
+            string,
+            "",
+            "BOOK",
+            logging=False,
+            path=self.PATH)
+        self.new_task(string, irc_client)
 
-    def new_search(self,keywords,format):
-        self.new_task(keywords,IRCClient(keywords,format,"SEARCH",logging=False,path=self.PATH))
+    def new_search(self, keywords, fmt):
+        """Search for a book"""
+        irc_client = IRCClient(
+            keywords,
+            fmt,
+            "SEARCH",
+            logging=False,
+            path=self.PATH)
+        self.new_task(keywords, irc_client)
 
-    def new_task(self,query,t):
-        nt = task(query,t,self.last_id)
-        self.last_id+=1
+    def new_task(self, query, t):
+        """ Add a new 'task' to the list """
+        nt = task(query, t, self.last_id)
+        self.last_id += 1
         self.tasks.append(nt)
 
     def task_status(self):
-        ret=[]
+        """ Return a dict with the status of each 'task' """
+        ret = []
         for t in self.tasks:
-            elapsed= time()-t.START_TIME
+            elapsed = time() - t.START_TIME
             ret.append({"ID": t.ID,
-                "STATUS": t.get_status(),
-                "PROGRESS": t.get_progress(),
-                "OUT": t.get_output(),
-                "ELAPSED": int(elapsed),
-                "QUERY": t.QUERY,
-                "EXTRA": t.get_extra(),
-                "TYPE": t.get_type()})
+                        "STATUS": t.get_status(),
+                        "PROGRESS": t.get_progress(),
+                        "OUT": t.get_output(),
+                        "ELAPSED": int(elapsed),
+                        "QUERY": t.QUERY,
+                        "EXTRA": t.get_extra(),
+                        "TYPE": t.get_type()})
 
         return ret
+
+
 class task:
-    CUR_STATUS=""
-    OUTPUT=""
-    ID=0
-    CLIENT=None
-    START_TIME=None
-    QUERY=""
-    def __init__(self,query,c,id):
-        self.QUERY=query
-        self.CLIENT=c
-        self.ID=id
-        self.START_TIME=time()
-        pass
+    """ Task class. Only returns a status"""
+    CUR_STATUS = ""
+    OUTPUT = ""
+    ID = 0
+    CLIENT = None
+    START_TIME = None
+    QUERY = ""
+
+    def __init__(self, query, c, _id):
+        """Initialize the task"""
+        self.QUERY = query
+        self.CLIENT = c
+        self.ID = _id
+        self.START_TIME = time()
 
     def get_progress(self):
-        return self.CLIENT.PROGRESS;
+        """ Return progress."""
+        return self.CLIENT.PROGRESS
 
     def get_extra(self):
-        return self.CLIENT.EXTRA_OUTPUT;
+        """ Return extra output."""
+        return self.CLIENT.EXTRA_OUTPUT
+
     def get_status(self):
-       if time() - self.START_TIME > 600:
-           self.CLIENT.do_timeout()
-       self.OUTPUT=self.CLIENT.OUTPUT
-       return self.CLIENT.STATUS 
+        """ Return status.  If timed out since last call, stop task """
+        if time() - self.START_TIME > 600:
+            self.CLIENT.do_timeout()
+        self.OUTPUT = self.CLIENT.OUTPUT
+        return self.CLIENT.STATUS
 
     def get_type(self):
-       return self.CLIENT.TYPE
+        """ Return type """
+        return self.CLIENT.TYPE
 
     def get_output(self):
-       if type(self.OUTPUT) == str or self.OUTPUT is None:
-           return self.OUTPUT
-       TAGS_R = re.compile('[\[(].*?[\])]|\.rar|v\d.*?\s')
-       books = self.OUTPUT
-       ret=[]
-       for book in books:
-           book = book.replace("---","")
-           groups = re.search("(?P<BOT>^!\w+)(?P<BOOK>.*?)(?P<INFO>::INFO.*)?$",book)
-           if groups is None:
-               #print("[No groups] %s" % book)
-               return books
-           groups=groups.groupdict()
-           groups["BOOK"]=re.sub(TAGS_R,"",groups["BOOK"]) #remove tags from book
-           groups["TAGS"]=[ r.strip("()[]") for r in re.findall(TAGS_R,book) ]
-           groups["TEXT"]=book
-           ret.append(groups)
-       return ret
-
-#BASIC CLI INTERFACE
+        """Return output:
+           If it's just a string, return it (Books).
+           If it's a list, parse it and return a representative dict.
+           Should cache this or something, doesn't make sense to
+           calculate on each call.
+        """
+        if isinstance(self.OUTPUT, str) or self.OUTPUT is None:
+            return self.OUTPUT
+        TAGS_R = re.compile(r'[\[(].*?[\])]|\.rar|v\d.*?\s')
+        books = self.OUTPUT
+        ret = []
+        for book in books:
+            book = book.replace("---", "")
+            groups = re.search(
+                r"(?P<BOT>^!\w+)(?P<BOOK>.*?)(?P<INFO>::INFO.*)?$", book)
+            if groups is None:
+                #print("[No groups] %s" % book)
+                return books
+            groups = groups.groupdict()
+
+            # remove tags from book
+            groups["BOOK"] = re.sub(TAGS_R, "", groups["BOOK"])
+            groups["TAGS"] = [r.strip("()[]")
+                              for r in re.findall(TAGS_R, book)]
+            groups["TEXT"] = book
+            ret.append(groups)
+        return ret
+
+# BASIC CLI INTERFACE
 if __name__ == "__main__":
     q = qManager()
     while True:
-        line=input()
-        words=shlex.split(line)
+        line = input()
+        words = shlex.split(line)
         print(q.task_status())
-        if len(words)<2:
+        if len(words) < 2:
             continue
-        comm=words[0]
-        if comm.lower()=="quit":
+        comm = words[0]
+        if comm.lower() == "quit":
             break
-        if comm.lower()=="search":
-            if len(words)<3:
+        if comm.lower() == "search":
+            if len(words) < 3:
                 continue
-            q.new_search(words[1],words[2])
+            q.new_search(words[1], words[2])
             continue
-        if comm.lower()=="dl":
+        if comm.lower() == "dl":
             q.new_dl(words[1])
             continue
-    
+
         print("search <'multiple keywords'> <format>|dl <line>|quit")

+ 36 - 26
unzipper.py

@@ -1,54 +1,64 @@
-import zipfile,os.path
+"""Try and uncompress a source file to a dest dir"""
+import zipfile
+import os.path
 import subprocess
 
-USE_UNRAR=True
+USE_UNRAR = True
 
-def unar(source,dest_dir):
+
+def unar(source, dest_dir):
+    """Split input into zip or rar and parse accordingly.
+    Return source if it's not a zip or rar"""
     if source.lower().endswith("zip"):
-        return unzip(source,dest_dir)
+        return unzip(source, dest_dir)
     if source.lower().endswith("rar"):
-        return unrar(source,dest_dir)
+        return unrar(source, dest_dir)
 
     print("NOT RAR? NOT ZIP? I'm panicking.")
     print("I got %s" % source)
     return [source]
 
+
 def unzip(source_filename, dest_dir):
-    out=[]
-    with zipfile.ZipFile(source_filename) as zf:
-        for member in zf.infolist():
+    """Unzip source_filename to dest_dir"""
+    out = []
+    with zipfile.ZipFile(source_filename) as zfile:
+        for member in zfile.infolist():
             # Path traversal defense copied from
             # http://hg.python.org/cpython/file/tip/Lib/http/server.py#l789
             words = member.filename.split('/')
             path = dest_dir
             for word in words[:-1]:
-                drive, word = os.path.splitdrive(word)
-                head, word = os.path.split(word)
-                if word in (os.curdir, os.pardir, ''): continue
+                _, word = os.path.splitdrive(word)
+                _, word = os.path.split(word)
+                if word in (os.curdir, os.pardir, ''):
+                    continue
                 path = os.path.join(path, word)
-            out.append(os.path.join(path,member.filename.split('/')[-1]))
-            zf.extract(member, path)
+            out.append(os.path.join(path, member.filename.split('/')[-1]))
+            zfile.extract(member, path)
     return out
 
 
 def unrar(source, dest_dir):
-    out=[]
-    list_files=[]
-    extract_files=[]
+    """Unzip source to dest_dir. Might use unar or unrar
+    based on the flag USE_UNRAR"""
+    out = []
+    list_files = []
+    extract_files = []
     if USE_UNRAR:
-        list_files=["unrar","lb",source]
-        extract_files=["unrar","x","-o+",source,dest_dir]
+        list_files = ["unrar", "lb", source]
+        extract_files = ["unrar", "x", "-o+", source, dest_dir]
     else:
-        list_files=["lsar",source]
-        extract_files=["unar","-f","-o",dest_dir,source]
+        list_files = ["lsar", source]
+        extract_files = ["unar", "-f", "-o", dest_dir, source]
 
     with subprocess.Popen(list_files, stdout=subprocess.PIPE) as proc:
-        out=proc.stdout.read().decode('utf-8').split("\n")
+        out = proc.stdout.read().decode('utf-8').split("\n")
         if out[0].endswith(": RAR"):
-            del out[0] #header
-        if len(out[-1])==0:
+            del out[0]  # header
+        if len(out[-1]) == 0:
             del out[-1]
-        #print(out)
-    subprocess.run(extract_files,stdout=subprocess.DEVNULL)
+        # print(out)
+    subprocess.run(extract_files, stdout=subprocess.DEVNULL)
 
-    return [os.path.join(dest_dir,file) for file in out ]
+    return [os.path.join(dest_dir, file) for file in out]