Flask session var erased in decorator - python

I have a login_require decorator that needs to set a session variable when email comes as QS.
decorators.py
def login_required(f):
#wraps(f)
def decorated_function(*args, **kwargs):
if session.get('is_logged', False):
email = request.args.get('email', False)
if not email:
session['my_var'] = 'Hello'
return redirect(url_for('login'))
return f(*args, **kwargs)
return decorated_function
views.py
#route('/')
#login_required
def home():
print session.get('my_var')
I am able to login and set the session var my_var for the login() view, but when I reach home() that var is None so I guess the decorator function is erasing it, but I know for sure it was set before in the login() view that the decorator redirected to.
How can I work it?

Related

Multi-Level decorator in Flask

I'm developing Flask application and I have a decorator called login_required which checks if the user is logged in, and sends the current user to the next function.
def login_required(function):
#wraps(function)
def decorator(*args, **kwargs):
token = None
if 'x-access-tokens' in request.headers:
token = request.headers['x-access-tokens']
if not token:
return jsonify({'message': 'a valid token is missing'}), 401
try:
data = jwt.decode(token, app.secret_key)
current_user = User.query.filter_by(username=data['username']).first()
except:
return jsonify({'message': 'token is invalid'}), 401
return function(current_user, *args, **kwargs)
return decorator
So, the callback function is declared like this.
#blueprint.route('/settings', methods=['GET'])
#login_required
def settings(current_user):
# here I check if the current_user is admin or not
...
# function body
Now I want to implement an admin_required decorator which depends on login_required decorator, to be easy using it within functions like this.
#blueprint.route('/settings', methods=['GET'])
#admin_required
def settings(current_user):
...
# function body
How can I achieve this?
So you can create your functionality like this
def decorator1(function):
def fun(*args, **kwargs):
user = function(*args, **kwargs)
print("here")
return user
return fun
def decorator2(function):
def fun(*args, **kwargs):
# firslty i will decorate function with decorator1
decor1 = decorator1(function)
decor1_result = decor1(*args, **kwargs)
# do operation with above result in your case now check if user is admin
print("here in decor2")
return decor1_result + 5
return fun
#decorator2
def test_function(a,b):
return a+b
# Now you want to access a+b and then add c to it
I have included comments for better understanding.
In your case decorator1 will be login_required decorator
and decorator2 will be admin_check decorator. So while creating admin check decorator you can access login_required decorator inside it.

Flask custom decorator not functioning as expected

I'm using Flask and have two decorators that I'm trying to use on certain routes to be more Pythonic, prevent code reuse, and improve readability.
The first of these are just checks to see if the user is logged in else redirect user to the login page. This works fine.
After this I check for the user 'class' (Admin, Manager, Standard User, etc), however this one isn't working and I'm not sure what is missing.
From the routes.py:
#app.route('/user/user-account.html', methods=['GET', 'POST'])
#login_required
def cna_account():
if valid_user() == True:
result = get_user_account()
return render_template('user/user-account.html', result=result)
else:
return redirect(url_for('index'))
#login_reqiured works fine, here is the code:
#wraps(f)
def wrap(*args, **kwargs):
if 'logged_in' in session:
return f(*args, **kwargs)
else:
return redirect(url_for('index'))
return wrap
What I tried for the other decorator that doesn't work:
def valid_user(f):
''' Makes sure that only Base Users can view the Base User pages '''
#wraps(f)
def wrap(*args, **kwargs):
if 'access' in session and session['access'] == 'c':
return f(*args, **kwargs)
else:
return redirect(url_for('index'))
These are stored in a separate module that is imported in routes.py, the only thing I can guess is perhaps the session isn't being carried across even though it's included from flask in the module and routes, but again I'm not sure how this would be remedied.
What I'm trying to do with it is: have the routes that use both decorators and don't require the if valid_user() == True check. Instead it should function and look like:
#app.route('/user/user-account.html', methods=['GET', 'POST'])
#login_required
#valid_user
def cna_account():
result = get_user_account()
return render_template('user/user-account.html', result=result)
Any help on what I'm missing here? Do I need to pass the session variable as an argument to #valid_user? I tried that a few different ways and still had it throw errors.
Much appreciated!
You need to return the inner function wrap in valid_user:
def valid_user(f):
#wraps(f)
def wrap(*args, **kwargs):
if 'access' in session and session['access'] == 'c':
return f(*args, **kwargs)
return redirect(url_for('index'))
return wrap
While it is perfectly valid to not return anything from an outer decorator function, bear in mind that the wrapped function is passed to the wrapper at run time, thus, the returned value will be None:
def foo(f):
print("inside decorator with '{}'".format(f.__name__))
def inner():
return 10
#foo
def bar():
return 'in bar'
"inside decorator with 'bar'"
>>>bar()
Traceback (most recent call last):
File "", line 1, in
TypeError: 'NoneType' object is not callable

