diff --git a/example/app.py b/example/app.py index 16cbbf7..ee2e38e 100644 --- a/example/app.py +++ b/example/app.py @@ -13,9 +13,11 @@ from flask.ext.mail import Mail from flask.ext.mongoengine import MongoEngine from flask.ext.sqlalchemy import SQLAlchemy from flask.ext.security import Security, LoginForm, login_required, \ - roles_required, roles_accepted, UserMixin, RoleMixin + roles_required, roles_accepted, UserMixin, RoleMixin from flask.ext.security.datastore import SQLAlchemyUserDatastore, \ - MongoEngineUserDatastore + MongoEngineUserDatastore +from flask.ext.security.decorators import http_auth_required, \ + auth_token_required def create_roles(): @@ -29,7 +31,8 @@ def create_users(): ('jill@lp.com', 'password', ['author'], True), ('tiya@lp.com', 'password', [], False)): current_app.security.datastore.create_user( - email=u[0], password=u[1], roles=u[2], active=u[3]) + email=u[0], password=u[1], roles=u[2], active=u[3], + authentication_token='123abc') def populate_data(): @@ -70,6 +73,16 @@ def create_app(auth_config): def post_login(): return render_template('index.html', content='Post Login') + @app.route('/http') + @http_auth_required + def http(): + return render_template('index.html', content='HTTP Authentication') + + @app.route('/token') + @auth_token_required + def token(): + return render_template('index.html', content='Token Authentication') + @app.route('/post_logout') def post_logout(): return render_template('index.html', content='Post Logout') @@ -116,6 +129,8 @@ def create_sqlalchemy_app(auth_config=None): confirmed_at = db.Column(db.DateTime()) reset_password_token = db.Column(db.String(255)) reset_password_sent_at = db.Column(db.DateTime()) + authentication_token = db.Column(db.String(255)) + authentication_token_created_at = db.Column(db.DateTime()) roles = db.relationship('Role', secondary=roles_users, backref=db.backref('users', lazy='dynamic')) @@ -151,6 +166,8 @@ def create_mongoengine_app(auth_config=None): confirmed_at = db.DateTimeField() reset_password_token = db.StringField(max_length=255) reset_password_sent_at = db.DateTimeField() + authentication_token = db.StringField(max_length=255) + authentication_token_created_at = db.DateTimeField() roles = db.ListField(db.ReferenceField(Role), default=[]) Security(app, MongoEngineUserDatastore(db, User, Role)) diff --git a/flask_security/decorators.py b/flask_security/decorators.py new file mode 100644 index 0000000..12b0461 --- /dev/null +++ b/flask_security/decorators.py @@ -0,0 +1,64 @@ + +from functools import wraps + +from flask import current_app as app, Response, request, abort + +_default_http_auth_msg = """ +
The server could not verify that you are authorized to access the URL + requested. You either supplied the wrong credentials (e.g. a bad password), + or your browser doesn't understand how to supply the credentials required.
+In case you are allowed to request the document, please check your + user-id and password and try again.
+ """ + + +def _check_token(): + header_key = app.security.token_authentication_header + args_key = app.security.token_authentication_key + + header_token = request.headers.get(header_key, None) + token = request.args.get(args_key, header_token) + + try: + app.security.datastore.find_user(authentication_token=token) + except: + return False + + return True + + +def _check_http_auth(): + auth = request.authorization or dict(username=None, password=None) + + try: + user = app.security.datastore.find_user(email=auth.username) + except: + return False + + return app.security.pwd_context.verify(auth.password, user.password) + + +def http_auth_required(fn): + headers = {'WWW-Authenticate': 'Basic realm="Login Required"'} + + @wraps(fn) + def decorated(*args, **kwargs): + if _check_http_auth(): + return fn(*args, **kwargs) + + return Response(_default_http_auth_msg, 401, headers) + + return decorated + + +def auth_token_required(fn): + + @wraps(fn) + def decorated(*args, **kwargs): + if _check_token(): + return fn(*args, **kwargs) + + abort(401) + + return decorated diff --git a/flask_security/tokens.py b/flask_security/tokens.py index f6695d2..19b027f 100644 --- a/flask_security/tokens.py +++ b/flask_security/tokens.py @@ -3,17 +3,17 @@ from datetime import datetime from flask import current_app from flask.ext.security.exceptions import BadCredentialsError, \ - UserNotFoundError, AuthenticationError + UserNotFoundError from flask.ext.security.utils import generate_token from werkzeug.local import LocalProxy -security = LocalProxy(lambda: current_app.security) +datastore = LocalProxy(lambda: current_app.security.datastore) def find_user_by_authentication_token(token): if not token: raise BadCredentialsError('Authentication token required') - return security.datastore.find_user(authentication_token=token) + return datastore.find_user(authentication_token=token) def generate_authentication_token(user): @@ -28,23 +28,21 @@ def generate_authentication_token(user): try: user['authentication_token'] = token - user['authentication_token_generated_at'] = now + user['authentication_token_created_at'] = now except TypeError: user.authentication_token = token - user.authentication_token_generated_at = now + user.authentication_token_created_at = now return user -def authenticate_by_token(token): - try: - return find_user_by_authentication_token(token) - except UserNotFoundError: - raise BadCredentialsError('Invalid authentication token') - except Exception, e: - raise AuthenticationError(str(e)) - - def reset_authentication_token(user): user = generate_authentication_token(user) - security.datastore._save_model(user) + datastore._save_model(user) + return user.authentication_token + + +def ensure_authentication_token(user): + if not user.authentication_token: + reset_authentication_token(user) + return user.authentication_token diff --git a/tests/__init__.py b/tests/__init__.py index e94077a..63cdbbc 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -18,9 +18,10 @@ class SecurityTest(TestCase): def _create_app(self, auth_config): return app.create_sqlalchemy_app(auth_config) - def _get(self, route, content_type=None, follow_redirects=None): + def _get(self, route, content_type=None, follow_redirects=None, headers=None): return self.client.get(route, follow_redirects=follow_redirects, - content_type=content_type or 'text/html') + content_type=content_type or 'text/html', + headers=headers) def _post(self, route, data=None, content_type=None, follow_redirects=True): return self.client.post(route, data=data, diff --git a/tests/functional_tests.py b/tests/functional_tests.py index 177832c..8ce1934 100644 --- a/tests/functional_tests.py +++ b/tests/functional_tests.py @@ -78,6 +78,22 @@ class DefaultSecurityTests(SecurityTest): r = self._get('/admin', follow_redirects=True) self.assertIn('