144 lines
3.9 KiB
Go
144 lines
3.9 KiB
Go
package ai
|
|
|
|
import (
|
|
"crypto/aes"
|
|
"crypto/cipher"
|
|
"crypto/rand"
|
|
"crypto/sha256"
|
|
"encoding/base64"
|
|
"errors"
|
|
"io"
|
|
"time"
|
|
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/clause"
|
|
"senlinai-agent/backend/internal/domain"
|
|
)
|
|
|
|
type Gateway struct {
|
|
db *gorm.DB
|
|
systemKey string
|
|
encryptionSecret string
|
|
}
|
|
|
|
type SelectedKey struct {
|
|
Provider string
|
|
APIKey string
|
|
KeyType string
|
|
}
|
|
|
|
func NewGateway(database *gorm.DB, systemKey string) *Gateway {
|
|
return NewGatewayWithSecret(database, systemKey, "development-ai-key-secret-change-me")
|
|
}
|
|
|
|
func NewGatewayWithSecret(database *gorm.DB, systemKey string, encryptionSecret string) *Gateway {
|
|
return &Gateway{db: database, systemKey: systemKey, encryptionSecret: encryptionSecret}
|
|
}
|
|
|
|
func (g *Gateway) SaveUserKey(userID uint, provider string, apiKey string) error {
|
|
encrypted, err := encryptAPIKey(apiKey, g.encryptionSecret)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
key := domain.AIKey{UserID: userID, Provider: provider, EncryptedAPIKey: encrypted}
|
|
return g.db.Clauses(clause.OnConflict{
|
|
Columns: []clause.Column{{Name: "user_id"}},
|
|
DoUpdates: clause.AssignmentColumns([]string{"provider", "encrypted_api_key", "updated_at"}),
|
|
}).Create(&key).Error
|
|
}
|
|
|
|
func (g *Gateway) SelectKey(userID uint) (SelectedKey, error) {
|
|
var userKey domain.AIKey
|
|
if err := g.db.Where("user_id = ?", userID).First(&userKey).Error; err == nil {
|
|
apiKey, err := decryptAPIKey(userKey.EncryptedAPIKey, g.encryptionSecret)
|
|
if err != nil {
|
|
return SelectedKey{}, err
|
|
}
|
|
return SelectedKey{Provider: userKey.Provider, APIKey: apiKey, KeyType: "user"}, nil
|
|
}
|
|
if g.systemKey == "" {
|
|
return SelectedKey{}, errors.New("no ai key available")
|
|
}
|
|
return SelectedKey{Provider: "openai", APIKey: g.systemKey, KeyType: "system"}, nil
|
|
}
|
|
|
|
func (g *Gateway) RecordCall(userID uint, provider string, usedKeyType string, action string, status string, errText string) error {
|
|
return g.db.Create(&domain.AICallLog{
|
|
UserID: userID,
|
|
Provider: provider,
|
|
UsedKeyType: usedKeyType,
|
|
Action: action,
|
|
Status: status,
|
|
Error: errText,
|
|
}).Error
|
|
}
|
|
|
|
func (g *Gateway) CheckRateLimit(userID uint, action string, limit int, window time.Duration) error {
|
|
if limit <= 0 {
|
|
return nil
|
|
}
|
|
var count int64
|
|
if err := g.db.Model(&domain.AICallLog{}).
|
|
Where("user_id = ? AND action = ? AND created_at >= ?", userID, action, time.Now().Add(-window)).
|
|
Count(&count).Error; err != nil {
|
|
return err
|
|
}
|
|
if count >= int64(limit) {
|
|
return errors.New("ai rate limit exceeded")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func encryptAPIKey(apiKey string, secret string) (string, error) {
|
|
if apiKey == "" {
|
|
return "", errors.New("api key is required")
|
|
}
|
|
block, err := aes.NewCipher(encryptionKey(secret))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
gcm, err := cipher.NewGCM(block)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
nonce := make([]byte, gcm.NonceSize())
|
|
if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
|
|
return "", err
|
|
}
|
|
ciphertext := gcm.Seal(nil, nonce, []byte(apiKey), nil)
|
|
return "v1:" + base64.RawURLEncoding.EncodeToString(append(nonce, ciphertext...)), nil
|
|
}
|
|
|
|
func decryptAPIKey(encrypted string, secret string) (string, error) {
|
|
if len(encrypted) < 3 || encrypted[:3] != "v1:" {
|
|
return encrypted, nil
|
|
}
|
|
raw, err := base64.RawURLEncoding.DecodeString(encrypted[3:])
|
|
if err != nil {
|
|
return "", errors.New("invalid encrypted api key")
|
|
}
|
|
block, err := aes.NewCipher(encryptionKey(secret))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
gcm, err := cipher.NewGCM(block)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if len(raw) < gcm.NonceSize() {
|
|
return "", errors.New("invalid encrypted api key")
|
|
}
|
|
nonce := raw[:gcm.NonceSize()]
|
|
ciphertext := raw[gcm.NonceSize():]
|
|
plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
|
|
if err != nil {
|
|
return "", errors.New("invalid encrypted api key")
|
|
}
|
|
return string(plaintext), nil
|
|
}
|
|
|
|
func encryptionKey(secret string) []byte {
|
|
sum := sha256.Sum256([]byte(secret))
|
|
return sum[:]
|
|
}
|