change version to 0.0.5

* Update README.md

* Update README.md

* Update README.md

* Update README.md

* Update README.md

* Update README.md

README.md

@@ -3,3 +3,40 @@ A python daemon to perform virus scans with uvscan (McAfee) over TCP socket, mai
 
 ## Developer information
 Everyone who wants to improve or extend this project is very welcome.
+
+## Installation
+git clone https://github.com/spacefreak86/uvscand  
+cd uvscand  
+python3 setup.py build  
+python3 setup.py install  
+  
+cp docs/uvscand.conf /etc/  
+  
+cat << 'EOF'>> /etc/systemd/system/uvscand.service  
+[Unit]  
+Description=uvscand Service  
+After=multi-user.target  
+[Service]  
+Type=simple  
+Restart=always  
+ExecStart=/usr/bin/python3 /usr/local/bin/uvscand  
+[Install]  
+WantedBy=multi-user.target  
+EOF  
+  
+systemctl restart uvscand  
+systemctl status uvscand  
+systemctl enable uvscand  
+  
+cat << 'EOF'>> /etc/rspamd/local.d/antivirus.conf
+uvscan {  
+  scan_mime_parts = true;  
+  scan_text_mime = true;  
+  scan_image_mime = true;  
+  type = "clamav";  
+  symbol = "MCAFEE_VIRUS";  
+  servers = "127.0.0.1:10060";  
+  action = "reject";  
+}  
+EOF  
+systemctl restart rspamd  

setup.py

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

uvscand/__init__.py

@@ -4,12 +4,12 @@
 # it under the terms of the GNU General Public License as published by
 # the Free Software Foundation, either version 3 of the License, or
 # (at your option) any later version.
-# 
+#
 # uvscand 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 
+# 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 PyQuarantineMilter.  If not, see <http://www.gnu.org/licenses/>.
 #
@@ -25,95 +25,158 @@ import struct
 import sys
 import time
 
-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):
-    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
+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 = f"stream: {name} FOUND"
+        else:
+            result = "stream: OK"
+        cb(result)
 
 
 class AIO(asyncio.Protocol):
     config = None
+    queue = asyncio.Queue()
     separator = b"\x00"
 
     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 _send_response(self, response):
+        response = response.encode() + AIO.separator
+        self.logger.debug(f"{self.peer} sending response: {response}")
+        self.transport.write(response)
+        self.transport.close()
 
     def connection_made(self, transport):
         self.peer = transport.get_extra_info("peername")
-        self.logger.debug("new connection from {}".format(self.peer))
+        self.logger.info(f"new connection from {self.peer}")
         self.transport = transport
         self.request_time = str(time.time())
+        self.buffer = bytearray()
+        self.data = bytearray()
+        self.command = None
+        self.length = None
+        self.all_chunks = False
 
     def data_received(self, data):
-        self.logger.debug("data received from {}".format(self.peer))
-        self.data.extend(data)
-        if self.data[-4:] == b"\x00\x00\x00\x00":
-            self.logger.debug("last data chunk received from {}".format(self.peer))
-            self.process_request()
-        else:
-            self.logger.debug("received data chunk from {}".format(self.peer))
-
-    def process_request(self):
         try:
-            if self.data[0] != ord(b"z"):
-                raise RuntimeError("protocol error")
-            pos = self.data.index(ord(AIO.separator))
-            # parse command
-            command = self.data[0:pos].decode()
-            pos += 1
-            if command == "zINSTREAM":
-                # save data chunks to temporary file
-                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:
-                    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)
-            else:
-                raise RuntimeError("unknown command")
+            nbytes = len(data)
+            if self.all_chunks:
+                self.logger.warning(f"{self.peer} received {nbytes} bytes of garbage after last chunk")
+                return
+            self.logger.debug(f"{self.peer} received {nbytes} bytes")
+            self.buffer.extend(data)
+
+            if not self.command:
+                if len(self.buffer) < 10:
+                    return
+                if self.buffer[0] != ord(b"z"):
+                    raise RuntimeError("protocol error")
+                pos = self.buffer.index(ord(AIO.separator))
+                # parse command
+                command = self.buffer[0:pos].decode()
+                if command != "zINSTREAM":
+                    raise RuntimeError("unknown command")
+                self.command = command
+                self.logger.debug(f"{self.peer} command is {command}")
+                pos += 1
+                self.buffer = self.buffer[pos:]
+            if self.command:
+                while True:
+                    if not self.length:
+                        if len(self.buffer) < 4:
+                            break
+                        self.length = struct.unpack(">I", self.buffer[0:4])[0]
+                        self.buffer = self.buffer[4:]
+                        if self.length == 0:
+                            self.all_chunks = True
+                            suffix = str(self.peer[1])
+                            tmpfile = os.path.join(AIO.config["tmpdir"], f"uvscan_{self.request_time}_{suffix}")
+                            self.logger.debug(f"{self.peer} got last chunk, save data to {tmpfile}")
+                            with open(tmpfile, "wb") as f:
+                                self.tmpfile = tmpfile
+                                f.write(self.data)
+                            AIO.queue.put_nowait((AIO.config["uvscan_path"], tmpfile, self.process_uvscan_result))
+                            queuesize = AIO.queue.qsize()
+                            self.logger.info(f"{self.peer} queued uvscan of {tmpfile}, queue size is {queuesize}")
+                            break
+                        self.logger.debug(f"{self.peer} got chunk size of {self.length} bytes")
+                    else:
+                        if len(self.buffer) < self.length:
+                            nbytes = len(self.buffer)
+                            self.logger.debug(f"{self.peer} got {nbytes} of {self.length} bytes")
+                            break
+                        self.logger.debug(f"{self.peer} chunk complete ({self.length} bytes)")
+                        self.data.extend(self.buffer[0:self.length])
+                        self.buffer = self.buffer[self.length:]
+                        self.length = None
+
         except (RuntimeError, IndexError, IOError, struct.error) as e:
