MailUI4Agents/gateway/internal/handler/auth.go

package handler

import (
	"errors"
	"net/http"
	"strings"
	"time"

	"github.com/agentmail/gateway/internal/middleware"
	"github.com/agentmail/gateway/internal/models"
	"github.com/agentmail/gateway/internal/repo"
	"github.com/go-chi/chi/v5"
	"github.com/google/uuid"
)

// ---------- 登录 / 登出 / 自身信息 ----------

type loginRequest struct {
	Username string `json:"username"`
	Password string `json:"password"`
}

type userOut struct {
	UserID        string   `json:"user_id"`
	Username      string   `json:"username"`
	DisplayName   string   `json:"display_name"`
	Role          string   `json:"role"`
	Status        string   `json:"status"`
	AllowedAgents []string `json:"allowed_agents"`
	AllowedPaths  []string `json:"allowed_paths"`
	LastLogin     string   `json:"last_login,omitempty"`
	CreatedAt     string   `json:"created_at,omitempty"`
}

func toUserOut(u *models.User) userOut {
	o := userOut{
		UserID:        u.ID.String(),
		Username:      u.Username,
		DisplayName:   u.DisplayName,
		Role:          u.Role,
		Status:        u.Status,
		AllowedAgents: emptySlice(u.AllowedAgents),
		AllowedPaths:  emptySlice(u.AllowedPaths),
		CreatedAt:     u.CreatedAt.Format("2006-01-02 15:04:05"),
	}
	if u.LastLogin != nil {
		o.LastLogin = u.LastLogin.Format("2006-01-02 15:04:05")
	}
	return o
}

// ---------- 首次初始化 ----------

// GET /api/v1/setup/status —— 公开：前端据此判断是否展示初始化向导
func SetupStatus(w http.ResponseWriter, r *http.Request) {
	needs, err := repo.NeedsSetup(r.Context())
	if err != nil {
		Error(w, http.StatusInternalServerError, "Failed to check setup status")
		return
	}
	JSON(w, http.StatusOK, map[string]bool{"needs_setup": needs})
}

type setupRequest struct {
	Username    string `json:"username"`
	Password    string `json:"password"`
	DisplayName string `json:"display_name"`
}

// POST /api/v1/setup/admin —— 公开，但仅在系统无任何用户时可用
func SetupAdmin(w http.ResponseWriter, r *http.Request) {
	var req setupRequest
	if !DecodeBody(w, r, &req) {
		return
	}
	if len(req.Password) < 8 {
		Error(w, http.StatusBadRequest, "密码自少 8 位")
		return
	}

	u, err := repo.SetupFirstAdmin(r.Context(), req.Username, req.Password, req.DisplayName)
	if err != nil {
		switch {
		case errors.Is(err, repo.ErrAlreadySetup):
			Error(w, http.StatusConflict, "系统已初始化，请直接登录")
		case errors.Is(err, repo.ErrInvalidUsername):
			Error(w, http.StatusBadRequest, "用户名只能是 2-64 位的小写字母、数字、点、下划线、连字符，且不能为 human")
		case errors.Is(err, repo.ErrNameTaken):
			Error(w, http.StatusConflict, "该名称已被 Agent 占用")
		default:
			Error(w, http.StatusInternalServerError, "初始化失败")
		}
		return
	}

	// 初始化后直接登录
	token, expires, err := repo.CreateUserSession(r.Context(), u.ID, r.UserAgent())
	if err == nil {
		maxAge := int(time.Until(expires).Seconds())
		if maxAge < 0 {
			maxAge = 0
		}
		middleware.SetSessionCookie(w, token, maxAge)
	}

	JSON(w, http.StatusOK, map[string]interface{}{"user": toUserOut(u)})
}

// POST /api/v1/auth/login
func Login(w http.ResponseWriter, r *http.Request) {
	var req loginRequest
	if !DecodeBody(w, r, &req) {
		return
	}
	name := strings.ToLower(strings.TrimSpace(req.Username))
	if name == "" || req.Password == "" {
		Error(w, http.StatusBadRequest, "Missing username or password")
		return
	}

	if locked, remain := limiter.Locked(r.Context(), name); locked {
		JSON(w, http.StatusTooManyRequests, map[string]interface{}{
			"error":       "尝试过于频繁，请稍后再试",
			"retry_after": remain,
		})
		return
	}

	u, err := repo.Authenticate(r.Context(), name, req.Password)
	if err != nil {
		switch {
		case errors.Is(err, repo.ErrBadCredentials):
			limiter.Fail(r.Context(), name)
			Error(w, http.StatusUnauthorized, "用户名或密码错误")
		case errors.Is(err, repo.ErrUserDisabled):
			Error(w, http.StatusForbidden, "账号已被禁用")
		default:
			Error(w, http.StatusInternalServerError, "登录失败")
		}
		return
	}
	limiter.Reset(r.Context(), name)

	token, expires, err := repo.CreateUserSession(r.Context(), u.ID, r.UserAgent())
	if err != nil {
		Error(w, http.StatusInternalServerError, "无法创建会话")
		return
	}
	maxAge := int(time.Until(expires).Seconds())
	if maxAge < 0 {
		maxAge = 0
	}
	middleware.SetSessionCookie(w, token, maxAge)

	JSON(w, http.StatusOK, map[string]interface{}{
		"user": toUserOut(u),
	})
}

