From 0faa732c9a3e1ffced2b26bee682f513b0e5f0ae Mon Sep 17 00:00:00 2001 From: Jon Bergli Heier Date: Wed, 28 Oct 2020 19:16:34 +0100 Subject: Use flask-sqlalchemy instead of using sqlalchemy directly This makes database access a bit easier and also greatly simplifies some upcoming changes. --- fbin-scanner.py | 72 ++++++++++++++++++++++++++++----------------------------- 1 file changed, 36 insertions(+), 36 deletions(-) (limited to 'fbin-scanner.py') 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() -- cgit v1.2.3