Refactor and improve jwt state manager

flask_discord/_http.py

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

flask_discord/client.py

@@ -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",
+    @staticmethod
-                       permissions: Union[discord.Permissions, int] = None,
+    def __save_state(state):
-                       guild_id: int = None, disable_guild_select: bool = None,
+        session["DISCORD_OAUTH2_STATE"] = state
-                       **params):
+
+    @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
+        data : dict, optional
-        permissions: discord.Permissions object or int, optional
+            A mapping of your any custom data which you want to access after authorization grant. Use
-        guild_id : int, optional
+            `:py:meth:flask_discord.DiscordOAuth2Session.callback` to retrieve this data in your callback view.
-        disable_guild_select : bool, optional
+        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.
+            Additional query parameters to append to authorization URL for customized OAuth flow.
-            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'.
 
         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)
+        state = jwt.encode(data or dict(), current_app.config["SECRET_KEY"]).decode(encoding="utf-8")
-                                            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"))
 
         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
+        self.__save_state(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)
 
-        authorization_url = add_params_to_uri(authorization_url, uri_params)
+        params = params or dict()
-        if permissions:
+        params["prompt"] = "consent" if prompt else "none"
-            authorization_url = add_params_to_uri(authorization_url, {'permissions': permissions})
+        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"]
-
+        state = self.__get_state()
-        # Decode JWT. This only works if the state matches.
+        token = self._fetch_token(state)
-        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()
         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

flask_discord/utils.py

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

requirements.txt

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

tests/test_app.py

@@ -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(".login_with_data"), "Login with custom data")} <br />
-        {HYPERLINK.format(url_for(".invite"), "Invite Bot")} <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/")
+@app.route("/login-data/")
-def login_with_params():
+def login_with_data():
-    return discord.create_session(redirect="/me/", coupon="15off", number=15, zero=0, status=False)
+    return discord.create_session(data=dict(redirect="/me/", coupon="15off", number=15, zero=0, status=False))
 
 
-@app.route("/invite/")
+@app.route("/invite-bot/")
-def invite():
+def invite_bot():
-    return discord.create_session(scope=['bot'], permissions=8)
+    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()
+    data = discord.callback()
-    redirect_to = params.get("redirect", "/")
+    redirect_to = data.get("redirect", "/")
     return redirect(redirect_to)