storing urls in database

1 files changed, 5 insertions, 6 deletions

diff --git a/check/app/models/sql_factory.py b/check/app/models/sql_factory.py
index 68c2e30..5433b67 100644
--- a/check/app/models/sql_factory.py
+++ b/check/app/models/sql_factory.py
@@ -32,7 +32,8 @@ class FileTable(Base):
   __tablename__ = 'files'
   id = Column(Integer, primary_key=True)
   sha256 = Column(String(64, convert_unicode=True), nullable=False)
-  phash = Column(BigInteger, 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)
   def toJSON(self):
     return {
@@ -64,11 +65,9 @@ def search_by_hash(hash):
   match = session.query(FileTable).filter(FileTable.sha256 == hash)
   return match.first()
 
-def add_phash(sha256, phash, ext):
+def add_phash(sha256, phash, ext, url):
   """Add a file to the table"""
-  rec = FileTable(
-    sha256=sha256, phash=phash, ext=ext,
-  )
+  rec = FileTable(sha256=sha256, phash=phash, ext=ext, url=url)
   session = Session()
   session.add(rec)
   session.commit()
@@ -93,4 +92,4 @@ def add_phash_by_filename(path):
 
   hash = sha256(path)
 
-  add_phash(sha256=hash, phash=phash, ext=ext)
+  add_phash(sha256=hash, phash=phash, ext=ext, url=path)