summaryrefslogtreecommitdiff
path: root/check
diff options
context:
space:
mode:
Diffstat (limited to 'check')
-rw-r--r--check/app/models/sql_factory.py7
-rw-r--r--check/app/server/api.py11
2 files changed, 10 insertions, 8 deletions
diff --git a/check/app/models/sql_factory.py b/check/app/models/sql_factory.py
index 5433b67..ad27f62 100644
--- a/check/app/models/sql_factory.py
+++ b/check/app/models/sql_factory.py
@@ -32,15 +32,16 @@ class FileTable(Base):
__tablename__ = 'files'
id = Column(Integer, primary_key=True)
sha256 = Column(String(64, convert_unicode=True), nullable=False)
- url = Column(String(255, convert_unicode=True), nullable=False)
phash = Column(BigInteger, nullable=False, index=True)
ext = Column(String(4, convert_unicode=True), nullable=False)
+ url = Column(String(255, convert_unicode=True), nullable=False)
def toJSON(self):
return {
'id': self.id,
'sha256': self.sha256,
'phash': self.phash,
'ext': self.ext,
+ 'url': self.url,
}
Base.metadata.create_all(engine)
@@ -56,7 +57,7 @@ def search_by_phash(phash, threshold=6, limit=1):
LIMIT :limit
"""
matches = connection.execute(text(cmd), phash=phash, threshold=threshold, limit=limit).fetchall()
- keys = ('id', 'sha256', 'phash', 'ext', 'score')
+ keys = ('id', 'sha256', 'phash', 'ext', 'url', 'score')
results = [ dict(zip(keys, values)) for values in matches ]
return results
@@ -65,7 +66,7 @@ def search_by_hash(hash):
match = session.query(FileTable).filter(FileTable.sha256 == hash)
return match.first()
-def add_phash(sha256, phash, ext, url):
+def add_phash(sha256=None, phash=None, ext=None, url=None):
"""Add a file to the table"""
rec = FileTable(sha256=sha256, phash=phash, ext=ext, url=url)
session = Session()
diff --git a/check/app/server/api.py b/check/app/server/api.py
index 3742b15..63de5b6 100644
--- a/check/app/server/api.py
+++ b/check/app/server/api.py
@@ -34,9 +34,9 @@ def match():
start = time.time()
try:
- threshold = int(request.args.get('threshold') or 6)
- limit = int(request.args.get('limit') or 1)
- add = str(request.args.get('add') or 'true') == 'true'
+ threshold = int(request.form.get('threshold') or 6)
+ limit = int(request.form.get('limit') or 1)
+ add = str(request.form.get('add') or 'true') == 'true'
except:
return jsonify({
'success': False,
@@ -61,7 +61,7 @@ def match():
im = Image.open(file.stream).convert('RGB')
else:
- url = request.args.get('url')
+ url = request.form.get('url')
if not url:
return jsonify({
'success': False,
@@ -88,7 +88,8 @@ def match():
if len(results) == 0:
if add and url:
- hash = sha256_stream(file)
+ # hash = sha256_stream(file)
+ hash = sha256_stream(io.BytesIO(raw))
add_phash(sha256=hash, phash=phash, ext=ext, url=url)
match = False
else: