Improved error handling, improved configuration capabilities, added n…

…etwork socket capability to clamav.py

Dockerfile

@@ -56,6 +56,7 @@ COPY LICENSE /app
 
 WORKDIR /app
 ENV PYTHONPATH=${PYTHONPATH}:/app
+ENV CLAMD_CONN=socket
 
 RUN pip3 install poetry==1.1.5
 RUN poetry config virtualenvs.create false

scancan/clamav.py

@@ -1,13 +1,13 @@
 """Clamav Connector"""
-from pyvalve import PyvalveSocket, PyvalveConnectionError
+from pyvalve import PyvalveSocket, PyvalveConnectionError, PyvalveNetwork
 
 
 class ClamAv:
     """
     ClamAv
     Provides an abstraction between Pyvalve and the application
     """
-    def __init__(self) -> None:
+    def __init__(self, conf) -> None:
         """
         ClamAv constructor
 
@@ -16,6 +16,7 @@ def __init__(self) -> None:
         """
         self.pvs = None
         self.logger = None
+        self.conf = conf
 
     def set_logger(self, logger):
         self.logger = logger
@@ -47,7 +48,13 @@ async def instream(self, file):
 
     async def connecting(self):
         self.logger.info("Connecting...")
-        self.pvs = await PyvalveSocket()
+        if self.conf.clamd_conn == 'net':
+            self.pvs = await PyvalveNetwork(
+                self.conf.clamd_host,
+                self.conf.clamd_port
+            )
+        else:
+            self.pvs = await PyvalveSocket(self.conf.clamd_socket)
         self.pvs.set_persistant_connection(True)
 
     async def check_connect(self):

scancan/config.py

@@ -1,5 +1,8 @@
 # Config File #
+import os
 
 upload_size_limit = 104857600
-clamd_host = "localhost"
-clamd_port = 3310
+clamd_conn = os.environ.get('CLAMD_CONN', "net")  # 'net' or 'socket'
+clamd_socket = "/tmp/clamd.socket"
+clamd_host = "127.0.0.1"
+clamd_port = 3310

scancan/main.py

@@ -1,19 +1,20 @@
 """ScanCan Main entry point"""
+import asyncio
 import os
 import re
 import urllib
 from io import BytesIO
 
-from fastapi import FastAPI, HTTPException, File, status
+import aiohttp
+from aiofile import async_open
+from fastapi import FastAPI, File, status
 from fastapi.responses import PlainTextResponse, JSONResponse, FileResponse
 from fastapi.staticfiles import StaticFiles
-import aiohttp
-import asyncio
+from pyvalve import PyvalveConnectionError, PyvalveScanningError
 
-from config import *
-from logger import Logger
+import config as conf
 from clamav import ClamAv
-from aiofile import async_open
+from logger import Logger
 
 logger = Logger(name='ScanCan').get_logger()
 
@@ -36,7 +37,7 @@ async def startup_event():
     """ Startup """
     logger.info("Starting up ScanCan")
     global clamav
-    clamav = ClamAv()
+    clamav = ClamAv(conf)
     clamav.set_logger(logger)
     await clamav.connecting()
 
@@ -60,7 +61,11 @@ async def health():
         Returns:
             result (object)
     """
-    ping_result = await clamav.ping()
+    try:
+        ping_result = await clamav.ping()
+    except PyvalveConnectionError as err:
+        return JSONResponse(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                            content={'result': 'ClamAV connection error.'})
 
     if not re.match(r'^PONG$', ping_result):
         logger.error(ping_result)
@@ -89,7 +94,10 @@ async def scan_path(path: str):
             result (Object)
     """
     logger.info("Scanning path: %s" % path)
-    result = await clamav.scan(path)
+    try:
+        result = await clamav.scan(path)
+    except PyvalveScanningError as err:
+        return JSONResponse(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, content=str(err))
 
     response = {"result": result}
     if re.match(r'^.*\sFOUND$', result):
@@ -127,11 +135,14 @@ async def scan_url(url: str):
         return JSONResponse(
             status_code=status.HTTP_406_NOT_ACCEPTABLE,
             content=response)
-    if len(data) > upload_size_limit:
-        response = {'result': 'Max size %d bytes limit exceeded' % upload_size_limit}
+    if len(data) > conf.upload_size_limit:
+        response = {'result': 'Max size %d bytes limit exceeded' % conf.upload_size_limit}
         return JSONResponse(status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, content=response)
 
-    result = await clamav.instream(BytesIO(data))
+    try:
+        result = await clamav.instream(BytesIO(data))
+    except PyvalveScanningError as err:
+        return JSONResponse(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, content=str(err))
 
     response = {"result": result}
     if re.match(r'^.*\sFOUND$', result):
@@ -149,7 +160,10 @@ async def cont_scan(path: str):
             result (Object)
     """
     logger.info("Scanning path: %s" % path)
-    result = await clamav.contscan(path)
+    try:
+        result = await clamav.contscan(path)
+    except PyvalveScanningError as err:
+        return JSONResponse(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, content=str(err))
 
     response = {"result": result}
     regex = re.compile(r'^.*\sFOUND', re.MULTILINE)
@@ -167,11 +181,14 @@ async def scan_upload_file(file: bytes = File()):
         Returns:
             result (Object)
     """
-    if len(file) > upload_size_limit:
+    if len(file) > conf.upload_size_limit:
         response = {'result': 'Max size %d bytes limit exceeded' % upload_size_limit}
         return JSONResponse(status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, content=response)
 
-    result = await clamav.instream(BytesIO(file))
+    try:
+        result = await clamav.instream(BytesIO(file))
+    except PyvalveScanningError as err:
+        return JSONResponse(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, content=str(err))
 
     response = {"result": result}
     if re.match(r'^.*\sFOUND$', result):