Protecting endpoints in Flask apps by requiring login
In our last blog post on Flask we looked at how to add login and signup to a Flask website. We did this by adding a form, receiving a username and password from the user, and then making sure that they existed in our database.
When a user logged in successfully, we saved their e-mail to a session. When we received an appropriate cookie from a user, we were able to retrieve their session. If there was an e-mail there, we knew the user had previously logged in—since that's the only way an e-mail could be there.
If you haven't read the previous post, I'd definitely recommend it!
Securing an endpoint
"Securing an endpoint" means making sure an endpoint can only be accessed by logged in users. To do this, you must check that the user is logged in before executing any code in the endpoint.
For example, let's take this endpoint as an example:
@app.route("/profile")
def profile():
return render_template("profile.html", name=session["username"])
Here our code is doing session["username"], so this code will fail if the username is not present in the session. The user must be logged in before they can execute this endpoint, but at the moment we aren't checking.
Let's add a check and secure this endpoint:
@app.route("/profile")
def profile():
if "email" not in session:
return redirect(url_for("login"))
return render_template("profile.html", name=session["username"])
We have added two lines that will send the user to the login page if we don't have an e-mail in the session. This is an easy check to do, but if you have many endpoints to secure, you'll end up having this same code in many endpoints.
Writing a decorator instead can help as your code grows.
Securing endpoints with a decorator
A decorator is a function that extends another function. It allows us to run some code before or after running a function, and in our case we'll use a decorator to run that e-mail check before we run the actual endpoint function.
The decorator will be used like this:
@app.route("/profile")
@login_required
def profile():
return render_template("profile.html", name=session["username"])
That @login_required is our decorator being applied to the profile endpoint. Let's learn how to write this decorator.
To define a decorator, we write a function that takes another function as an argument. Here's a decorator that doesn't modify the function at all.
def login_required(func):
return func
When we type @login_required, Python will instantly pass the function underneath the decorator to the login_required function. In our code, the profile function will be passed to the decorator. What the decorator then does is replace the profile function by the return value of the decorator.
So we will want our decorator to return a function that will:
- Check the session contains the e-mail;
- Call the original function if it does;
- Redirect the user to the login page if it does not.
Something like this:
def login_required(func):
def secure_function():
if "email" not in session:
return redirect(url_for("login"))
return func()
return secure_function
Here our decorator takes a function as an argument. It then defines another function called secure_function, and inside secure_function it performs the e-mail check. If no e-mail is found in session, it redirects the user to the login endpoint. Otherwise, it calls the original function (which is the profile function!).
Finally it returns secure_function, and that's what the decorated function will become.
Again, remember the decorator is applied here:
@app.route("/profile")
@login_required
def profile():
return render_template("profile.html", name=session["username"])
So the end result is that our profile function will be called only when the e-mail is in the session.
Keeping function names
Flask uses the function names for some things, such as when using url_for, so when we use this decorator in more than one endpoint we will find problems—because all the decorated endpoints will be replaced by the secure_function function.
We want to make sure that when we replace a function with this new function, we keep the original function's name.
This is required because even if you don't use url_for, Flask requires function names used in endpoints all be unique.
We will keep the original function name with a built-in module called functools, specifically the function wraps:
import functools
def login_required(func):
@functools.wraps(func)
def secure_function():
if "email" not in session:
return redirect(url_for("login"))
return func()
return secure_function
The wraps function modifies secure_function and essentially renames it, giving it the name of func.
Allowing parameters
Now that we're keeping the original function name, we should also keep the original function parameters.
Some of our endpoints may have parameters (they don't at the moment, but they might in the future). It's therefore good practice in decorated functions to allow for parameters.
In order to allow our secure_function to have any number of parameters, we can make use of *args and **kwargs. If you're not familiar with these, here's a nice explanation of what they mean!
import functools
def login_required(func):
@functools.wraps(func)
def secure_function(*args, **kwargs):
if "email" not in session:
return redirect(url_for("login"))
return func(*args, **kwargs)
return secure_function
Here we've allowed for any number of positional arguments (*args), and any number of keyword arguments (**kwargs). This means that when our endpoint function is replaced by secure_function, it can still take the arguments it usually would.
Keeping track of the endpoint that was requested
When we decorate a function with this decorator, we may end up sending the user to the login page instead of to the endpoint they requested.
It's good practice for our application to remember where the user wanted to go, so that when they log in we can continue sending them there, instead of wherever the login endpoint sends them by default.
For example, if a user requests their "dashboard", we should send them there after their log in—and not to their profile!
What we can do for this is, when we send the user to the login endpoint, also include in the URL what part of our application they wanted to access:
def login_required(func):
@functools.wraps(func)
def secure_function(*args, **kwargs):
if "email" not in session:
return redirect(url_for("login", next=request.url))
return func(*args, **kwargs)
return secure_function
Here, next=request.url adds a query string parameter to the URL we're sending the user to—so instead of sending them to something like http://127.0.0.1:5000/login, we'll send them to http://127.0.0.1:5000/login?next=/dashboard. With this information, we could add a few changes to our application to handle that after they log in successfully.
However, we won't be learning how to that in this blog post—it's getting a bit long! Keep an eye out for future blog posts to learn how we'd handle the next parameter there, and send users to their original intended destination! Alternatively, sign up to our e-mail list (there's a form at the bottom of the page) or follow us on Twitter and we'll let you know when new blog posts come out.
Using the decorator
Now that we've defined our decorator, using it is pretty simple!
On any endpoint that you want to protect (i.e. make sure users are logged in before accessing it), just place the decorator above the endpoint definition.
Just make sure the decorator is under the @app.route... line!
@app.route("/profile")
@login_required
def profile():
return render_template("profile.html", name=session["username"])
Wrapping up
I hope you've enjoyed this blog post. We're doing some quite advanced Python now, and this stuff can get confusing very quickly! Make sure to take your time to understand how decorators work if you haven't used them before. They come in handy in many situations!
If you want to learn more about web development with Flask, check our our Complete Python Web Course. If you're newer to Python and would like a more complete understanding of the fundamentals (including decorators!), check our Complete Python Course. Also, sign up for our mailing list below, as we'll be sharing discount codes with our subscribers every month!
Thanks for reading, and I'll see you next time!