Cleanup and tests

Author: jvasile

demo.py

@@ -94,19 +94,15 @@ class TokenApiHandler(ForcePost):
     def post(self):
         if not session.get('logged_in', False):
             ret = {'message': False}
-            if app.debug:
-                ret['reason'] = "No login session."
             return ret, 403
 
         username = session.get('username', '')
         if not username:
             # Shouldn't be possible to get here, but just in case
             ret = {'message': False}
-            if app.debug:
-                ret['reason'] = "No username"
             return ret, 403
 
-        valid = api.csrfHandler.token_valid_p(request.values.get('csrf', ''), username)
+        valid = api.csrfHandler.token_valid(request.values.get('csrf', ''), username)
         ret = {'message': valid}
         return ret
 
@@ -129,7 +125,7 @@ def get(self):
         if process_login():
             ret = {
                 'message': 'Login success',
-                'csrf':api.csrfHandler.generate_token(username)
+                'csrf':api.csrfHandler.generate_token()
             }
             return ret
         else:

test_demo.py

@@ -0,0 +1,115 @@
+import simplejson as json
+import requests
+import sys
+
+# Run me with 'pytest test_demo.py'
+
+# Specify a test with 'pytest -k test_foo test_demo.py'
+
+def test_force_post():
+    resp = requests.get("http://localhost:5000/api/secured_endpoint")
+    assert(resp.status_code == 200)
+
+    data = json.loads(resp.text)
+    assert(data['message'] == "Resource only available via POST. Please try again.")
+
+def test_login_fail():
+    resp = requests.get("http://localhost:5000/api/login")
+    assert(resp.status_code == 401)
+
+    data = json.loads(resp.text)
+    assert(data['message'] == "Login fail")
+
+def test_login(session=None):
+    if not session:
+        session = requests.session()
+        
+    params = {'username': 'admin', 'passphrase': 'secret'}
+    resp = session.get("http://localhost:5000/api/login", params=params)
+    assert(resp.status_code == 200)
+
+    data = json.loads(resp.text)
+    assert(data['message'] == "Login success")
+    assert(data['csrf'])
+    
+    cookies = resp.cookies
+    assert(cookies['session'])
+
+    return resp, data, cookies
+
+def test_logged_in():
+    resp = requests.get("http://localhost:5000/api/logged_in")
+    assert(resp.status_code == 200)
+
+    data = json.loads(resp.text)
+    assert(data['message'] == False)
+
+    # Now login and try again
+    resp, data, cookies = test_login()
+    session = requests.session()
+    resp = session.get("http://localhost:5000/api/logged_in", cookies=cookies)
+    data = json.loads(resp.text)
+    assert(data['message'] == True)
+
+def test_token_valid():
+    # No cookie sent
+    session = requests.session()
+    resp = session.post("http://localhost:5000/api/token_valid")
+    assert(resp.status_code == 403)
+    data = json.loads(resp.text)
+    assert(data['message'] == False)
+    
+    # Now, log in and try again
+    resp, data, cookies = test_login(session)
+    params = {'csrf':data['csrf']}
+    print(data)
+    resp = session.post("http://localhost:5000/api/token_valid", cookies=cookies,
+                        params=params)
+    data = json.loads(resp.text)
+    print(data)
+    assert(data['message'] == True)
+
+    # Now with the wrong cookie
+    params = {'csrf':params['csrf']+"wrong"}
+    resp = session.post("http://localhost:5000/api/token_valid", cookies=cookies,
+                        params=params)
+    data = json.loads(resp.text)
+    assert(data['message'] == False)
+    
+def test_csrf_return():
+    resp, data, cookies = test_login()
+    params = {'csrf':data['csrf']}
+
+    # The logged_in endpoint doesn't require a csrf, but should still return
+    # one, since we're logged in.  Let's test it.
+    session = requests.session()
+    resp = session.get("http://localhost:5000/api/logged_in", cookies=cookies, params=params)
+    data = json.loads(resp.text)
+    assert('csrf' in data)
+
+def test_logout():
+    resp, data, cookies = test_login()
+    
+    session = requests.session()
+    resp = session.get("http://localhost:5000/api/logout", cookies=cookies)
+    data = json.loads(resp.text)
+    assert(data['message'] == "Logout success")
+
+def test_secured_endpoint():
+    # Try without logging in
+    session = requests.session()
+    resp = session.post("http://localhost:5000/api/secured_endpoint")
+    assert(resp.status_code == 403)
+    data = json.loads(resp.text)
+    assert(data['message'] == 'csrf fail')
+    
+
+    # Try logged in
+    resp, data, cookies = test_login()
+    params = {'csrf':data['csrf']}
+    resp = session.post("http://localhost:5000/api/secured_endpoint",
+                        cookies=cookies,
+                        params=params)
+    assert(resp.status_code == 200)
+    data = json.loads(resp.text)
+    assert(data['message'] == 'Secret information')