Python: how to make this more generic - probably with a decorator?

currently I have this code
#app.route("/protect1")
def protect1():
if not session.get('logged_in'):
session['next'] = "/protect1"
return render_template('login.html')
else:
return "This is the first protected page protect1"
#app.route("/protect2")
def protect2():
if not session.get('logged_in'):
session['next'] = "/protect2"
return render_template('login.html')
else:
return "This is the second protected page protect2"
in my flask app, everything working fine. Only it is not nice that I will need to repeat for each function (view) the if/else combination.
I would prefer to have some generic way, like this pseude-code:
#checklogin
#app.route("/protect1")
def protect1():
return "This is the first protected page protect1"
#checklogin
#app.route("/protect2")
def protect2():
return "This is the second protected page protect2"
One challenge here is that the #checklogin decorator would need to know the app.route path (e.g. "/protect1") in order to be able to set session['next'] correctly. I have no idea how to pass this parameter to the decorator, especially how to find it out in the first place. In other words, how does the function protect1() know that it is decorated with #app.route and which parameter ("/protect1") has been passed to that app.route decorator?
The decorator can look up the path on request; either using the URL that was loaded (available as request.url) or the request.endpoint attribute:
from functools import wraps
from flask import request, session
def checklogin(f):
#wraps(f)
def wrapper(*args, **kwargs):
if not session.get('logged_in'):
session['next'] = request.url
return render_template('login.html')
return f(*args, **kwargs)
return wrapper
Do place the decorator after the app.route() decorator or it'll not be registered as the handler for the route:
#app.route("/protect1")
#checklogin
def protect1():
return "This is the first protected page protect1"

Use decorators check user login in flask?

I defined a check user method:
from functools import wraps
def check_user(func):
#wraps(func)
def wrapper(*args, **kwargs):
if session['logged_in']:
return func(*args, **kwargs)
else:
return 'Log in'
return wrapper
#app.route('/test')
#check_user
def test():
return "Hello"
It does not work. how can I correct it?
It seems you don't know how to create decorators in python. There are many helpful answers on this question: How can I make a chain of function decorators in Python?
Below is how you can create a decorator that checks if a user is logged in.
from functools import wraps
def checkuser(func):
"""Checks whether user is logged in or raises error 401."""
#wraps(func)
def wrapper(*args, **kwargs):
if not g.user:
abort(401)
return func(*args, **kwargs)
return wrapper
The decorator above will raise a 401 error if a user is not logged in. It will return the view function otherwise.

Unit-testing a flask-principal application

