Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import unittest
- from functools import wraps
- import hmac
- from hashlib import sha1
- import flask
- from flask.ext.principal import (Principal, Permission, RoleNeed, Identity,
- identity_changed, identity_loaded,
- current_app)
- def roles_required(*roles):
- """Decorator which specifies that a user must have all the specified roles.
- Example::
- @app.route('/dashboard')
- @roles_required('admin', 'editor')
- def dashboard():
- return 'Dashboard'
- The current user must have both the `admin` role and `editor` role in order
- to view the page.
- :param args: The required roles.
- Source: https://github.com/mattupstate/flask-security/
- """
- def wrapper(fn):
- @wraps(fn)
- def decorated_view(*args, **kwargs):
- perms = [Permission(RoleNeed(role)) for role in roles]
- for perm in perms:
- if not perm.can():
- # return _get_unauthorized_view()
- flask.abort(403)
- return fn(*args, **kwargs)
- return decorated_view
- return wrapper
- def roles_accepted(*roles):
- """Decorator which specifies that a user must have at least one of the
- specified roles. Example::
- @app.route('/create_post')
- @roles_accepted('editor', 'author')
- def create_post():
- return 'Create Post'
- The current user must have either the `editor` role or `author` role in
- order to view the page.
- :param args: The possible roles.
- """
- def wrapper(fn):
- print "roles_accepted: ", roles
- @wraps(fn)
- def decorated_view(*args, **kwargs):
- perm = Permission(*[RoleNeed(role) for role in roles])
- print "roles_accepted.permission:", perm, perm.can()
- if perm.can():
- return fn(*args, **kwargs)
- # return _get_unauthorized_view()
- flask.abort(403)
- return decorated_view
- return wrapper
- def _on_principal_init(sender, identity):
- if identity.id == 'admin':
- identity.provides.add(RoleNeed('admin'))
- identity.provides.add(RoleNeed('member'))
- def create_app():
- app = flask.Flask(__name__)
- app.debug = True
- app.config.update(SECRET_KEY='secret',
- TESTING=True)
- principal = Principal(app)
- identity_loaded.connect(_on_principal_init)
- #
- @app.route('/')
- def index():
- return "OK"
- #
- @app.route('/member')
- @roles_accepted('admin', 'member')
- def role_needed(self):
- return "OK"
- @app.route('/admin')
- @roles_required('admin')
- def connect_admin(self):
- return "OK"
- @app.route('/admin_b')
- @admin_permission.require()
- def connect_admin_alt(self):
- return "OK"
- return app
- admin_permission = Permission(RoleNeed('admin'))
- class WorkshopTest(unittest.TestCase):
- #
- @classmethod
- def setUpClass(cls):
- app = create_app()
- cls.app = app
- cls.client = app.test_client()
- cls.testing = app.test_request_context()
- #
- def test_basic(self):
- r = self.client.get('/')
- print r.status_code
- self.assertEqual(r.data, "OK")
- def test_member(self):
- r = self.client.get('/member')
- self.assertEqual(r.status_code, 403)
- #
- identity_changed.send(current_app, identity=Identity('admin'))
- r = self.client.get('/member')
- self.assertEqual(r.status_code, 200)
- self.assertEqual(r.data, "OK")
- def test_admin_b(self):
- with self.testing as c:
- identity_changed.send(self.app, identity=Identity('admin'))
- r = self.client.get('/admin_b')
- self.assertEqual(r.status_code, 200)
- self.assertEqual(r.data, "OK")
- if __name__ == '__main__':
- unittest.main()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement