139 lines
4.3 KiB
Python
139 lines
4.3 KiB
Python
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/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/verify',
|
|
'authy', '', app.config['HTTPCLIENT'],
|
|
JWTAuthTestCase.public_key.name)
|
|
|
|
def tearDown(self):
|
|
finished_database_for_testing(self)
|
|
self.ctx.__exit__(True, None, None)
|
|
|
|
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()
|