decorators for basic http auth and token auth

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))

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 = """
+    <h1>Unauthorized</h1>
+    <p>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.</p>
+    <p>In case you are allowed to request the document, please check your
+    user-id and password and try again.</p>
+    """
+
+
+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

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

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,

tests/functional_tests.py

@@ -78,6 +78,22 @@ class DefaultSecurityTests(SecurityTest):
         r = self._get('/admin', follow_redirects=True)
         self.assertIn('<input id="next"', r.data)
 
+    def test_token_auth_via_querystring_valid_token(self):
+        r = self._get('/token?auth_token=123abc')
+        self.assertIn('Token Authentication', r.data)
+
+    def test_token_auth_via_header_valid_token(self):
+        r = self._get('/token', headers={"X-Auth-Token": '123abc'})
+        self.assertIn('Token Authentication', r.data)
+
+    def test_token_auth_via_querystring_invalid_token(self):
+        r = self._get('/token?auth_token=X')
+        self.assertEqual(401, r.status_code)
+
+    def test_token_auth_via_header_invalid_token(self):
+        r = self._get('/token', headers={"X-Auth-Token": 'X'})
+        self.assertEqual(401, r.status_code)
+
 
 class ConfiguredURLTests(SecurityTest):