From 8c424fd55360ef6ca8cc3b85539af5d076ed0453 Mon Sep 17 00:00:00 2001
From: liuhuapiaoyuan <278780765@qq.com>
Date: Sun, 21 Jan 2024 19:37:58 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E8=B4=A6=E5=8F=B7?=
 =?UTF-8?q?=E5=AF=86=E7=A0=81=E7=99=BB=E5=BD=95=E7=9A=84=E9=94=99=E8=AF=AF?=
 =?UTF-8?q?=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 api/login.go                               | 13 ++++++++-
 api/session.go                             | 18 ++++++++----
 modules/chatgpt/service/chatgpt_session.go | 32 ++++++++++++++++++++++
 modules/chatgpt/tasks/refresh_session.go   | 12 ++++----
 4 files changed, 61 insertions(+), 14 deletions(-)

diff --git a/api/login.go b/api/login.go
index 4001889..d9e8729 100644
--- a/api/login.go
+++ b/api/login.go
@@ -77,6 +77,7 @@ func LoginPost(r *ghttp.Request) {
 			return
 		}
 		officialSession := record["officialSession"].String()
+
 		r.Session.Set("offical-session", officialSession)
 		r.Session.Set("userToken", userToken)
 		r.Response.RedirectTo("/")
@@ -149,8 +150,18 @@ func LoginPost(r *ghttp.Request) {
 			})
 		}
 	}
-
 	officialSession := record2["officialSession"].String()
+	if officialSession == "" && record2["mode"].Int() == 0 {
+		err = ChatgptSessionService.RefreshSession(ctx, record2)
+		officialSession = record2["officialSession"].String()
+		if err != nil {
+			r.Response.WriteTpl("login.html", g.Map{
+				"username": r.Get("username").String(),
+				"error":    "后台登录失败，请联系管理员检查日志",
+			})
+			return
+		}
+	}
 	r.Session.Set("offical-session", officialSession)
 	r.Session.Set("userToken", user["userToken"].String())
 	r.Response.RedirectTo("/")
diff --git a/api/session.go b/api/session.go
index 25c28ef..e84c2ed 100644
--- a/api/session.go
+++ b/api/session.go
@@ -50,15 +50,21 @@ func Session(r *ghttp.Request) {
 
 		r.Response.WriteJsonExit(sessionJson)
 	} else {
+		refreshFlag := false
+		refreshCookie := ""
+		sessionJson := gjson.New("{}")
+		if record["officialSession"].String() == "" {
+			refreshFlag = true
+		} else {
+			sessionJson := gjson.New(record["officialSession"].String())
+			refreshCookie = sessionJson.Get("refreshCookie").String()
+			expires := sessionJson.Get("expires").Time()
+			refreshFlag = expires.Before(time.Now())
+		}
 
-		sessionJson := gjson.New(record["officialSession"].String())
 		getSessionUrl := config.CHATPROXY(ctx) + "/getsession"
-		refreshCookie := sessionJson.Get("refreshCookie").String()
-
-		expires := sessionJson.Get("expires").Time()
-
 		// 判断是否过期
-		if expires.Before(time.Now()) {
+		if refreshFlag {
 			g.Log().Info(ctx, "session 过期，重新获取")
 			sessionVar := g.Client().SetHeader("authkey", config.AUTHKEY(ctx)).PostVar(ctx, getSessionUrl, g.Map{
 				"username":      record["email"].String(),
diff --git a/modules/chatgpt/service/chatgpt_session.go b/modules/chatgpt/service/chatgpt_session.go
index 4c6ff2b..6583022 100644
--- a/modules/chatgpt/service/chatgpt_session.go
+++ b/modules/chatgpt/service/chatgpt_session.go
@@ -90,6 +90,31 @@ func (s *ChatgptSessionService) ModifyAfter(ctx g.Ctx, method string, param map[
 
 }
 
+// 写一个函数 刷新数据
+func (s *ChatgptSessionService) RefreshSession(ctx g.Ctx, chatgpt_session gdb.Record) (err error) {
+	getSessionUrl := config.CHATPROXY(ctx) + "/getsession"
+	sessionVar := g.Client().SetHeader("authkey", config.AUTHKEY(ctx)).PostVar(ctx, getSessionUrl, g.Map{
+		"username": chatgpt_session["email"],
+		"password": chatgpt_session["password"],
+		"authkey":  config.AUTHKEY(ctx),
+	})
+	sessionJson := gjson.New(sessionVar)
+	models := sessionJson.Get("models").Array()
+	_, err = cool.DBM(s.Model).Where("email=?", chatgpt_session["email"]).Update(g.Map{
+		"officialSession": sessionJson.String(),
+		"isPlus":          len(models) > 1,
+		"status":          1,
+	})
+	chatgpt_session["officialSession"] = (g.NewVar(sessionJson.String()))
+	if err == nil {
+		g.Log().Info(ctx, "GPT账号登录成功", "邮箱：", chatgpt_session["email"])
+		return
+	}
+	g.Log().Error(ctx, err)
+	err = gerror.New("后台账号登录失败，检查登录日志")
+	return
+}
+
 // GetSessionByUserToken 根据userToken获取session
 func (s *ChatgptSessionService) GetSessionByUserToken(ctx g.Ctx, userToken string) (record gdb.Record, expireTime string, err error) {
 
@@ -112,6 +137,13 @@ func (s *ChatgptSessionService) GetSessionByUserToken(ctx g.Ctx, userToken strin
 		return nil, "", gerror.New("没有可用的ChatGpt账号,请联系管理员")
 	}
 
+	officialSession := record["officialSession"].String()
+	if officialSession == "" && record["mode"].Int() == 0 {
+		// 没有登录尝试登录
+		err = s.RefreshSession(ctx, record)
+
+	}
+
 	record["user_username"] = user["username"]
 	return
 }
diff --git a/modules/chatgpt/tasks/refresh_session.go b/modules/chatgpt/tasks/refresh_session.go
index acb4096..09b865b 100644
--- a/modules/chatgpt/tasks/refresh_session.go
+++ b/modules/chatgpt/tasks/refresh_session.go
@@ -38,14 +38,12 @@ func RefreshSession(ctx g.Ctx) {
 
 		getSessionUrl := config.CHATPROXY(ctx) + "/getsession"
 		refreshCookie := gjson.New(v["officialSession"]).Get("refreshCookie").String()
-		if refreshCookie == "" {
-			continue
-		}
+
 		sessionVar := g.Client().SetHeader("authkey", config.AUTHKEY(ctx)).PostVar(ctx, getSessionUrl, g.Map{
-			"username": v["email"],
-			"password": v["password"],
-			"authkey":  config.AUTHKEY(ctx),
-			// "refreshCookie": refreshCookie,
+			"username":      v["email"],
+			"password":      v["password"],
+			"authkey":       config.AUTHKEY(ctx),
+			"refreshCookie": refreshCookie,
 		})
 		g.Log().Info(ctx, "账号", v["email"], "sessionVar", sessionVar)
 		sessionJson := gjson.New(sessionVar)