diff options
Diffstat (limited to 'check/app/models/sql_factory.py')
| -rw-r--r-- | check/app/models/sql_factory.py | 12 |
1 files changed, 9 insertions, 3 deletions
diff --git a/check/app/models/sql_factory.py b/check/app/models/sql_factory.py index 1d32a68..68c2e30 100644 --- a/check/app/models/sql_factory.py +++ b/check/app/models/sql_factory.py @@ -44,11 +44,17 @@ class FileTable(Base): Base.metadata.create_all(engine) -def search_by_phash(phash, threshold=6): +def search_by_phash(phash, threshold=6, limit=1): """Search files for a particular phash""" connection = engine.connect() - cmd = 'SELECT files.*, BIT_COUNT(phash ^ :phash) as hamming_distance FROM files HAVING hamming_distance < :threshold ORDER BY hamming_distance ASC LIMIT 1' - matches = connection.execute(text(cmd), phash=phash, threshold=threshold).fetchall() + cmd = """ + SELECT files.*, BIT_COUNT(phash ^ :phash) + AS hamming_distance FROM files + HAVING hamming_distance < :threshold + ORDER BY hamming_distance ASC + LIMIT :limit + """ + matches = connection.execute(text(cmd), phash=phash, threshold=threshold, limit=limit).fetchall() keys = ('id', 'sha256', 'phash', 'ext', 'score') results = [ dict(zip(keys, values)) for values in matches ] return results |