// POST /api/v1/auth/logout
func Logout(w http.ResponseWriter, r *http.Request) {
	if token := middleware.SessionToken(r); token != "" {
		_ = repo.DeleteUserSession(r.Context(), token)
	}
	middleware.ClearSessionCookie(w)
	JSON(w, http.StatusOK, map[string]string{"status": "logged_out"})
}

// GET /api/v1/auth/me
func Me(w http.ResponseWriter, r *http.Request) {
	u := middleware.GetUser(r)
	if u == nil {
		Error(w, http.StatusUnauthorized, "not authenticated")
		return
	}
	JSON(w, http.StatusOK, map[string]interface{}{"user": toUserOut(u)})
}

type changePasswordRequest struct {
	OldPassword string `json:"old_password"`
	NewPassword string `json:"new_password"`
}

// POST /api/v1/auth/password
func ChangePassword(w http.ResponseWriter, r *http.Request) {
	u := middleware.GetUser(r)
	if u == nil {
		Error(w, http.StatusUnauthorized, "not authenticated")
		return
	}

	var req changePasswordRequest
	if !DecodeBody(w, r, &req) {
		return
	}
	if len(req.NewPassword) < 8 {
		Error(w, http.StatusBadRequest, "新密码至少 8 位")
		return
	}
	if _, err := repo.Authenticate(r.Context(), u.Username, req.OldPassword); err != nil {
		Error(w, http.StatusUnauthorized, "原密码错误")
		return
	}
	if err := repo.SetPassword(r.Context(), u.ID, req.NewPassword); err != nil {
		Error(w, http.StatusInternalServerError, "修改密码失败")
		return
	}
	middleware.ClearSessionCookie(w)
	JSON(w, http.StatusOK, map[string]string{"status": "password_changed"})
}

// ---------- 管理员：用户管理 ----------

// GET /api/v1/admin/users
func AdminListUsers(w http.ResponseWriter, r *http.Request) {
	users, err := repo.ListUsers(r.Context())
	if err != nil {
		Error(w, http.StatusInternalServerError, "Failed to list users")
		return
	}
	out := make([]userOut, 0, len(users))
	for i := range users {
		out = append(out, toUserOut(&users[i]))
	}
	JSON(w, http.StatusOK, map[string]interface{}{"users": out})
}

type createUserRequest struct {
	Username      string   `json:"username"`
	Password      string   `json:"password"`
	DisplayName   string   `json:"display_name"`
	Role          string   `json:"role"`
	AllowedAgents []string `json:"allowed_agents"`
	AllowedPaths  []string `json:"allowed_paths"`
}

// POST /api/v1/admin/users
func AdminCreateUser(w http.ResponseWriter, r *http.Request) {
	var req createUserRequest
	if !DecodeBody(w, r, &req) {
		return
	}
	if len(req.Password) < 8 {
		Error(w, http.StatusBadRequest, "密码至少 8 位")
		return
	}

	u, err := repo.CreateUser(r.Context(), req.Username, req.Password, req.DisplayName, req.Role,
		req.AllowedAgents, req.AllowedPaths)
	if err != nil {
		switch {
		case errors.Is(err, repo.ErrNameTaken):
			Error(w, http.StatusConflict, "该名称已被用户或 Agent 占用")
		case errors.Is(err, repo.ErrInvalidUsername):
			Error(w, http.StatusBadRequest, "用户名只能是 2-64 位的小写字母、数字、点、下划线、连字符，且不能为 human")
		default:
			Error(w, http.StatusInternalServerError, "创建用户失败")
		}
		return
	}
	JSON(w, http.StatusOK, map[string]interface{}{"user": toUserOut(u)})
}

type updateUserRequest struct {
	DisplayName   *string   `json:"display_name"`
	Role          *string   `json:"role"`
	Status        *string   `json:"status"`
	AllowedAgents *[]string `json:"allowed_agents"`
	AllowedPaths  *[]string `json:"allowed_paths"`
}

