iqudo-top1/internal/auth/service.go

package auth

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"git.iwork-ai.com/xdc/iqudo-top1/internal/config"
	"git.iwork-ai.com/xdc/iqudo-top1/internal/model"
	"github.com/golang-jwt/jwt/v5"
	"github.com/google/uuid"
	"golang.org/x/crypto/bcrypt"
	"gorm.io/datatypes"
	"gorm.io/gorm"
)

var ErrInvalidCredentials = errors.New("invalid username or password")

type Principal struct {
	UserID      uint64   `json:"user_id"`
	TenantID    uint64   `json:"tenant_id"`
	RoleCode    string   `json:"role_code"`
	Permissions []string `json:"permissions"`
	Username    string   `json:"username"`
	DisplayName string   `json:"display_name"`
}

type Claims struct {
	TenantID    uint64 `json:"tenant_id"`
	RoleCode    string `json:"role_code"`
	Username    string `json:"username"`
	DisplayName string `json:"display_name"`
	jwt.RegisteredClaims
}

type TokenPair struct {
	AccessToken  string    `json:"access_token"`
	RefreshToken string    `json:"refresh_token"`
	ExpiresAt    time.Time `json:"expires_at"`
	User         Principal `json:"user"`
}

type Service struct {
	db  *gorm.DB
	cfg config.AuthConfig
}

func NewService(db *gorm.DB, cfg config.AuthConfig) *Service {
	return &Service{db: db, cfg: cfg}
}

func (s *Service) Login(username, password string) (TokenPair, error) {
	var user model.User
	if err := s.db.Where("username = ? AND status = ?", username, "active").First(&user).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return TokenPair{}, ErrInvalidCredentials
		}
		return TokenPair{}, err
	}
	if err := bcrypt.CompareHashAndPassword([]byte(user.PasswordHash), []byte(password)); err != nil {
		return TokenPair{}, ErrInvalidCredentials
	}
	principal, err := s.principalForUser(user)
	if err != nil {
		return TokenPair{}, err
	}
	return s.issueTokenPair(principal)
}

func (s *Service) Refresh(rawToken string) (TokenPair, error) {
	hash := tokenHash(rawToken)
	var stored model.RefreshToken
	err := s.db.Where("token_hash = ? AND revoked_at IS NULL AND expires_at > ?", hash, time.Now()).First(&stored).Error
	if err != nil {
		return TokenPair{}, ErrInvalidCredentials
	}
	var user model.User
	if err := s.db.Where("id = ? AND status = ?", stored.UserID, "active").First(&user).Error; err != nil {
		return TokenPair{}, err
	}
	principal, err := s.principalForTenant(user, stored.TenantID)
	if err != nil {
		return TokenPair{}, err
	}
	now := time.Now()
	if err := s.db.Model(&stored).Update("revoked_at", &now).Error; err != nil {
		return TokenPair{}, err
	}
	return s.issueTokenPair(principal)
}

func (s *Service) ParseAccessToken(rawToken string) (Principal, error) {
	claims := &Claims{}
	token, err := jwt.ParseWithClaims(rawToken, claims, func(token *jwt.Token) (interface{}, error) {
		if token.Method.Alg() != jwt.SigningMethodHS256.Alg() {
			return nil, fmt.Errorf("unexpected signing method: %s", token.Method.Alg())
		}
		return []byte(s.cfg.JWTSecret), nil
	})
	if err != nil || !token.Valid {
		return Principal{}, ErrInvalidCredentials
	}
	userID, err := parseUint(claims.Subject)
	if err != nil {
		return Principal{}, ErrInvalidCredentials
	}
	return Principal{UserID: userID, TenantID: claims.TenantID, RoleCode: claims.RoleCode, Username: claims.Username, DisplayName: claims.DisplayName}, nil
}

func (s *Service) RefreshPrincipal(principal Principal) (Principal, error) {
	var user model.User
	if err := s.db.Where("id = ? AND status = ?", principal.UserID, "active").First(&user).Error; err != nil {
		return Principal{}, ErrInvalidCredentials
	}
	current, err := s.principalForTenant(user, principal.TenantID)
	if err != nil {
		return Principal{}, ErrInvalidCredentials
	}
	return current, nil
}

