store: add support for pagination

WithSecureLabs · Jul 9, 2020 · 14c7cf6 · 14c7cf6
1 parent 1b82c71
commit 14c7cf6
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 11 deletions.
diff --git a/snake/engines/mongo/file.py b/snake/engines/mongo/file.py
@@ -157,7 +157,7 @@ def select_many(self, sha256_digest=None, file_type=None):
             del data[k]
         return self.db.files.find(data)
 
-    def select_all(self, filter_=None, order=pymongo.DESCENDING, sort=None):
+    def select_all(self, filter_=None, order=pymongo.DESCENDING, sort=None, limit=0, skip=0):
         """Select all files.
 
         Args:
@@ -170,9 +170,9 @@ def select_all(self, filter_=None, order=pymongo.DESCENDING, sort=None):
         """
         documents = []
         if filter_:
-            documents = self.db.files.find(filter_)
+            documents = self.db.files.find(filter_, limit=limit, skip=skip)
         else:
-            documents = self.db.files.find()
+            documents = self.db.files.find(limit=limit, skip=skip)
         if sort:
             documents = documents.sort([(sort, order)]).collation(collation.Collation(locale="en"))
         return documents

diff --git a/snake/routes/store.py b/snake/routes/store.py
@@ -37,7 +37,8 @@ class StoreHandler(snake_handler.SnakeHandler):
     @tornadoparser.use_args({
         # filter[field]: str
         'file_type': fields.Enum(type=enums.FileType, required=False, missing=None),
-        'limit': fields.Int(required=False, missing=None),
+        'from': fields.Int(required=False, missing=0),
+        'limit': fields.Int(required=False, missing=10),
         'operator': fields.Str(required=False, missing='and'),
         'order': fields.Int(required=False, missing=-1),
         'sort': fields.Str(required=False, missing=None)
@@ -51,17 +52,14 @@ async def get(self, data):
                 filter_['$and'] += [{'file_type': data['file_type']}]
         elif data['file_type']:
             filter_ = {'file_type': data['file_type']}
-        cursor = db.async_file_collection.select_all(filter_, data['order'], data['sort'])
-        index = 0
+        # NOTE: With async (motor) there is no count() on cursor so we have to work around that
+        total = await db.async_file_collection.db.files.count_documents(filter_ if filter_ else {})
+        cursor = db.async_file_collection.select_all(filter_, data['order'], data['sort'], data['limit'], data['from'])
         while await cursor.fetch_next:
-            if data['limit']:
-                if index >= data['limit']:
-                    break
-                index += 1
             documents += [cursor.next_object()]
 
         documents = schema.FileSchema(many=True).dump(schema.FileSchema(many=True).load(documents))
-        self.jsonify({'samples': documents})
+        self.jsonify({'samples': documents, 'total': total})
         self.finish()