Merge pull request #23 from williamhatcher/dev

Encode params into state

.gitignore

@@ -1,4 +1,4 @@
-
+tests/.env
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

flask_discord/client.py

@@ -1,7 +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
+from oauthlib.common import add_params_to_uri, generate_token
+import discord
+import jwt
 
 
 class DiscordOAuth2Session(_http.DiscordOAuth2HttpClient):
@@ -12,40 +16,43 @@ 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: bool = True, params: dict = None):
+    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):
         """Primary method used to create OAuth2 session and redirect users for
         authorization code grant.
 
@@ -54,11 +61,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 : 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.
-        params : dict, optional
+        prompt : str, optional
+        permissions: discord.Permissions object or int, optional
+        guild_id : int, optional
+        disable_guild_select : bool, optional
+        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'.
 
         Returns
         -------
@@ -68,17 +82,37 @@ class DiscordOAuth2Session(_http.DiscordOAuth2HttpClient):
         """
         scope = scope or request.args.get("scope", str()).split() or configs.DISCORD_OAUTH_DEFAULT_SCOPES
 
-        if not prompt and set(scope) & set(configs.DISCORD_PASSTHROUGH_SCOPES):
+        if prompt != "consent" and set(scope) & set(configs.DISCORD_PASSTHROUGH_SCOPES):
             raise ValueError("You should use explicit OAuth grant for passthrough scopes like bot.")
 
-        discord_session = self._make_session(scope=scope)
-        authorization_url, state = discord_session.authorization_url(configs.DISCORD_AUTHORIZATION_BASE_URL)
-        session["DISCORD_OAUTH2_STATE"] = state
+        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)}.")
 
-        prompt = "consent" if prompt else "none"
-        params = params or dict()
-        params.update(prompt=prompt)
-        authorization_url = add_params_to_uri(authorization_url, params)
+        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
+        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)
+        if permissions:
+            authorization_url = add_params_to_uri(authorization_url, {'permissions': permissions})
 
         return redirect(authorization_url)
 
@@ -109,11 +143,24 @@ 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()
         self.save_authorization_token(token)
+        return decoded
 
     def revoke(self):
         """This method clears current discord token, state and all session data from flask

requirements.txt

@@ -1,7 +1,8 @@
-Flask
-cachetools
-setuptools
-requests
-oauthlib
-discord.py
+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
 requests_oauthlib
+PyJWT~=1.7.1

tests/test_app.py

@@ -1,6 +1,6 @@
 import os
 
-from flask import Flask, redirect, url_for
+from flask import Flask, redirect, url_for, session
 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"] = 490732332240863233
+app.config["DISCORD_CLIENT_ID"] = 732337836619202590
 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"
@@ -22,7 +22,13 @@ HYPERLINK = '<a href="{}">{}</a>'
 @app.route("/")
 def index():
     if not discord.authorized:
-        return HYPERLINK.format(url_for(".login"), "Login")
+        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(".invite_oauth"), "Authorize with oauth and bot invite")}
+        """
+
     return f"""
     {HYPERLINK.format(url_for(".me"), "@ME")}<br />
     {HYPERLINK.format(url_for(".logout"), "Logout")}<br />
@@ -36,10 +42,26 @@ 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("/invite/")
+def invite():
+    return discord.create_session(scope=['bot'], permissions=8)
+
+
+@app.route("/invite_oauth/")
+def invite_oauth():
+    return discord.create_session(scope=['bot', 'identify'], permissions=8)
+
+
 @app.route("/callback/")
 def callback():
-    discord.callback()
-    return redirect(url_for(".index"))
+    params = discord.callback()
+    redirect_to = params.get("redirect", "/")
+    return redirect(redirect_to)
 
 
 @app.route("/me/")