1 files changed, 24 insertions, 8 deletions
diff --git a/src/wsgi.py b/src/wsgi.py
index 98efc6f..49c1c17 100755
--- a/src/wsgi.py
+++ b/src/wsgi.py
@@ -1,15 +1,16 @@
 #! /usr/bin/env python3
 
+import os
 import sys
 from wsgiref.simple_server import make_server
 import falcon
 import json
 from db import DictDB
-import time
 from base64 import b64decode
 
 import authn
 
+
 class CollectorResource():
     def __init__(self, db, users):
         self._db = db
@@ -28,7 +29,7 @@ class CollectorResource():
         try:
             user = userbytes.decode('utf-8')
             pw = pwbytes.decode('utf-8')
-        except:
+        except Exception:
             return None, None   # Fail
         return authfun(user, pw)
 
@@ -38,15 +39,20 @@ class EPGet(CollectorResource):
         resp.status = falcon.HTTP_200
         resp.content_type = falcon.MEDIA_JSON
         out = []
-
         orgs = self.user_auth(req.auth, self._users.read_perms)
+
         if not orgs:
             resp.status = falcon.HTTP_401
             resp.text = 'Invalid username or password\n'
             return
 
+        # We really should rely on req.params in its pure form since
+        # it might contain garbage.
+        selectors = req.params
+
         for org in orgs:
-            out += [{time.ctime(key): dict} for (key, dict) in self._db.search('domain', dict_val=org)]
+            selectors['domain'] = org
+            out.append(self._db.search(**selectors))
 
         resp.text = json.dumps(out) + '\n'
 
@@ -75,7 +81,7 @@ class EPAdd(CollectorResource):
         rawin = req.bounded_stream.read()
         try:
             decodedin = rawin.decode('UTF-8')
-        except:
+        except Exception:
             resp.status = falcon.HTTP_400
             resp.text = 'Need UTF-8\n'
             return
@@ -97,11 +103,10 @@ class EPAdd(CollectorResource):
         resp.text = repr(key) + '\n'
 
 
-def init(url_res_map, addr = '', port = 8000):
+def init(url_res_map, addr='', port=8000):
     app = falcon.App()
     for url, res in url_res_map:
         app.add_route(url, res)
-
     return make_server(addr, port, app)
 
 
@@ -123,7 +128,17 @@ def main():
     #
     #    $ curl -s -u user1:pw1 http://localhost:8000/sc/v0/get | json_pp -json_opt utf8,pretty
 
-    db = DictDB('wsgi_demo.db')
+    try:
+        database = os.environ['DB_NAME']
+        hostname = os.environ['DB_HOSTNAME']
+        username = os.environ['DB_USERNAME']
+        password = os.environ['DB_PASSWORD']
+    except KeyError:
+        print('The environment variables DB_NAME, DB_HOSTNAME, DB_USERNAME ' +
+              'and DB_PASSWORD must be set.')
+        sys.exit(-1)
+
+    db = DictDB(database, hostname, username, password)
     users = authn.UserDB('wsgi_demo_users.yaml')
 
     httpd = init([('/sc/v0/add', EPAdd(db, users)),
@@ -131,5 +146,6 @@ def main():
     print('Serving on port 8000...')
     httpd.serve_forever()
 
+
 if __name__ == '__main__':
     sys.exit(main())