fbin-scanner: Add a polling mode for continuous scans

This should make it easier to do scans more often and with less delay
after a file has been uploaded. Activated with the -p/--poll argument.
Without this argument the scanner runs in the default one-shot mode.

1 files changed, 25 insertions, 2 deletions

diff --git a/fbin-scanner.py b/fbin-scanner.py
index 3867bac..2653903 100644
--- a/fbin-scanner.py
+++ b/fbin-scanner.py
@@ -2,6 +2,7 @@
 
 import argparse
 from collections import deque
+import datetime
 import fcntl
 import hashlib
 import importlib
@@ -29,6 +30,7 @@ logger = logging.getLogger('fbin-scanner')
 parser = argparse.ArgumentParser()
 parser.add_argument('--lock-file', default='/tmp/fbin-scanner.lock')
 parser.add_argument('-c', '--config-file', default='fbin/fbin.cfg')
+parser.add_argument('-p', '--poll', action='store_true')
 args = parser.parse_args()
 
 lock_file = open(args.lock_file, 'w')
@@ -80,10 +82,12 @@ def get_report(dbfile, digest, fileobj):
     logger.warning('Unknown response: %s', data)
 
 
-def main():
+def process_files(filters=None):
     storage = importlib.import_module(current_app.config.get('STORAGE_MODULE', 'fbin.file_storage.filesystem')) \
             .Storage(current_app)
-    files = deque(db.session.query(File).filter(File.scanned == False).all())  # noqa: E712
+    files = deque(db.session.query(File).filter(*filters).all())
+    if not files:
+        return False
     while len(files):
         dbfile = files.pop()
         if not dbfile.size:
@@ -124,6 +128,25 @@ def main():
         db.session.commit()
         time.sleep(FILE_DELAY)
     logger.info('No more files to scan')
+    return True
+
+
+def poll_loop():
+    now = datetime.datetime(1970, 1, 1)
+    next_now = datetime.datetime.utcnow()
+    while True:
+        process_files([File.scanned == False, File.date >= now])  # noqa: E712
+        now = next_now
+        next_now = datetime.datetime.utcnow()
+        time.sleep(60)
+
+
+def main():
+    if args.poll:
+        poll_loop()
+    else:
+        if not process_files([File.scanned == False]):  # noqa: E712
+            logger.info('No files to scan were found')
 
 
 app = Flask('scanner')