import unittest import requests import jwt import base64 from app import app from flask import Flask, abort, jsonify, request, make_response from flask.ext.testing import LiveServerTestCase from initdb import setup_database_for_testing, finished_database_for_testing from data.users import JWTAuthUsers from tempfile import NamedTemporaryFile from Crypto.PublicKey import RSA from datetime import datetime, timedelta class JWTAuthTestCase(LiveServerTestCase): maxDiff = None @classmethod def setUpClass(cls): public_key = NamedTemporaryFile(delete=True) key = RSA.generate(1024) private_key_data = key.exportKey('PEM') pubkey = key.publickey() public_key.write(pubkey.exportKey('OpenSSH')) public_key.seek(0) JWTAuthTestCase.public_key = public_key JWTAuthTestCase.private_key_data = private_key_data def create_app(self): users = [ { 'name': 'cooluser', 'email': 'user@domain.com', 'password': 'password' }, { 'name': 'some.neat.user', 'email': 'neat@domain.com', 'password': 'foobar'} ] jwt_app = Flask('testjwt') private_key = JWTAuthTestCase.private_key_data def _get_basic_auth(): data = base64.b64decode(request.headers['Authorization'][len('Basic '):]) return data.split(':', 1) @jwt_app.route('/user/exists', methods=['GET']) def user_exists(): username, _ = _get_basic_auth() for user in users: if user['name'] == username or user['email'] == username: return 'OK' abort(404) @jwt_app.route('/user/verify', methods=['GET']) def verify_user(): username, password = _get_basic_auth() if username == 'disabled': return make_response('User is currently disabled', 401) for user in users: if user['name'] == username or user['email'] == username: if password != user['password']: return make_response('', 404) token_data = { 'iss': 'authy', 'aud': 'quay.io/jwtauthn', 'nbf': datetime.utcnow(), 'exp': datetime.utcnow() + timedelta(seconds=60), 'sub': user['name'], 'email': user['email'] } encoded = jwt.encode(token_data, private_key, 'RS256') return jsonify({ 'token': encoded }) return make_response('', 404) jwt_app.config['TESTING'] = True return jwt_app def setUp(self): setup_database_for_testing(self) self.app = app.test_client() self.ctx = app.test_request_context() self.ctx.__enter__() self.session = requests.Session() self.jwt_auth = JWTAuthUsers( self.get_server_url() + '/user/exists', self.get_server_url() + '/user/verify', 'authy', JWTAuthTestCase.public_key.name) def tearDown(self): finished_database_for_testing(self) self.ctx.__exit__(True, None, None) def test_user_exists(self): self.assertFalse(self.jwt_auth.user_exists('testuser')) self.assertFalse(self.jwt_auth.user_exists('anotheruser')) self.assertTrue(self.jwt_auth.user_exists('cooluser')) self.assertTrue(self.jwt_auth.user_exists('user@domain.com')) def test_verify_user(self): result, error_message = self.jwt_auth.verify_user('invaliduser', 'foobar') self.assertEquals('Invalid username or password', error_message) self.assertIsNone(result) result, _ = self.jwt_auth.verify_user('cooluser', 'invalidpassword') self.assertIsNone(result) result, _ = self.jwt_auth.verify_user('cooluser', 'password') self.assertIsNotNone(result) self.assertEquals('cooluser', result.username) result, _ = self.jwt_auth.verify_user('some.neat.user', 'foobar') self.assertIsNotNone(result) self.assertEquals('some_neat_user', result.username) def test_confirm_existing_user(self): # Create the users in the DB. result, _ = self.jwt_auth.verify_user('cooluser', 'password') self.assertIsNotNone(result) result, _ = self.jwt_auth.verify_user('some.neat.user', 'foobar') self.assertIsNotNone(result) # Confirm a user with the same internal and external username. result, _ = self.jwt_auth.confirm_existing_user('cooluser', 'password') self.assertIsNotNone(result) self.assertEquals('cooluser', result.username) # Fail to confirm the *external* username, which should return nothing. result, _ = self.jwt_auth.confirm_existing_user('some.neat.user', 'password') self.assertIsNone(result) # Now confirm the internal username. result, _ = self.jwt_auth.confirm_existing_user('some_neat_user', 'foobar') self.assertIsNotNone(result) self.assertEquals('some_neat_user', result.username) def test_disabled_user_custom_erorr(self): result, error_message = self.jwt_auth.verify_user('disabled', 'password') self.assertIsNone(result) self.assertEquals('User is currently disabled', error_message) if __name__ == '__main__': unittest.main()