api now fetches url

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: