refact & update new model
This commit is contained in:
Sakurasan
182
pkg/team/key.go
Normal file
182
pkg/team/key.go
Normal file
@@ -0,0 +1,182 @@
|
||||
package team
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"opencatd-open/pkg/azureopenai"
|
||||
"opencatd-open/store"
|
||||
"strings"
|
||||
|
||||
"github.com/Sakurasan/to"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type Key struct {
|
||||
ID int `json:"id,omitempty"`
|
||||
Key string `json:"key,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
ApiType string `json:"api_type,omitempty"`
|
||||
Endpoint string `json:"endpoint,omitempty"`
|
||||
UpdatedAt string `json:"updatedAt,omitempty"`
|
||||
CreatedAt string `json:"createdAt,omitempty"`
|
||||
}
|
||||
|
||||
func HandleKeys(c *gin.Context) {
|
||||
keys, err := store.GetAllKeys()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, keys)
|
||||
}
|
||||
|
||||
func HandleAddKey(c *gin.Context) {
|
||||
var body Key
|
||||
if err := c.BindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{
|
||||
"message": err.Error(),
|
||||
}})
|
||||
return
|
||||
}
|
||||
body.Name = strings.ToLower(strings.TrimSpace(body.Name))
|
||||
body.Key = strings.TrimSpace(body.Key)
|
||||
if strings.HasPrefix(body.Name, "azure.") {
|
||||
keynames := strings.Split(body.Name, ".")
|
||||
if len(keynames) < 2 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{
|
||||
"message": "Invalid Key Name",
|
||||
}})
|
||||
return
|
||||
}
|
||||
k := &store.Key{
|
||||
ApiType: "azure",
|
||||
Name: body.Name,
|
||||
Key: body.Key,
|
||||
ResourceNmae: keynames[1],
|
||||
EndPoint: body.Endpoint,
|
||||
}
|
||||
if err := store.CreateKey(k); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": gin.H{
|
||||
"message": err.Error(),
|
||||
}})
|
||||
return
|
||||
}
|
||||
} else if strings.HasPrefix(body.Name, "claude.") {
|
||||
keynames := strings.Split(body.Name, ".")
|
||||
if len(keynames) < 2 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{
|
||||
"message": "Invalid Key Name",
|
||||
}})
|
||||
return
|
||||
}
|
||||
if body.Endpoint == "" {
|
||||
body.Endpoint = "https://api.anthropic.com"
|
||||
}
|
||||
k := &store.Key{
|
||||
// ApiType: "anthropic",
|
||||
ApiType: "claude",
|
||||
Name: body.Name,
|
||||
Key: body.Key,
|
||||
ResourceNmae: keynames[1],
|
||||
EndPoint: body.Endpoint,
|
||||
}
|
||||
if err := store.CreateKey(k); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": gin.H{
|
||||
"message": err.Error(),
|
||||
}})
|
||||
return
|
||||
}
|
||||
} else if strings.HasPrefix(body.Name, "google.") {
|
||||
keynames := strings.Split(body.Name, ".")
|
||||
if len(keynames) < 2 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{
|
||||
"message": "Invalid Key Name",
|
||||
}})
|
||||
return
|
||||
}
|
||||
|
||||
k := &store.Key{
|
||||
// ApiType: "anthropic",
|
||||
ApiType: "google",
|
||||
Name: body.Name,
|
||||
Key: body.Key,
|
||||
ResourceNmae: keynames[1],
|
||||
EndPoint: body.Endpoint,
|
||||
}
|
||||
if err := store.CreateKey(k); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": gin.H{
|
||||
"message": err.Error(),
|
||||
}})
|
||||
return
|
||||
}
|
||||
} else if strings.HasPrefix(body.Name, "github.") {
|
||||
keynames := strings.Split(body.Name, ".")
|
||||
if len(keynames) < 2 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{
|
||||
"message": "Invalid Key Name",
|
||||
}})
|
||||
return
|
||||
}
|
||||
|
||||
k := &store.Key{
|
||||
ApiType: "github",
|
||||
Name: body.Name,
|
||||
Key: body.Key,
|
||||
ResourceNmae: keynames[1],
|
||||
EndPoint: body.Endpoint,
|
||||
}
|
||||
if err := store.CreateKey(k); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": gin.H{
|
||||
"message": err.Error(),
|
||||
}})
|
||||
return
|
||||
}
|
||||
} else {
|
||||
if body.ApiType == "" {
|
||||
if err := store.AddKey("openai", body.Key, body.Name); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": gin.H{
|
||||
"message": err.Error(),
|
||||
}})
|
||||
return
|
||||
}
|
||||
} else {
|
||||
k := &store.Key{
|
||||
ApiType: body.ApiType,
|
||||
Name: body.Name,
|
||||
Key: body.Key,
|
||||
ResourceNmae: azureopenai.GetResourceName(body.Endpoint),
|
||||
EndPoint: body.Endpoint,
|
||||
}
|
||||
if err := store.CreateKey(k); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": gin.H{
|
||||
"message": err.Error(),
|
||||
}})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
k, err := store.GetKeyrByName(body.Name)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": gin.H{
|
||||
"message": err.Error(),
|
||||
}})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, k)
|
||||
}
|
||||
|
||||
func HandleDelKey(c *gin.Context) {
|
||||
id := to.Int(c.Param("id"))
|
||||
if id < 1 {
|
||||
c.JSON(http.StatusOK, gin.H{"error": "invalid key id"})
|
||||
return
|
||||
}
|
||||
if err := store.DeleteKey(uint(id)); err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{"error": "invalid key id"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "ok"})
|
||||
}
|
||||
104
pkg/team/me.go
Normal file
104
pkg/team/me.go
Normal file
@@ -0,0 +1,104 @@
|
||||
package team
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"opencatd-open/store"
|
||||
"time"
|
||||
|
||||
"github.com/Sakurasan/to"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func Handleinit(c *gin.Context) {
|
||||
user, err := store.GetUserByID(1)
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
u := store.User{Name: "root", Token: uuid.NewString()}
|
||||
u.ID = 1
|
||||
if err := store.CreateUser(&u); err != nil {
|
||||
c.JSON(http.StatusForbidden, gin.H{
|
||||
"error": err.Error(),
|
||||
})
|
||||
return
|
||||
} else {
|
||||
rootToken = u.Token
|
||||
resJSON := User{
|
||||
false,
|
||||
int(u.ID),
|
||||
u.UpdatedAt.Format(time.RFC3339),
|
||||
u.Name,
|
||||
u.Token,
|
||||
u.CreatedAt.Format(time.RFC3339),
|
||||
}
|
||||
c.JSON(http.StatusOK, resJSON)
|
||||
return
|
||||
}
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"error": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
if user.ID == uint(1) {
|
||||
c.JSON(http.StatusForbidden, gin.H{
|
||||
"error": "super user already exists, use cli to reset password",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func HandleMe(c *gin.Context) {
|
||||
token := c.GetHeader("Authorization")
|
||||
u, err := store.GetUserByToken(token[7:])
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
|
||||
resJSON := User{
|
||||
false,
|
||||
int(u.ID),
|
||||
u.UpdatedAt.Format(time.RFC3339),
|
||||
u.Name,
|
||||
u.Token,
|
||||
u.CreatedAt.Format(time.RFC3339),
|
||||
}
|
||||
c.JSON(http.StatusOK, resJSON)
|
||||
}
|
||||
|
||||
func HandleMeUsage(c *gin.Context) {
|
||||
token := c.GetHeader("Authorization")
|
||||
fromStr := c.Query("from")
|
||||
toStr := c.Query("to")
|
||||
getMonthStartAndEnd := func() (start, end string) {
|
||||
loc, _ := time.LoadLocation("Local")
|
||||
now := time.Now().In(loc)
|
||||
|
||||
year, month, _ := now.Date()
|
||||
|
||||
startOfMonth := time.Date(year, month, 1, 0, 0, 0, 0, loc)
|
||||
endOfMonth := startOfMonth.AddDate(0, 1, 0)
|
||||
|
||||
start = startOfMonth.Format("2006-01-02")
|
||||
end = endOfMonth.Format("2006-01-02")
|
||||
return
|
||||
}
|
||||
if fromStr == "" || toStr == "" {
|
||||
fromStr, toStr = getMonthStartAndEnd()
|
||||
}
|
||||
user, err := store.GetUserByToken(token)
|
||||
if err != nil {
|
||||
c.AbortWithError(http.StatusForbidden, err)
|
||||
return
|
||||
}
|
||||
usage, err := store.QueryUserUsage(to.String(user.ID), fromStr, toStr)
|
||||
if err != nil {
|
||||
c.AbortWithError(http.StatusForbidden, err)
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(200, usage)
|
||||
}
|
||||
59
pkg/team/middleware.go
Normal file
59
pkg/team/middleware.go
Normal file
@@ -0,0 +1,59 @@
|
||||
package team
|
||||
|
||||
import (
|
||||
"log"
|
||||
"net/http"
|
||||
"opencatd-open/store"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
var (
|
||||
rootToken string
|
||||
)
|
||||
|
||||
func AuthMiddleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if rootToken == "" {
|
||||
u, err := store.GetUserByID(uint(1))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized"})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
rootToken = u.Token
|
||||
}
|
||||
token := c.GetHeader("Authorization")
|
||||
if token == "" || token[:7] != "Bearer " {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized"})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if store.IsExistAuthCache(token[7:]) {
|
||||
if strings.HasPrefix(c.Request.URL.Path, "/1/me") {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
}
|
||||
if token[7:] != rootToken {
|
||||
u, err := store.GetUserByID(uint(1))
|
||||
if err != nil {
|
||||
log.Println(err)
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized"})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if token[:7] != u.Token {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized"})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
rootToken = u.Token
|
||||
store.LoadAuthCache()
|
||||
}
|
||||
// 可以在这里对 token 进行验证并检查权限
|
||||
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
38
pkg/team/usage.go
Normal file
38
pkg/team/usage.go
Normal file
@@ -0,0 +1,38 @@
|
||||
package team
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"opencatd-open/store"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func HandleUsage(c *gin.Context) {
|
||||
fromStr := c.Query("from")
|
||||
toStr := c.Query("to")
|
||||
getMonthStartAndEnd := func() (start, end string) {
|
||||
loc, _ := time.LoadLocation("Local")
|
||||
now := time.Now().In(loc)
|
||||
|
||||
year, month, _ := now.Date()
|
||||
|
||||
startOfMonth := time.Date(year, month, 1, 0, 0, 0, 0, loc)
|
||||
endOfMonth := startOfMonth.AddDate(0, 1, 0)
|
||||
|
||||
start = startOfMonth.Format("2006-01-02")
|
||||
end = endOfMonth.Format("2006-01-02")
|
||||
return
|
||||
}
|
||||
if fromStr == "" || toStr == "" {
|
||||
fromStr, toStr = getMonthStartAndEnd()
|
||||
}
|
||||
|
||||
usage, err := store.QueryUsage(fromStr, toStr)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(200, usage)
|
||||
}
|
||||
89
pkg/team/user.go
Normal file
89
pkg/team/user.go
Normal file
@@ -0,0 +1,89 @@
|
||||
package team
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"opencatd-open/store"
|
||||
|
||||
"github.com/Sakurasan/to"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
type User struct {
|
||||
IsDelete bool `json:"IsDelete,omitempty"`
|
||||
ID int `json:"id,omitempty"`
|
||||
UpdatedAt string `json:"updatedAt,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Token string `json:"token,omitempty"`
|
||||
CreatedAt string `json:"createdAt,omitempty"`
|
||||
}
|
||||
|
||||
func HandleUsers(c *gin.Context) {
|
||||
users, err := store.GetAllUsers()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, users)
|
||||
}
|
||||
|
||||
func HandleAddUser(c *gin.Context) {
|
||||
var body User
|
||||
if err := c.BindJSON(&body); err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if len(body.Name) == 0 {
|
||||
c.JSON(http.StatusOK, gin.H{"error": "invalid user name"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := store.AddUser(body.Name, uuid.NewString()); err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
u, err := store.GetUserByName(body.Name)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, u)
|
||||
}
|
||||
|
||||
func HandleDelUser(c *gin.Context) {
|
||||
id := to.Int(c.Param("id"))
|
||||
if id <= 1 {
|
||||
c.JSON(http.StatusOK, gin.H{"error": "invalid user id"})
|
||||
return
|
||||
}
|
||||
if err := store.DeleteUser(uint(id)); err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "ok"})
|
||||
}
|
||||
|
||||
func HandleResetUserToken(c *gin.Context) {
|
||||
id := to.Int(c.Param("id"))
|
||||
newtoken := c.Query("token")
|
||||
if newtoken == "" {
|
||||
newtoken = uuid.NewString()
|
||||
}
|
||||
|
||||
if err := store.UpdateUser(uint(id), newtoken); err != nil {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
u, err := store.GetUserByID(uint(id))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if u.ID == uint(1) {
|
||||
rootToken = u.Token
|
||||
}
|
||||
c.JSON(http.StatusOK, u)
|
||||
}
|
||||
Reference in New Issue
Block a user