1 files changed, 28 insertions, 17 deletions
diff --git a/src/wsgi.py b/src/wsgi.py
index 62db3c9..b547ffa 100755
--- a/src/wsgi.py
+++ b/src/wsgi.py
@@ -33,9 +33,11 @@ class CollectorResource():
     def user_auth(self, auth_header, authfun):
         if not auth_header:
             return None, None   # Fail.
+
         BAlit, b64 = auth_header.split()
         if BAlit != "Basic":
             return None, None   # Fail
+
         userbytes, pwbytes = b64decode(b64).split(b':')
         try:
             user = userbytes.decode('utf-8')
@@ -50,14 +52,19 @@ class EPGet(CollectorResource):
         out = list()
         selectors = dict()
 
+        limit = 25
+        skip = 0
+
         resp.status = falcon.HTTP_200
         resp.content_type = falcon.MEDIA_JSON
         orgs = self.user_auth(req.auth, self._users.read_perms)
 
         if not orgs:
             resp.status = falcon.HTTP_401
-            resp.text = json.dumps(
-                {'status': 'error', 'message': 'Invalid username or password\n'})
+            resp.text = json.dumps({
+                'status': 'error',
+                'message': 'Invalid username or password\n'
+            })
             return
 
         if key:
@@ -66,6 +73,10 @@ class EPGet(CollectorResource):
             return
 
         for param in req.params:
+            if param == 'limit':
+                limit = req.params['limit']
+            elif param == 'skip':
+                skip = req.params['skip']
             for i in index.indexes:
                 for j in i['index']['fields']:
                     if j == param:
@@ -73,7 +84,7 @@ class EPGet(CollectorResource):
 
         for org in orgs:
             selectors['domain'] = org
-            data = self._db.search(**selectors)
+            data = self._db.search(**selectors, limit=limit, skip=skip)
             if data:
                 out += data
 
@@ -103,6 +114,7 @@ class EPAdd(CollectorResource):
         # NOTE: Reading the whole body in one go instead of streaming
         # it nicely.
         rawin = req.bounded_stream.read()
+
         try:
             decodedin = rawin.decode('UTF-8')
         except Exception:
@@ -134,25 +146,24 @@ class EPAdd(CollectorResource):
 def main(port=8000, wsgi_helper=False):
     db = DictDB(database, hostname, username, password)
     users = authn.UserDB('wsgi_demo_users.yaml')
-    resources_map = [
-        ('/sc/v0/add', EPAdd(db, users)),
-        ('/sc/v0/get', EPGet(db, users)),
-        ('/sc/v0/get/{key}', EPGet(db, users))
-    ]
-    app = falcon.App(cors_enable=True)
 
-    for url, res in resources_map:
-        app.add_route(url, res)
+    app = falcon.App(cors_enable=True)
+    app.add_route('/sc/v0/add', EPAdd(db, users))
+    app.add_route('/sc/v0/get', EPGet(db, users))
+    app.add_route('/sc/v0/get/{key}', EPGet(db, users))
 
-    if not wsgi_helper:
-        print('Serving on port 8000...')
-        httpd = make_server('', port, app)
-        httpd.serve_forever()
+    if wsgi_helper:
+        return app
 
-    return app
+    print('Serving on port 8000...')
+    httpd = make_server('', port, app)
+    httpd.serve_forever()
 
 
 if __name__ == '__main__':
-    sys.exit(main())
+    try:
+        sys.exit(main())
+    except KeyboardInterrupt:
+        print('\nBye!')
 else:
     app = main(port=8000, wsgi_helper=True)