1 files changed, 36 insertions, 36 deletions
diff --git a/fbin-scanner.py b/fbin-scanner.py
index 2e60522..76bc33b 100644
--- a/fbin-scanner.py
+++ b/fbin-scanner.py
@@ -80,47 +80,47 @@ def get_report(dbfile, digest, fileobj):
 
 def main():
     storage = importlib.import_module(current_app.config.get('STORAGE_MODULE', 'fbin.file_storage.filesystem')).Storage(current_app)
-    with session_scope() as session:
-        files = deque(session.query(File).filter(File.scanned == False).all())
-        while len(files):
-            dbfile = files.pop()
-            if not dbfile.get_size():
-                logger.info('Ignoring file %s/%s due to unknown size', dbfile.filename, dbfile.hash)
-                continue
-            if dbfile.get_size() > 32*10**6:
-                logger.info('Ignoring file %s/%s due to size (%s)', dbfile.filename, dbfile.hash, dbfile.formatted_size)
-                continue
-            logger.info('Checking file %s/%s (%s)', dbfile.filename, dbfile.hash, dbfile.formatted_size)
-            try:
-                with storage.temp_file(dbfile) as f:
-                    h = hashlib.sha256()
+    files = deque(db.session.query(File).filter(File.scanned == False).all())
+    while len(files):
+        dbfile = files.pop()
+        if not dbfile.get_size():
+            logger.info('Ignoring file %s/%s due to unknown size', dbfile.filename, dbfile.hash)
+            continue
+        if dbfile.get_size() > 32*10**6:
+            logger.info('Ignoring file %s/%s due to size (%s)', dbfile.filename, dbfile.hash, dbfile.formatted_size)
+            continue
+        logger.info('Checking file %s/%s (%s)', dbfile.filename, dbfile.hash, dbfile.formatted_size)
+        try:
+            with storage.temp_file(dbfile) as f:
+                h = hashlib.sha256()
+                chunk = f.read(2**10*16)
+                while chunk:
+                    h.update(chunk)
                     chunk = f.read(2**10*16)
-                    while chunk:
-                        h.update(chunk)
-                        chunk = f.read(2**10*16)
-                    f.seek(0)
-                    digest = h.hexdigest()
-                    logger.info('SHA-256: %s', digest)
-                    report = get_report(dbfile, digest, f)
-            except:
-                logger.exception('Failed to get report for %s/%s', dbfile.filename, dbfile.hash)
-                # Most likely an error from virustotal, so just break here and retry later.
-                break
-            dbfile.scanned = True
-            if report and any(r.get('detected', False) for r in report['scans'].values()):
-                logger.warning('Positive match')
-                dbfile.blocked_reason = report
-            else:
-                logger.info('No match')
-            session.add(dbfile)
-            session.commit()
-            time.sleep(FILE_DELAY)
-        logger.info('No more files to scan')
+                f.seek(0)
+                digest = h.hexdigest()
+                logger.info('SHA-256: %s', digest)
+                report = get_report(dbfile, digest, f)
+        except:
+            logger.exception('Failed to get report for %s/%s', dbfile.filename, dbfile.hash)
+            # Most likely an error from virustotal, so just break here and retry later.
+            break
+        dbfile.scanned = True
+        if report and any(r.get('detected', False) for r in report['scans'].values()):
+            logger.warning('Positive match')
+            dbfile.blocked_reason = report
+        else:
+            logger.info('No match')
+        db.session.add(dbfile)
+        db.session.commit()
+        time.sleep(FILE_DELAY)
+    logger.info('No more files to scan')
 
 app = Flask('scanner')
 with app.app_context():
     app.config.from_pyfile(args.config_file)
-    from fbin.db import session_scope, File
+    from fbin.db import db, File
+    db.init_app(app)
     config = app.config
     main()