Refactor and improve jwt state manager

This commit is contained in:
thec0sm0s
2020-08-10 11:07:09 +05:30
parent fbb7e23116
commit 1b505fec69
5 changed files with 109 additions and 92 deletions

View File

@@ -64,8 +64,8 @@ class DiscordOAuth2HttpClient(abc.ABC):
def get_authorization_token() -> dict:
raise NotImplementedError
def _fetch_token(self):
discord = self._make_session(state=session.get("DISCORD_OAUTH2_STATE"))
def _fetch_token(self, state):
discord = self._make_session(state=state)
return discord.fetch_token(
configs.DISCORD_TOKEN_URL,
client_secret=self.__client_secret,

View File

@@ -1,11 +1,11 @@
from typing import Union
from . import configs, _http, models
from flask import request, session, redirect
from oauthlib.common import add_params_to_uri, generate_token
import discord
import jwt
import typing
import discord
from . import configs, _http, models, utils
from oauthlib.common import add_params_to_uri
from flask import request, session, redirect, current_app
class DiscordOAuth2Session(_http.DiscordOAuth2HttpClient):
@@ -16,43 +16,51 @@ class DiscordOAuth2Session(_http.DiscordOAuth2HttpClient):
Parameters
----------
`app` : Flask
app : Flask
An instance of your `flask application <http://flask.pocoo.org/docs/1.0/api/#flask.Flask>`_.
`client_id` : int, optional
client_id : int, optional
The client ID of discord application provided. Can be also set to flask config
with key ``DISCORD_CLIENT_ID``.
`client_secret` : str, optional
client_secret : str, optional
The client secret of discord application provided. Can be also set to flask config
with key ``DISCORD_CLIENT_SECRET``.
`redirect_uri` : str, optional
redirect_uri : str, optional
The default URL to use to redirect user to after authorization. Can be also set to flask config
with key ``DISCORD_REDIRECT_URI``.
`bot_token` : str, optional
bot_token : str, optional
The bot token of the application. This is required when you also need to access bot scope resources
beyond the normal resources provided by the OAuth. Can be also set to flask config with
key ``DISCORD_BOT_TOKEN``.
`users_cache` : cachetools.LFUCache, optional
users_cache : cachetools.LFUCache, optional
Any dict like mapping to internally cache the authorized users. Preferably an instance of
cachetools.LFUCache or cachetools.TTLCache. If not specified, default cachetools.LFUCache is used.
Uses the default max limit for cache if ``DISCORD_USERS_CACHE_MAX_LIMIT`` isn't specified in app config.
Attributes
----------
`client_id` : int
client_id : int
The client ID of discord application provided.
`redirect_uri` : str
redirect_uri : str
The default URL to use to redirect user to after authorization.
`users_cache` : cachetools.LFUCache
users_cache : cachetools.LFUCache
A dict like mapping to internally cache the authorized users. Preferably an instance of
cachetools.LFUCache or cachetools.TTLCache. If not specified, default cachetools.LFUCache is used.
Uses the default max limit for cache if ``DISCORD_USERS_CACHE_MAX_LIMIT`` isn't specified in app config.
"""
def create_session(self, scope: list = None, prompt: str = "consent",
permissions: Union[discord.Permissions, int] = None,
guild_id: int = None, disable_guild_select: bool = None,
**params):
@staticmethod
def __save_state(state):
session["DISCORD_OAUTH2_STATE"] = state
@staticmethod
def __get_state():
return session.pop("DISCORD_OAUTH2_STATE", str())
def create_session(
self, scope: list = None, *, data: dict = None, prompt: bool = True,
permissions: typing.Union[discord.Permissions, int] = 0, **params
):
"""Primary method used to create OAuth2 session and redirect users for
authorization code grant.
@@ -61,18 +69,18 @@ class DiscordOAuth2Session(_http.DiscordOAuth2HttpClient):
scope : list, optional
An optional list of valid `Discord OAuth2 Scopes
<https://discordapp.com/developers/docs/topics/oauth2#shared-resources-oauth2-scopes>`_.
prompt : str, optional
permissions: discord.Permissions object or int, optional
guild_id : int, optional
disable_guild_select : bool, optional
data : dict, optional
A mapping of your any custom data which you want to access after authorization grant. Use
`:py:meth:flask_discord.DiscordOAuth2Session.callback` to retrieve this data in your callback view.
prompt : bool, optional
Determines if the OAuth2 grant should be explicitly prompted and re-approved. Defaults to True.
Specify False for implicit grant which will skip the authorization screen and redirect to redirect URI.
permissions: typing.Union[discord.Permissions, int], optional
An optional parameter determining guild permissions of the bot while adding it to a guild using
discord OAuth2 with `bot` scope. It is same as generating so called *bot invite link* which redirects
to your callback endpoint after bot authorization flow. Defaults to 0 or no permissions.
params : kwargs, optional
An optional mapping of query parameters to supply to the authorization URL.
Since query parameters aren't passed through Discord Oauth2, these get added to the state.
Use `:py:meth:`flask_discord.DiscordOAuth2Session.callback()` to retrieve the params passed in.
Notes
-----
`prompt` has been changed. You must specify the raw value ('consent' or 'none'). Defaults to 'consent'.
Additional query parameters to append to authorization URL for customized OAuth flow.
Returns
-------
@@ -82,37 +90,29 @@ class DiscordOAuth2Session(_http.DiscordOAuth2HttpClient):
"""
scope = scope or request.args.get("scope", str()).split() or configs.DISCORD_OAUTH_DEFAULT_SCOPES
if prompt != "consent" and set(scope) & set(configs.DISCORD_PASSTHROUGH_SCOPES):
if not prompt and set(scope) & set(configs.DISCORD_PASSTHROUGH_SCOPES):
raise ValueError("You should use explicit OAuth grant for passthrough scopes like bot.")
if permissions is not None and not (isinstance(permissions, discord.Permissions)
or isinstance(permissions, int)):
raise ValueError(f"permissions must be an int or discord.Permissions, not {type(permissions)}.")
if isinstance(permissions, discord.Permissions):
permissions = permissions.value
# Encode any params into a jwt with the state as the key
# Use generate_token in case state is None
session['DISCORD_JWT_KEY'] = session.get("DISCORD_JWT_KEY", generate_token())
state = jwt.encode(params, session.get("DISCORD_JWT_KEY"))
state = jwt.encode(data or dict(), current_app.config["SECRET_KEY"]).decode(encoding="utf-8")
discord_session = self._make_session(scope=scope, state=state)
authorization_url, state = discord_session.authorization_url(configs.DISCORD_AUTHORIZATION_BASE_URL)
session['DISCORD_OAUTH2_STATE'] = state.decode("utf-8")
# Add special parameters to uri instead of state
uri_params = {'prompt': prompt}
if permissions:
uri_params.update(permissions=permissions)
if guild_id:
uri_params.update(guild_id=guild_id)
if disable_guild_select is not None:
uri_params.update(disable_guild_select=disable_guild_select)
self.__save_state(state)
authorization_url = add_params_to_uri(authorization_url, uri_params)
if permissions:
authorization_url = add_params_to_uri(authorization_url, {'permissions': permissions})
params = params or dict()
params["prompt"] = "consent" if prompt else "none"
if "bot" in scope:
if not isinstance(permissions, (discord.Permissions, int)):
raise ValueError(f"Passed permissions must be an int or discord.Permissions, not {type(permissions)}.")
if isinstance(permissions, discord.Permissions):
permissions = permissions.value
params["permissions"] = permissions
try:
params["disable_guild_select"] = utils.json_bool(params["disable_guild_select"])
except KeyError:
pass
authorization_url = add_params_to_uri(authorization_url, params)
return redirect(authorization_url)
@@ -143,24 +143,14 @@ class DiscordOAuth2Session(_http.DiscordOAuth2HttpClient):
It fetches the authorization token and saves it flask
`session <http://flask.pocoo.org/docs/1.0/api/#flask.session>`_ object.
Raises
------
oauthlib.oauth2.rfc6749.errors.MismatchingStateError
jwt.exceptions.InvalidSignatureError
"""
if request.values.get("error"):
return request.values["error"]
# Decode JWT. This only works if the state matches.
passed_state = request.args.get("state")
jwt_key = session.get("DISCORD_JWT_KEY")
decoded = jwt.decode(passed_state, jwt_key)
# Now that we've decoded the state, we can continue the oauth2 process
token = self._fetch_token()
state = self.__get_state()
token = self._fetch_token(state)
self.save_authorization_token(token)
return decoded
return jwt.decode(state, current_app.config["SECRET_KEY"])
def revoke(self):
"""This method clears current discord token, state and all session data from flask

View File

@@ -1,10 +1,37 @@
"""Few utility functions and decorators."""
import functools
from . import exceptions
from flask import current_app
class JSONBool(object):
def __init__(self, value):
self.value = bool(value)
def __bool__(self):
return self.value
def __str__(self):
return "true" if self else "false"
@classmethod
def from_string(cls, value):
if value.lower() == "true":
return cls(True)
if value.lower() == "false":
return cls(False)
raise ValueError
def json_bool(value):
if isinstance(value, str):
return str(JSONBool.from_string(value))
return str(JSONBool(value))
# Decorators.
def requires_authorization(view):

View File

@@ -1,8 +1,8 @@
Flask~=1.1.2
cachetools~=4.1.1
setuptools~=49.2.1
requests~=2.24.0
oauthlib~=3.1.0
discord.py~=1.4.0
Flask
pyjwt
requests
oauthlib
discord.py
cachetools
setuptools
requests_oauthlib
PyJWT~=1.7.1

View File

@@ -1,6 +1,6 @@
import os
from flask import Flask, redirect, url_for, session
from flask import Flask, redirect, url_for
from flask_discord import DiscordOAuth2Session, requires_authorization
@@ -8,7 +8,7 @@ app = Flask(__name__)
app.secret_key = b"%\xe0'\x01\xdeH\x8e\x85m|\xb3\xffCN\xc9g"
app.config["DISCORD_CLIENT_ID"] = 732337836619202590
app.config["DISCORD_CLIENT_ID"] = 490732332240863233
app.config["DISCORD_CLIENT_SECRET"] = os.getenv("DISCORD_CLIENT_SECRET")
app.config["DISCORD_BOT_TOKEN"] = os.getenv("DISCORD_BOT_TOKEN")
app.config["DISCORD_REDIRECT_URI"] = "http://127.0.0.1:5000/callback"
@@ -24,8 +24,8 @@ def index():
if not discord.authorized:
return f"""
{HYPERLINK.format(url_for(".login"), "Login")} <br />
{HYPERLINK.format(url_for(".login_with_params"), "Login with params")} <br />
{HYPERLINK.format(url_for(".invite"), "Invite Bot")} <br />
{HYPERLINK.format(url_for(".login_with_data"), "Login with custom data")} <br />
{HYPERLINK.format(url_for(".invite_bot"), "Invite Bot with permissions 8")} <br />
{HYPERLINK.format(url_for(".invite_oauth"), "Authorize with oauth and bot invite")}
"""
@@ -42,25 +42,25 @@ def login():
return discord.create_session()
@app.route("/login-params/")
def login_with_params():
return discord.create_session(redirect="/me/", coupon="15off", number=15, zero=0, status=False)
@app.route("/login-data/")
def login_with_data():
return discord.create_session(data=dict(redirect="/me/", coupon="15off", number=15, zero=0, status=False))
@app.route("/invite/")
def invite():
return discord.create_session(scope=['bot'], permissions=8)
@app.route("/invite-bot/")
def invite_bot():
return discord.create_session(scope=["bot"], permissions=8, guild_id=464488012328468480, disable_guild_select=True)
@app.route("/invite_oauth/")
@app.route("/invite-oauth/")
def invite_oauth():
return discord.create_session(scope=['bot', 'identify'], permissions=8)
return discord.create_session(scope=["bot", "identify"], permissions=8)
@app.route("/callback/")
def callback():
params = discord.callback()
redirect_to = params.get("redirect", "/")
data = discord.callback()
redirect_to = data.get("redirect", "/")
return redirect(redirect_to)