package auth import ( "encoding/json" "fmt" "net/http" "strconv" "time" "github.com/gin-gonic/gin" "github.com/pkg/errors" "github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/ctxkey" "github.com/songquanpeng/one-api/controller" "github.com/songquanpeng/one-api/model" ) type wechatLoginResponse struct { Success bool `json:"success"` Message string `json:"message"` Data string `json:"data"` } func getWeChatIdByCode(code string) (string, error) { if code == "" { return "", errors.New("Invalid parameter") } req, err := http.NewRequest("GET", fmt.Sprintf("%s/api/wechat/user?code=%s", config.WeChatServerAddress, code), nil) if err != nil { return "", err } req.Header.Set("Authorization", config.WeChatServerToken) client := http.Client{ Timeout: 5 * time.Second, } httpResponse, err := client.Do(req) if err != nil { return "", err } defer httpResponse.Body.Close() var res wechatLoginResponse err = json.NewDecoder(httpResponse.Body).Decode(&res) if err != nil { return "", err } if !res.Success { return "", errors.New(res.Message) } if res.Data == "" { return "", errors.New("Verification code error or expired") } return res.Data, nil } func WeChatAuth(c *gin.Context) { ctx := c.Request.Context() if !config.WeChatAuthEnabled { c.JSON(http.StatusOK, gin.H{ "message": "The administrator has not enabled login and registration via WeChat", "success": false, }) return } code := c.Query("code") wechatId, err := getWeChatIdByCode(code) if err != nil { c.JSON(http.StatusOK, gin.H{ "message": err.Error(), "success": false, }) return } user := model.User{ WeChatId: wechatId, } if model.IsWeChatIdAlreadyTaken(wechatId) { err := user.FillUserByWeChatId() if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } } else { if config.RegisterEnabled { user.Username = "wechat_" + strconv.Itoa(model.GetMaxUserId()+1) user.DisplayName = "WeChat User" user.Role = model.RoleCommonUser user.Status = model.UserStatusEnabled if err := user.Insert(ctx, 0); err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } } else { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "The administrator has turned off new user registration", }) return } } if user.Status != model.UserStatusEnabled { c.JSON(http.StatusOK, gin.H{ "message": "User has been banned", "success": false, }) return } controller.SetupLogin(&user, c) } func WeChatBind(c *gin.Context) { if !config.WeChatAuthEnabled { c.JSON(http.StatusOK, gin.H{ "message": "The administrator has not enabled login and registration via WeChat", "success": false, }) return } code := c.Query("code") wechatId, err := getWeChatIdByCode(code) if err != nil { c.JSON(http.StatusOK, gin.H{ "message": err.Error(), "success": false, }) return } if model.IsWeChatIdAlreadyTaken(wechatId) { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "The WeChat account has been bound", }) return } id := c.GetInt(ctxkey.Id) user := model.User{ Id: id, } err = user.FillUserById() if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } user.WeChatId = wechatId err = user.Update(false) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), }) return } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", }) return }