func (s *Service) issueTokenPair(principal Principal) (TokenPair, error) {
	now := time.Now()
	expiresAt := now.Add(s.cfg.AccessTokenTTL)
	claims := Claims{
		TenantID: principal.TenantID, RoleCode: principal.RoleCode, Username: principal.Username, DisplayName: principal.DisplayName,
		RegisteredClaims: jwt.RegisteredClaims{Subject: fmt.Sprintf("%d", principal.UserID), IssuedAt: jwt.NewNumericDate(now), ExpiresAt: jwt.NewNumericDate(expiresAt)},
	}
	access, err := jwt.NewWithClaims(jwt.SigningMethodHS256, claims).SignedString([]byte(s.cfg.JWTSecret))
	if err != nil {
		return TokenPair{}, err
	}
	refresh := uuid.NewString() + uuid.NewString()
	record := model.RefreshToken{TenantID: principal.TenantID, UserID: principal.UserID, TokenHash: tokenHash(refresh), ExpiresAt: now.Add(s.cfg.RefreshTokenTTL)}
	if err := s.db.Create(&record).Error; err != nil {
		return TokenPair{}, err
	}
	return TokenPair{AccessToken: access, RefreshToken: refresh, ExpiresAt: expiresAt, User: principal}, nil
}

func (s *Service) principalForUser(user model.User) (Principal, error) {
	return s.principalForTenant(user, 0)
}

func (s *Service) principalForTenant(user model.User, tenantID uint64) (Principal, error) {
	type row struct {
		TenantID    uint64
		RoleCode    string
		Permissions datatypes.JSON
	}
	var membership row
	query := s.db.Table("tenant_members tm").Select("tm.tenant_id, r.code AS role_code, r.permissions").Joins("JOIN roles r ON r.id = tm.role_id").Where("tm.user_id = ? AND tm.status = ?", user.ID, "active")
	if tenantID != 0 {
		query = query.Where("tm.tenant_id = ?", tenantID)
	}
	err := query.Order("tm.id").First(&membership).Error
	if err != nil {
		return Principal{}, err
	}
	var permissions []string
	if err := json.Unmarshal(membership.Permissions, &permissions); err != nil {
		return Principal{}, fmt.Errorf("decode role permissions: %w", err)
	}
	return Principal{UserID: user.ID, TenantID: membership.TenantID, RoleCode: membership.RoleCode, Permissions: permissions, Username: user.Username, DisplayName: user.DisplayName}, nil
}

func Seed(db *gorm.DB, cfg config.SeedConfig) error {
	if cfg.AdminUsername == "" || cfg.AdminPassword == "" {
		return nil
	}
	return db.Transaction(func(tx *gorm.DB) error {
		var tenant model.Tenant
		if err := tx.Where("slug = ?", "default").FirstOrCreate(&tenant, model.Tenant{Name: "默认企业", Slug: "default", Status: "active"}).Error; err != nil {
			return err
		}
		roleDefinitions := []struct {
			Name, Code  string
			Permissions []string
		}{
			{Name: "管理员", Code: "admin", Permissions: []string{"*"}},
			{Name: "SOP 编辑者", Code: "editor", Permissions: []string{"dashboard.view", "scenario.view", "scenario.team_view", "scenario.edit", "sop.view", "sop.edit", "scriptkit.view", "scriptkit.edit", "runs.view_all"}},
			{Name: "一线执行者", Code: "operator", Permissions: []string{"dashboard.view", "scenario.view", "sop.execute", "scriptkit.view", "runs.view_own", "runs.feedback"}},
		}
		roles := make(map[string]model.Role, len(roleDefinitions))
		for _, definition := range roleDefinitions {
			permissions, _ := json.Marshal(definition.Permissions)
			var role model.Role
			if err := tx.Where("tenant_id = ? AND code = ?", tenant.ID, definition.Code).Assign(model.Role{Name: definition.Name, Permissions: datatypes.JSON(permissions)}).FirstOrCreate(&role, model.Role{TenantID: tenant.ID, Code: definition.Code}).Error; err != nil {
				return err
			}
			roles[definition.Code] = role
		}
		var user model.User
		err := tx.Where("username = ?", cfg.AdminUsername).First(&user).Error
		if errors.Is(err, gorm.ErrRecordNotFound) {
			hash, hashErr := bcrypt.GenerateFromPassword([]byte(cfg.AdminPassword), bcrypt.DefaultCost)
			if hashErr != nil {
				return hashErr
			}
			user = model.User{Username: cfg.AdminUsername, PasswordHash: string(hash), DisplayName: cfg.AdminDisplayName, Status: "active"}
			if err := tx.Create(&user).Error; err != nil {
				return err
			}
		} else if err != nil {
			return err
		}
		var member model.TenantMember
		return tx.Where("tenant_id = ? AND user_id = ?", tenant.ID, user.ID).FirstOrCreate(&member, model.TenantMember{TenantID: tenant.ID, UserID: user.ID, RoleID: roles["admin"].ID, Status: "active"}).Error
	})
}

func tokenHash(value string) string {
	sum := sha256.Sum256([]byte(value))
	return hex.EncodeToString(sum[:])
}

func parseUint(value string) (uint64, error) {
	var result uint64
	_, err := fmt.Sscanf(value, "%d", &result)
	return result, err
}