From 2eee6f2f472c07f6f1d59dad2129cfe01a4b445e Mon Sep 17 00:00:00 2001
From: Joe <1264204425@qq.com>
Date: Mon, 2 Sep 2024 14:50:45 +0800
Subject: [PATCH] feat: add oauth invite token

---
 api/controllers/console/auth/oauth.py | 5 ++---
 api/libs/oauth.py                     | 9 +++++++--
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/api/controllers/console/auth/oauth.py b/api/controllers/console/auth/oauth.py
index fe64b322d0de17..dc1e17c785c499 100644
--- a/api/controllers/console/auth/oauth.py
+++ b/api/controllers/console/auth/oauth.py
@@ -42,8 +42,7 @@ def get_oauth_providers():
 
 
 class OAuthLogin(Resource):
-    def get(self, provider: str):
-        invite_token = request.args.get("invite_token") or None
+    def get(self, provider: str, invite_toke: Optional[str] = None):
         OAUTH_PROVIDERS = get_oauth_providers()
         with current_app.app_context():
             oauth_provider = OAUTH_PROVIDERS.get(provider)
@@ -51,7 +50,7 @@ def get(self, provider: str):
         if not oauth_provider:
             return {"error": "Invalid provider"}, 400
 
-        auth_url = oauth_provider.get_authorization_url(invite_token=invite_token)
+        auth_url = oauth_provider.get_authorization_url(invite_toke)
         return redirect(auth_url)
 
 
diff --git a/api/libs/oauth.py b/api/libs/oauth.py
index d8ce1a1e6633e6..6b6919de24f90f 100644
--- a/api/libs/oauth.py
+++ b/api/libs/oauth.py
@@ -1,5 +1,6 @@
 import urllib.parse
 from dataclasses import dataclass
+from typing import Optional
 
 import requests
 
@@ -40,12 +41,14 @@ class GitHubOAuth(OAuth):
     _USER_INFO_URL = "https://api.github.com/user"
     _EMAIL_INFO_URL = "https://api.github.com/user/emails"
 
-    def get_authorization_url(self):
+    def get_authorization_url(self, invite_token: Optional[str] = None):
         params = {
             "client_id": self.client_id,
             "redirect_uri": self.redirect_uri,
             "scope": "user:email",  # Request only basic user information
         }
+        if invite_token:
+            params["state"] = invite_token
         return f"{self._AUTH_URL}?{urllib.parse.urlencode(params)}"
 
     def get_access_token(self, code: str):
@@ -90,13 +93,15 @@ class GoogleOAuth(OAuth):
     _TOKEN_URL = "https://oauth2.googleapis.com/token"
     _USER_INFO_URL = "https://www.googleapis.com/oauth2/v3/userinfo"
 
-    def get_authorization_url(self):
+    def get_authorization_url(self, invite_token: Optional[str] = None):
         params = {
             "client_id": self.client_id,
             "response_type": "code",
             "redirect_uri": self.redirect_uri,
             "scope": "openid email",
         }
+        if invite_token:
+            params["state"] = invite_token
         return f"{self._AUTH_URL}?{urllib.parse.urlencode(params)}"
 
     def get_access_token(self, code: str):