-            self.send_response(str(e))
+            self.logger.warning(f"{self.peer} warning: {e}")
+            self._send_response(str(e))
 
-    def handle_uvscan_result(self, task):
-        self.send_response(task.result())
+    def process_uvscan_result(self, result):
+        if not self.tmpfile:
+            return
 
-    def send_response(self, response):
-        response = response.encode()
-        response += AIO.separator
-        self.logger.debug("sending response to {}: {}".format(self.peer, response))
-        self.transport.write(response)
-        self.transport.close()
+        self.logger.info(f"{self.peer} received uvscan result of {self.tmpfile}: {result}")
+        self._send_response(result)
+        self.tmpfile = None
+
+    def connection_lost(self, exc):
+        if not self.tmpfile:
+            self.logger.info(f"closed connection to {self.peer}")
+            return
+
+        entries = []
+        try:
+            for entry in iter(AIO.queue.get_nowait, None):
+                if not entry:
+                    continue
+                if self.tmpfile and entry[1] == self.tmpfile:
+                    self.logger.warning(f"{self.peer} client prematurely closed connection, skipping scan of {self.tmpfile}")
+                    os.remove(self.tmpfile)
+                    self.tmpfile = None
+                    continue
+                entries.append(entry)
+        except asyncio.QueueEmpty:
+            pass
+        for entry in entries:
+            AIO.queue.put_nowait(entry)
+
+        if self.tmpfile:
+            self.logger.warning(f"{self.peer} client prematurely closed connection, but scan is already running")
+            self.tmpfile = None
 
 
 def main():
     "Run uvscand."
     # parse command line
-    parser = argparse.ArgumentParser(description="uvscand daemon",
+    parser = argparse.ArgumentParser(
+            description="uvscand daemon",
             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("-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()
 
@@ -123,8 +186,8 @@ def main():
     syslog_name = logname
     if args.debug:
         loglevel = logging.DEBUG
-        logname = "{}[%(name)s]".format(logname)
-        syslog_name = "{}: [%(name)s] %(levelname)s".format(syslog_name)
+        logname = f"{logname}[%(name)s]"
+        syslog_name = f"{syslog_name}: [%(name)s] %(levelname)s"
 
     root_logger = logging.getLogger()
     root_logger.setLevel(loglevel)
@@ -132,14 +195,14 @@ def main():
     # setup console log
     stdouthandler = logging.StreamHandler(sys.stdout)
     stdouthandler.setLevel(loglevel)
-    formatter = logging.Formatter("%(asctime)s {}: [%(levelname)s] %(message)s".format(logname), datefmt="%Y-%m-%d %H:%M:%S")
+    formatter = logging.Formatter(f"%(asctime)s {logname}: [%(levelname)s] %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
     stdouthandler.setFormatter(formatter)
     root_logger.addHandler(stdouthandler)
 
     # setup syslog
     sysloghandler = logging.handlers.SysLogHandler(address="/dev/log")
     sysloghandler.setLevel(loglevel)
-    formatter = logging.Formatter("{}: %(message)s".format(syslog_name))
+    formatter = logging.Formatter(f"{syslog_name}: %(message)s")
     sysloghandler.setFormatter(formatter)
     root_logger.addHandler(sysloghandler)
 
@@ -156,15 +219,27 @@ def main():
     config = dict(parser.items("uvscand"))
     for option in ["bind_address", "bind_port", "tmpdir", "uvscan_path", "loglevel"]:
         if option not in config.keys():
-            logger.error("option '{}' not present in config section 'uvscand'".format(option))
+            logger.error(f"option '{option}' not present in config section 'uvscand'")
             sys.exit(1)
 
+    if not args.debug:
+        # 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"]))
+        logger.error(f"uvscan binary '{config['uvscan_path']}' does not exist or is not executable")
         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 +249,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")