add checking for arkose token

linweiyuan · Jul 11, 2023 · eb45ffe · eb45ffe
1 parent 6ecc6c2
commit eb45ffe
Show file tree

Hide file tree

Showing 5 changed files with 52 additions and 12 deletions.
diff --git a/api/chatgpt/api.go b/api/chatgpt/api.go
@@ -54,13 +54,7 @@ func CreateConversation(c *gin.Context) {
 
 	if strings.HasPrefix(request.Model, gpt4Model) {
 		if arkoseTokenUrl == "" {
-			var arkoseToken string
-			var err error
-			if bx == "" {
-				arkoseToken, err = funcaptcha.GetOpenAIToken()
-			} else {
-				arkoseToken, err = funcaptcha.GetOpenAITokenWithBx(bx)
-			}
+			arkoseToken, err := funcaptcha.GetOpenAITokenWithBx(bx)
 			if err != nil {
 				c.AbortWithStatusJSON(http.StatusInternalServerError, api.ReturnMessage(getArkoseTokenErrorMessage))
 				return
@@ -83,8 +77,13 @@ func CreateConversation(c *gin.Context) {
 			defer resp.Body.Close()
 			responseMap := make(map[string]interface{})
 			json.NewDecoder(resp.Body).Decode(&responseMap)
-			arkoseToken := responseMap["token"].(string)
-			request.ArkoseToken = arkoseToken
+			token, ok := responseMap["token"]
+			if !ok || token == "" {
+				c.AbortWithStatusJSON(http.StatusForbidden, api.ReturnMessage(getArkoseTokenErrorMessage))
+				return
+			}
+
+			request.ArkoseToken = token.(string)
 		}
 	}
 

diff --git a/api/token/arkose.go b/api/token/arkose.go
@@ -0,0 +1,32 @@
+package token
+
+import (
+	"os"
+
+	"github.com/gin-gonic/gin"
+	"github.com/linweiyuan/funcaptcha"
+	"github.com/linweiyuan/go-chatgpt-api/api"
+
+	http "github.com/bogdanfinn/fhttp"
+)
+
+var (
+	bx string
+)
+
+//goland:noinspection SpellCheckingInspection
+func init() {
+	bx = os.Getenv("GO_CHATGPT_API_BX")
+}
+
+func GetArkoseToken(c *gin.Context) {
+	token, err := funcaptcha.GetOpenAITokenWithBx(bx)
+	if err != nil {
+		c.AbortWithStatusJSON(http.StatusInternalServerError, api.ReturnMessage("Failed to get arkose token."))
+		return
+	}
+
+	c.JSON(http.StatusOK, gin.H{
+		"token": token,
+	})
+}
diff --git a/go.mod b/go.mod
@@ -10,7 +10,7 @@ require (
 	github.com/gin-gonic/gin v1.9.1
 	github.com/google/uuid v1.3.0
 	github.com/joho/godotenv v1.5.1
-	github.com/linweiyuan/funcaptcha v0.0.0-20230704212036-45a7d90e69e5
+	github.com/linweiyuan/funcaptcha v0.0.0-20230711153901-3330bd4101dd
 	github.com/linweiyuan/go-logger v0.0.0-20230709142852-da1f090a7d4c
 )
 

diff --git a/go.sum b/go.sum
@@ -53,8 +53,8 @@ github.com/klauspost/cpuid/v2 v2.2.4 h1:acbojRNwl3o09bUq+yDCtZFc1aiwaAAxtcn8YkZX
 github.com/klauspost/cpuid/v2 v2.2.4/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY=
 github.com/leodido/go-urn v1.2.4 h1:XlAE/cm/ms7TE/VMVoduSpNBoyc2dOxHs5MZSwAN63Q=
 github.com/leodido/go-urn v1.2.4/go.mod h1:7ZrI8mTSeBSHl/UaRyKQW1qZeMgak41ANeCNaVckg+4=
-github.com/linweiyuan/funcaptcha v0.0.0-20230704212036-45a7d90e69e5 h1:eALnriDN2p7Z8QXT8AqVZ4kytd+Ob4Zcm/QjHX3u1HQ=
-github.com/linweiyuan/funcaptcha v0.0.0-20230704212036-45a7d90e69e5/go.mod h1:ecLx/KqJH1MqzbYwCTyFl8lXqHwGqyO/3ujhT/c/Pq4=
+github.com/linweiyuan/funcaptcha v0.0.0-20230711153901-3330bd4101dd h1:d1ouKWOYZcG4h+7v9ySusLg2tngbd1YH1+z3XKkfxj0=
+github.com/linweiyuan/funcaptcha v0.0.0-20230711153901-3330bd4101dd/go.mod h1:ecLx/KqJH1MqzbYwCTyFl8lXqHwGqyO/3ujhT/c/Pq4=
 github.com/linweiyuan/go-logger v0.0.0-20230709142852-da1f090a7d4c h1:KJqkWkepk+PGSCC3i+ANrLnthM/x9qyfG4uCTGg2B8E=
 github.com/linweiyuan/go-logger v0.0.0-20230709142852-da1f090a7d4c/go.mod h1:U37nbW0kovzaMnrX5sBmxOwidA7P4ymM3/K64oaTQJc=
 github.com/mattn/go-isatty v0.0.19 h1:JITubQf0MOLdlGRuRq+jtsDlekdYPia9ZFsB8h/APPA=

diff --git a/main.go b/main.go
@@ -10,6 +10,7 @@ import (
 	"github.com/linweiyuan/go-chatgpt-api/api/chatgpt"
 	"github.com/linweiyuan/go-chatgpt-api/api/imitate"
 	"github.com/linweiyuan/go-chatgpt-api/api/platform"
+	"github.com/linweiyuan/go-chatgpt-api/api/token"
 	_ "github.com/linweiyuan/go-chatgpt-api/env"
 	"github.com/linweiyuan/go-chatgpt-api/middleware"
 
@@ -32,6 +33,7 @@ func main() {
 	setupPlatformAPIs(router)
 	setupPandoraAPIs(router)
 	setupImitateAPIs(router)
+	setupTokenAPIs(router)
 	router.NoRoute(api.Proxy)
 
 	router.GET("/", func(c *gin.Context) {
@@ -99,3 +101,10 @@ func setupImitateAPIs(router *gin.Engine) {
 		}
 	}
 }
+
+func setupTokenAPIs(router *gin.Engine) {
+	tokenGroup := router.Group("/token")
+	{
+		tokenGroup.GET("/arkose", token.GetArkoseToken)
+	}
+}