// PUT /api/v1/admin/users/{id}
func AdminUpdateUser(w http.ResponseWriter, r *http.Request) {
	id, ok := pathUUID(w, r, "id")
	if !ok {
		return
	}
	var req updateUserRequest
	if !DecodeBody(w, r, &req) {
		return
	}

	// 不允许把最后一个管理员降级或禁用
	if err := guardLastAdmin(r, id, req.Role, req.Status); err != nil {
		Error(w, http.StatusConflict, err.Error())
		return
	}

	u, err := repo.UpdateUser(r.Context(), id, repo.UserUpdate{
		DisplayName:   req.DisplayName,
		Role:          req.Role,
		Status:        req.Status,
		AllowedAgents: req.AllowedAgents,
		AllowedPaths:  req.AllowedPaths,
	})
	if err != nil {
		if errors.Is(err, repo.ErrUserNotFound) {
			Error(w, http.StatusNotFound, "用户不存在")
			return
		}
		Error(w, http.StatusInternalServerError, "更新用户失败")
		return
	}
	JSON(w, http.StatusOK, map[string]interface{}{"user": toUserOut(u)})
}

// GET /api/v1/admin/scopes —— 可授权的 Agent 与目录候选
func AdminListScopes(w http.ResponseWriter, r *http.Request) {
	agents, err := repo.ListAgents(r.Context(), "")
	if err != nil {
		Error(w, http.StatusInternalServerError, "Failed to list agents")
		return
	}
	names := make([]string, 0, len(agents))
	for _, a := range agents {
		names = append(names, a.Name)
	}
	paths, _ := repo.AllWorkspaceNames(r.Context())

	JSON(w, http.StatusOK, map[string]interface{}{
		"agents": emptySlice(names),
		"paths":  emptySlice(paths),
	})
}

// DELETE /api/v1/admin/users/{id} —— 禁用而非物理删除，保留邮件历史
func AdminDisableUser(w http.ResponseWriter, r *http.Request) {
	id, ok := pathUUID(w, r, "id")
	if !ok {
		return
	}
	disabled := "disabled"
	if err := guardLastAdmin(r, id, nil, &disabled); err != nil {
		Error(w, http.StatusConflict, err.Error())
		return
	}
	if err := repo.DisableUser(r.Context(), id); err != nil {
		if errors.Is(err, repo.ErrUserNotFound) {
			Error(w, http.StatusNotFound, "用户不存在")
			return
		}
		Error(w, http.StatusInternalServerError, "禁用用户失败")
		return
	}
	JSON(w, http.StatusOK, map[string]string{"status": "disabled"})
}

type resetPasswordRequest struct {
	NewPassword string `json:"new_password"`
}

// POST /api/v1/admin/users/{id}/reset
func AdminResetPassword(w http.ResponseWriter, r *http.Request) {
	id, ok := pathUUID(w, r, "id")
	if !ok {
		return
	}
	var req resetPasswordRequest
	if !DecodeBody(w, r, &req) {
		return
	}
	if len(req.NewPassword) < 8 {
		Error(w, http.StatusBadRequest, "密码至少 8 位")
		return
	}
	if err := repo.SetPassword(r.Context(), id, req.NewPassword); err != nil {
		if errors.Is(err, repo.ErrUserNotFound) {
			Error(w, http.StatusNotFound, "用户不存在")
			return
		}
		Error(w, http.StatusInternalServerError, "重置密码失败")
		return
	}
	JSON(w, http.StatusOK, map[string]string{"status": "password_reset"})
}

// ---------- 辅助 ----------

func pathUUID(w http.ResponseWriter, r *http.Request, key string) (uuid.UUID, bool) {
	id, err := uuid.Parse(chi.URLParam(r, key))
	if err != nil {
		Error(w, http.StatusBadRequest, "Invalid "+key)
		return uuid.Nil, false
	}
	return id, true
}

// guardLastAdmin 阻止把系统里最后一个可用管理员降级或禁用
func guardLastAdmin(r *http.Request, id uuid.UUID, role, status *string) error {
	demoting := role != nil && *role != "admin"
	disabling := status != nil && *status != "active"
	if !demoting && !disabling {
		return nil
	}

	target, err := repo.GetUserByID(r.Context(), id)
	if err != nil || !target.IsAdmin() || target.Status != "active" {
		return nil
	}
	n, err := repo.CountAdmins(r.Context())
	if err != nil {
		return nil
	}
	if n <= 1 {
		return errors.New("系统至少需要保留一个可用管理员")
	}
	return nil
}