Fix queue handling and log level config

Signed-off-by: Thomas Oettli <spacefreak@noop.ch>

setup.py

@@ -5,7 +5,7 @@ def read_file(fname):
         return f.read()
 
 setup(name = "uvscand",
-    version = "0.0.1",
+    version = "0.0.3",
     author = "Thomas Oettli",
     author_email = "spacefreak@noop.ch",
     description = "A python daemon to perform virus scans with uvscan (McAfee) over TCP socket.",

uvscand/__init__.py

@@ -31,32 +31,42 @@ from subprocess import Popen, PIPE
 uvscan_regex = re.compile(r"Found:?(?: the| potentially unwanted program| (?:virus|trojan) or variant)? (.+?)(?:\.| (?:virus |trojan )?)", re.MULTILINE)
 
 
-async def run(uvscan, filename):
+async def uvscan_worker(queue):
+    while True:
+        job = await queue.get()
+        if job is None:
+            await queue.put(None)
+            break
+        uvscan, filename, cb = job
         proc = await asyncio.create_subprocess_exec(uvscan, "--secure", "--mime", "--noboot", "--panalyse", "--manalyse", filename, stdout=asyncio.subprocess.PIPE)
         stdout, _ = await proc.communicate()
-    os.remove(filename)
         if proc.returncode == 13:
             match = uvscan_regex.search(stdout.decode())
             name = match.group(1) if match else "UNKNOWN"
             result = "stream: {} FOUND".format(name)
         else:
             result = "stream: OK"
-    return result
+        cb(result)
 
 
 class AIO(asyncio.Protocol):
     config = None
+    queue = asyncio.Queue()
     separator = b"\x00"
+    completed = False
 
     def __init__(self):
         if not AIO.config:
             raise RuntimeError("configuration not set")
+        if not AIO.queue:
+            raise RuntimeError("queue not set")
         self.logger = logging.getLogger(__name__)
         self.data = bytearray()
+        self.tmpfile = None
 
     def connection_made(self, transport):
         self.peer = transport.get_extra_info("peername")
-        self.logger.debug("new connection from {}".format(self.peer))
+        self.logger.info("new connection from {}".format(self.peer))
         self.transport = transport
         self.request_time = str(time.time())
 
@@ -82,30 +92,51 @@ class AIO(asyncio.Protocol):
                 tmpfile = os.path.join(AIO.config["tmpdir"], "uvscan_{}_{}".format(self.request_time, str(self.peer[1])))
                 self.logger.debug("save data from {} in temporary file {}".format(self.peer, tmpfile))
                 with open(tmpfile, "wb") as f:
+                    self.tmpfile = tmpfile
                     while True:
                         length = struct.unpack(">I", self.data[pos:pos + 4])[0]
                         if length == 0: break
                         pos += 4
                         f.write(self.data[pos:pos + length])
                         pos += length
-                self.logger.debug("starting uvscan for file {}".format(tmpfile))
-                task = asyncio.async(run(AIO.config["uvscan_path"], tmpfile))
-                task.add_done_callback(self.handle_uvscan_result)
+                AIO.queue.put_nowait((AIO.config["uvscan_path"], tmpfile, self.process_uvscan_result))
+                self.logger.info("queued uvscan of {}, queue size is {}".format(tmpfile, AIO.queue.qsize()))
             else:
                 raise RuntimeError("unknown command")
         except (RuntimeError, IndexError, IOError, struct.error) as e:
+            self.logger.warning("warning: {}".format(e))
             self.send_response(str(e))
 
-    def handle_uvscan_result(self, task):
-        self.send_response(task.result())
-
     def send_response(self, response):
-        response = response.encode()
-        response += AIO.separator
+        response = response.encode() + AIO.separator
         self.logger.debug("sending response to {}: {}".format(self.peer, response))
         self.transport.write(response)
         self.transport.close()
 
+    def process_uvscan_result(self, result):
+        self.logger.info("received uvscan result of {}: {}".format(self.tmpfile, result))
+        self.completed = True
+        self.send_response(result)
+
+    def connection_lost(self, exc):
+        if self.tmpfile:
+            if not self.completed:
+                self.logger.warning("client prematurely closed connection, removing {} from scan queue".format(self.tmpfile))
+                entries = []
+                try:
+                    for entry in iter(AIO.queue.get_nowait, None):
+                        if not entry:
+                            continue
+                        if entry[1] != self.tmpfile:
+                            entries.append(entry)
+                except asyncio.QueueEmpty:
+                    pass
+                for entry in entries:
+                    AIO.queue.put_nowait(entry)
+            self.logger.debug("removing temporary file {}".format(self.tmpfile))
+            os.remove(self.tmpfile)
+        self.logger.info("closed connection to {}".format(self.peer))
+
 
 def main():
     "Run uvscand."
@@ -114,6 +145,8 @@ def main():
             formatter_class=lambda prog: argparse.HelpFormatter(prog, max_help_position=45, width=140))
     parser.add_argument("-c", "--config", help="List of config files to read.", nargs="+",
             default=["/etc/uvscand.conf"])
+    parser.add_argument("-m", "--maxprocs", help="Maximum number of parallel scan processes.",
+            type=int, default=8)
     parser.add_argument("-d", "--debug", help="Log debugging messages.", action="store_true")
     args = parser.parse_args()
 
@@ -159,12 +192,23 @@ def main():
             logger.error("option '{}' not present in config section 'uvscand'".format(option))
             sys.exit(1)
 
+    # set loglevel according to config
+    stdouthandler.setLevel(int(config["loglevel"]))
+    sysloghandler.setLevel(int(config["loglevel"]))
+
     # check if uvscan binary exists and is executable
     if not os.path.isfile(config["uvscan_path"]) or not os.access(config["uvscan_path"], os.X_OK):
         logger.error("uvscan binary '{}' does not exist or is not executable".format(config["uvscan_path"]))
         sys.exit(1)
-    loop = asyncio.get_event_loop()
+
+    # setup protocol
     AIO.config = config
+
+    # start uvscan workers
+    loop = asyncio.get_event_loop()
+    workers = [loop.create_task(uvscan_worker(AIO.queue)) for _ in range(args.maxprocs)]
+
+    # start server
     coro = loop.create_server(AIO, config["bind_address"], config["bind_port"])
     server = loop.run_until_complete(coro)
     logger.info("uvscand started")
@@ -174,8 +218,13 @@ def main():
     except KeyboardInterrupt:
         pass
 
+    # close server
     server.close()
     loop.run_until_complete(server.wait_closed())
+
+    # shutdown uvscan workers
+    loop.run_until_complete(AIO.queue.put(None))
+    loop.run_until_complete(asyncio.wait(workers))
     loop.close()
     logger.info("uvscand stopped")