All, I'm writing a flask application that depends on flask-principal for managing user roles. I'd like to write some simple unit tests to check which views can be accessed by which user. An example of code is posted on pastebin to avoid cluttering this post. In short, I define a few routes, decorating some so that they can be accessed only by users with the proper role, then try to access them in a test.
In the code pasted, the test_member and test_admin_b both fail, complaining about a PermissionDenied. Obviously, I'm failing to declare the user properly; at least, the info about the user roles is not in the right context.
Any help or insight about the complexities of context processing will be deeply appreciated.
Flask-Principal does not store information for you between requests. It's up to you to do this however you like. Keep that in mind and think about your tests for a moment. You call the test_request_context method in the setUpClass method. This creates a new request context. You are also making test client calls with self.client.get(..) in your tests. These calls create additional request contexts that are not shared between each other. Thus, your calls to identity_changed.send(..) do not happen with the context of the requests that are checking for permissions. I've gone ahead and edited your code to make the tests pass in hopes that it will help you understand. Pay special attention to the before_request filter I added in the create_app method.
import hmac
import unittest
from functools import wraps
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):
#wraps(fn)
def decorated_view(*args, **kwargs):
perm = Permission(*[RoleNeed(role) for role in roles])
if perm.can():
return fn(*args, **kwargs)
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.before_request
def determine_identity():
# This is where you get your user authentication information. This can
# be done many ways. For instance, you can store user information in the
# session from previous login mechanism, or look for authentication
# details in HTTP headers, the querystring, etc...
identity_changed.send(current_app._get_current_object(), identity=Identity('admin'))
#app.route('/')
def index():
return "OK"
#app.route('/member')
#roles_accepted('admin', 'member')
def role_needed():
return "OK"
#app.route('/admin')
#roles_required('admin')
def connect_admin():
return "OK"
#app.route('/admin_b')
#admin_permission.require()
def connect_admin_alt():
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()
def test_basic(self):
r = self.client.get('/')
self.assertEqual(r.data, "OK")
def test_member(self):
r = self.client.get('/member')
self.assertEqual(r.status_code, 200)
self.assertEqual(r.data, "OK")
def test_admin_b(self):
r = self.client.get('/admin_b')
self.assertEqual(r.status_code, 200)
self.assertEqual(r.data, "OK")
if __name__ == '__main__':
unittest.main()
As Matt explained, it's only a matter of context. Thanks to his explanations, I came with two different ways to switch identities during unit tests.
Before all, let's modify a bit the application creation:
def _on_principal_init(sender, identity):
"Sets the roles for the 'admin' and 'member' identities"
if identity.id:
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():
return "OK"
#
#app.route('/admin')
#roles_required('admin')
def connect_admin():
return "OK"
# Using `flask.ext.principal` `Permission.require`...
# ... instead of Matt's decorators
#app.route('/admin_alt')
#admin_permission.require()
def connect_admin_alt():
return "OK"
return app
A first possibility is to create a function that loads an identity before each request in our test. The easiest is to declare it in the setUpClass of the test suite after the app is created, using the app.before_request decorator:
class WorkshopTestOne(unittest.TestCase):
#
#classmethod
def setUpClass(cls):
app = create_app()
cls.app = app
cls.client = app.test_client()
#app.before_request
def get_identity():
idname = flask.request.args.get('idname', '') or None
print "Notifying that we're using '%s'" % idname
identity_changed.send(current_app._get_current_object(),
identity=Identity(idname))
Then, the tests become:
def test_admin(self):
r = self.client.get('/admin')
self.assertEqual(r.status_code, 403)
#
r = self.client.get('/admin', query_string={'idname': "member"})
self.assertEqual(r.status_code, 403)
#
r = self.client.get('/admin', query_string={'idname': "admin"})
self.assertEqual(r.status_code, 200)
self.assertEqual(r.data, "OK")
#
def test_admin_alt(self):
try:
r = self.client.get('/admin_alt')
except flask.ext.principal.PermissionDenied:
pass
#
try:
r = self.client.get('/admin_alt', query_string={'idname': "member"})
except flask.ext.principal.PermissionDenied:
pass
#
try:
r = self.client.get('/admin_alt', query_string={'idname': "admin"})
except flask.ext.principal.PermissionDenied:
raise
self.assertEqual(r.data, "OK")
(Incidentally, the very last test shows that Matt's decorator are far easier to use....)
A second approach uses the test_request_context function with a with ... to create a temporary context. No need to define a function decorated by #app.before_request, just pass the route to test as argument of test_request_context, send the identity_changed signal in the context and use the .full_dispatch_request method
class WorkshopTestTwo(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_admin(self):
with self.testing("/admin") as c:
r = c.app.full_dispatch_request()
self.assertEqual(r.status_code, 403)
#
with self.testing("/admin") as c:
identity_changed.send(c.app, identity=Identity("member"))
r = c.app.full_dispatch_request()
self.assertEqual(r.status_code, 403)
#
with self.testing("/admin") as c:
identity_changed.send(c.app, identity=Identity("admin"))
r = c.app.full_dispatch_request()
self.assertEqual(r.status_code, 200)
self.assertEqual(r.data, "OK")
Along Matt's response, I've created a context manager to make the determine_identity a little cleaner:
#contextmanager
def identity_setter(app, user):
#app.before_request
def determine_identity():
#see http://stackoverflow.com/questions/16712321/unit-testing-a-flask-principal-application for details
identity_changed.send(current_app._get_current_object(), identity=Identity(user.id))
determine_identity.remove_after_identity_test = True
try:
yield
finally:
#if there are errors in the code under trest I need this to be run or the addition of the decorator could affect other tests
app.before_request_funcs = {None: [e for e in app.before_request_funcs[None] if not getattr(e,'remove_after_identity_test', False)]}
So when I run my test it looks like:
with identity_setter(self.app,user):
with user_set(self.app, user):
with self.app.test_client() as c:
response = c.get('/orders/' + order.public_key + '/review')
I hope this helps, and I would welcome any feedback :)
~Victor

Categories

Resources