refactor: OAuth - redirect direct to provider if just one provider ex…

…ists (#1618) * Redirect direct to provider if just one provider exists * Format using black Co-authored-by: Daniel Vaz Gaspar <[email protected]>
dpgaspar · Jun 11, 2021 · 46b7050 · 46b7050
1 parent 8fe7872
commit 46b7050
Showing 1 changed file with 57 additions and 59 deletions.
diff --git a/flask_appbuilder/security/views.py b/flask_appbuilder/security/views.py
@@ -67,7 +67,7 @@ class PermissionViewModelView(ModelView):
 
 class ResetMyPasswordView(SimpleFormView):
     """
-        View for resetting own user password
+    View for resetting own user password
     """
 
     route_base = "/resetmypassword"
@@ -83,7 +83,7 @@ def form_post(self, form):
 
 class ResetPasswordView(SimpleFormView):
     """
-        View for reseting all users password
+    View for reseting all users password
     """
 
     route_base = "/resetpassword"
@@ -248,49 +248,49 @@ def userinfoedit(self, item):
 
 class UserOIDModelView(UserModelView):
     """
-        View that add OID specifics to User view.
-        Override to implement your own custom view.
-        Then override useroidmodelview property on SecurityManager
+    View that add OID specifics to User view.
+    Override to implement your own custom view.
+    Then override useroidmodelview property on SecurityManager
     """
 
     pass
 
 
 class UserLDAPModelView(UserModelView):
     """
-        View that add LDAP specifics to User view.
-        Override to implement your own custom view.
-        Then override userldapmodelview property on SecurityManager
+    View that add LDAP specifics to User view.
+    Override to implement your own custom view.
+    Then override userldapmodelview property on SecurityManager
     """
 
     pass
 
 
 class UserOAuthModelView(UserModelView):
     """
-        View that add OAUTH specifics to User view.
-        Override to implement your own custom view.
-        Then override userldapmodelview property on SecurityManager
+    View that add OAUTH specifics to User view.
+    Override to implement your own custom view.
+    Then override userldapmodelview property on SecurityManager
     """
 
     pass
 
 
 class UserRemoteUserModelView(UserModelView):
     """
-        View that add REMOTE_USER specifics to User view.
-        Override to implement your own custom view.
-        Then override userldapmodelview property on SecurityManager
+    View that add REMOTE_USER specifics to User view.
+    Override to implement your own custom view.
+    Then override userldapmodelview property on SecurityManager
     """
 
     pass
 
 
 class UserDBModelView(UserModelView):
     """
-        View that add DB specifics to User view.
-        Override to implement your own custom view.
-        Then override userdbmodelview property on SecurityManager
+    View that add DB specifics to User view.
+    Override to implement your own custom view.
+    Then override userdbmodelview property on SecurityManager
     """
 
     add_form_extra_fields = {
@@ -646,50 +646,48 @@ def login(self, provider=None, register=None):
         if g.user is not None and g.user.is_authenticated:
             log.debug("Already authenticated {0}".format(g.user))
             return redirect(self.appbuilder.get_url_for_index)
+
         if provider is None:
-            return self.render_template(
-                self.login_template,
-                providers=self.appbuilder.sm.oauth_providers,
-                title=self.title,
-                appbuilder=self.appbuilder,
-            )
-        else:
-            log.debug("Going to call authorize for: {0}".format(provider))
-            state = jwt.encode(
-                request.args.to_dict(flat=False),
-                self.appbuilder.app.config["SECRET_KEY"],
-                algorithm="HS256",
-            )
-            try:
-                if register:
-                    log.debug("Login to Register")
-                    session["register"] = True
-                if provider == "twitter":
-                    return self.appbuilder.sm.oauth_remotes[
-                        provider
-                    ].authorize_redirect(
-                        redirect_uri=url_for(
-                            ".oauth_authorized",
-                            provider=provider,
-                            _external=True,
-                            state=state,
-                        )
-                    )
-                else:
-                    return self.appbuilder.sm.oauth_remotes[
-                        provider
-                    ].authorize_redirect(
-                        redirect_uri=url_for(
-                            ".oauth_authorized", provider=provider, _external=True
-                        ),
-                        state=state.decode("ascii")
-                        if isinstance(state, bytes)
-                        else state,
+            if len(self.appbuilder.sm.oauth_providers) > 1:
+                return self.render_template(
+                    self.login_template,
+                    providers=self.appbuilder.sm.oauth_providers,
+                    title=self.title,
+                    appbuilder=self.appbuilder,
+                )
+            else:
+                provider = self.appbuilder.sm.oauth_providers[0]["name"]
+
+        log.debug("Going to call authorize for: {0}".format(provider))
+        state = jwt.encode(
+            request.args.to_dict(flat=False),
+            self.appbuilder.app.config["SECRET_KEY"],
+            algorithm="HS256",
+        )
+        try:
+            if register:
+                log.debug("Login to Register")
+                session["register"] = True
+            if provider == "twitter":
+                return self.appbuilder.sm.oauth_remotes[provider].authorize_redirect(
+                    redirect_uri=url_for(
+                        ".oauth_authorized",
+                        provider=provider,
+                        _external=True,
+                        state=state,
                     )
-            except Exception as e:
-                log.error("Error on OAuth authorize: {0}".format(e))
-                flash(as_unicode(self.invalid_login_message), "warning")
-                return redirect(self.appbuilder.get_url_for_index)
+                )
+            else:
+                return self.appbuilder.sm.oauth_remotes[provider].authorize_redirect(
+                    redirect_uri=url_for(
+                        ".oauth_authorized", provider=provider, _external=True
+                    ),
+                    state=state.decode("ascii") if isinstance(state, bytes) else state,
+                )
+        except Exception as e:
+            log.error("Error on OAuth authorize: {0}".format(e))
+            flash(as_unicode(self.invalid_login_message), "warning")
+            return redirect(self.appbuilder.get_url_for_index)
 
     @expose("/oauth-authorized/<provider>")
     def oauth_authorized(self, provider):