MailUI4Agents/server/internal/repo/keys.go

package repo

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"time"

	"github.com/agentmail/gateway/internal/db"
	"github.com/agentmail/gateway/internal/models"
	"github.com/google/uuid"
)

// ---------- 密钥认证 ----------
//
// 两类密钥共享一个全局唯一的 token 命名空间：验证时先查 agent_keys 再查 user_keys。
// 这样一个 token 永远只有一种身份，不会出现「同一串字符既能注册 Agent 又能读人类邮箱」。

var (
	// ErrKeyNotFound 密钥不存在
	ErrKeyNotFound = errors.New("key not found")
	// ErrKeyUsed 一次性密钥已被使用
	ErrKeyUsed = errors.New("key already used")
	// ErrKeyExpired 定时密钥已过期
	ErrKeyExpired = errors.New("key expired")
	// ErrKeyTypeInvalid 密钥类型不受支持
	ErrKeyTypeInvalid = errors.New("invalid key type")
	// ErrKeyNeedsExpiry timed 密钥缺少有效的 expires_hours
	ErrKeyNeedsExpiry = errors.New("timed key requires positive expires_hours")
	// ErrKeyTooShort 登记的客户端密钥长度不足
	ErrKeyTooShort = errors.New("key token too short")
)

// expiryFor 依据密钥类型算出过期时间。
// 只有 timed 需要 expires_at；permanent 与 one_time 都是 NULL，
// 各自的失效条件由 key_type 本身表达，不混用 expires_at。
func expiryFor(keyType string, hours int) (*time.Time, error) {
	if !models.ValidKeyType(keyType) {
		return nil, ErrKeyTypeInvalid
	}
	if keyType != models.KeyTimed {
		return nil, nil
	}
	if hours <= 0 {
		return nil, ErrKeyNeedsExpiry
	}
	t := time.Now().Add(time.Duration(hours) * time.Hour)
	return &t, nil
}

// checkKeyUsable 判断一条密钥记录当前是否可用。
func checkKeyUsable(keyType string, expiresAt, usedAt *time.Time) error {
	switch keyType {
	case models.KeyOneTime:
		if usedAt != nil {
			return ErrKeyUsed
		}
	case models.KeyTimed:
		if expiresAt == nil || time.Now().After(*expiresAt) {
			return ErrKeyExpired
		}
	}
	return nil
}

// ---------- Agent 密钥（管理员签发） ----------

// ErrKeyTokenTaken 登记的密钥已被占用
var ErrKeyTokenTaken = errors.New("key token already registered")

// CreateAgentKey 签发一条 Agent 接入密钥。agentName 为空表示待绑定。
//
// presetToken 非空时登记客户端已在本地生成的密钥（插件首装场景），
// 这样密钥全文只从客户端往服务器走一次，不需要反方向传递；留空则由服务器生成。
func CreateAgentKey(ctx context.Context, agentName, keyType, label string, expiresHours int, createdBy uuid.UUID, presetToken string) (*models.AgentKey, error) {
	expires, err := expiryFor(keyType, expiresHours)
	if err != nil {
		return nil, err
	}

	token := presetToken
	if token == "" {
		if token, err = newToken(); err != nil {
			return nil, err
		}
	} else if len(token) < 32 {
		// 太短的客户端密钥不接受，否则等于把弱口令当凭证
		return nil, ErrKeyTooShort
	}

	// 已退役的名字不可重建 —— 登记密钥会级联建 agents 行，绕过注册检查。
	if agentName != "" {
		if retired, rErr := IsRetiredAgentName(ctx, agentName); rErr != nil {
			return nil, rErr
		} else if retired {
			return nil, fmt.Errorf("该名字已退役，不可重建")
		}
		// ★ 顺带确保 agents 行存在（2026-10-04 补）。
		//
		// 上面那句注释一直说「登记密钥会级联建 agents 行」，但实现里
		// **从来没有建过** —— 于是 admin 签出一把 key，MCP 客户端拿它连上来
		// 得到的是：
		//
		//     {"error":"密钥尚未绑定 Agent，请先调用 /agent/register 完成注册"}
		//
		// 而它**永远不会**去调 /agent/register —— MCP 宿主只按配置发
		// tools/list。所以「admin 建身份 → 填进 MCP 配置 → 用」这条
		// 路是断的，而断点没有任何报错指向 admin 侧。
		if err := ensureAgentRow(ctx, agentName); err != nil {
			return nil, err
		}
	}

	var namePtr *string
	if agentName != "" {
		namePtr = &agentName
	}

	k := &models.AgentKey{
		Token:     token,
		TokenHint: models.TokenHint(token),
		AgentName: namePtr,
		KeyType:   keyType,
		Label:     label,
		ExpiresAt: expires,
		CreatedBy: &createdBy,
	}
	err = db.DB.QueryRowContext(ctx, `
		INSERT INTO agent_keys (key_token, agent_name, key_type, label, expires_at, created_by)
		VALUES ($1, $2, $3, $4, $5, $6)
		RETURNING key_id, created_at
	`, token, namePtr, keyType, label, expires, createdBy).Scan(&k.ID, &k.CreatedAt)
	if db.IsUniqueViolation(err) {
		return nil, ErrKeyTokenTaken
	}
	if err != nil {
		return nil, err
	}
	return k, nil
}

// ListAgentKeys 列出 Agent 密钥；agentName 非空时按 Agent 过滤。
// 返回值不含 token 全文，只有 hint。
func ListAgentKeys(ctx context.Context, agentName string) ([]models.AgentKey, error) {
	q := `SELECT key_id, key_token, agent_name, key_type, label, expires_at, used_at, created_by, created_at
	      FROM agent_keys`
	args := []any{}
	if agentName != "" {
		q += ` WHERE agent_name = $1`
		args = append(args, agentName)
	}
	q += ` ORDER BY created_at DESC`

	rows, err := db.DB.QueryContext(ctx, q, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	keys := []models.AgentKey{}
	for rows.Next() {
		var k models.AgentKey
		var token string
		if err := rows.Scan(&k.ID, &token, &k.AgentName, &k.KeyType, &k.Label,
			&k.ExpiresAt, &k.UsedAt, &k.CreatedBy, &k.CreatedAt); err != nil {
			return nil, err
		}
		k.TokenHint = models.TokenHint(token) // 不回传全文
		keys = append(keys, k)
	}
	return keys, rows.Err()
}

// DeleteAgentKey 吊销一条 Agent 密钥。
func DeleteAgentKey(ctx context.Context, id uuid.UUID) error {
	tag, err := db.DB.ExecContext(ctx, `DELETE FROM agent_keys WHERE key_id = $1`, id)
	if err != nil {
		return err
	}
	if n, _ := tag.RowsAffected(); n == 0 {
		return ErrKeyNotFound
	}
	return nil
}

// BindAgentKey 把一条密钥绑定到指定 Agent 名。
func BindAgentKey(ctx context.Context, id uuid.UUID, agentName string) error {
	tag, err := db.DB.ExecContext(ctx,
		`UPDATE agent_keys SET agent_name = $2 WHERE key_id = $1`, id, agentName)
	if err != nil {
		return err
	}
	if n, _ := tag.RowsAffected(); n == 0 {
		return ErrKeyNotFound
	}
	return nil
}

// VerifyAgentKey 校验 Agent 密钥并返回它绑定的 Agent 名（未绑定时返回空串）。
//
// 一次性密钥在校验通过时立刻写 used_at —— 用 WHERE used_at IS NULL 保证并发下
// 只有一个请求能把它标记掉，避免两个 Agent 拿同一把一次性密钥同时注册成功。
func VerifyAgentKey(ctx context.Context, token string) (string, error) {
	var (
		id        uuid.UUID
		agentName *string
		keyType   string
		expiresAt *time.Time
		usedAt    *time.Time
	)
	err := db.DB.QueryRowContext(ctx, `
		SELECT key_id, agent_name, key_type, expires_at, used_at
		FROM agent_keys WHERE key_token = $1
	`, token).Scan(&id, &agentName, &keyType, &expiresAt, &usedAt)
	if errors.Is(err, sql.ErrNoRows) {
		return "", ErrKeyNotFound
	}
	if err != nil {
		return "", err
	}
	if err := checkKeyUsable(keyType, expiresAt, usedAt); err != nil {
		return "", err
	}

	if keyType == models.KeyOneTime {
		tag, err := db.DB.ExecContext(ctx,
			`UPDATE agent_keys SET used_at = NOW() WHERE key_id = $1 AND used_at IS NULL`, id)
		if err != nil {
			return "", err
		}
		if n, _ := tag.RowsAffected(); n == 0 {
			return "", ErrKeyUsed // 并发下被别人抢先用掉了
		}
	}

	if agentName == nil {
		return "", nil
	}
	// ★ 2026-10-03 修（与 VerifyAgent 同一处缺口）：密钥本身可能没过期，
	//   但**它所属的 Agent 可能已被停用**。原来这里直接返回名字，
	//   于是 Bearer 令牌能绕过「停用 = 撤销一切权限」——
	//   而 key_token 正是注释里标「推荐」的那条鉴权路径。
	//
	//   停用后返回空身份（不是错误）：调用方看到的是「这个令牌不是某个
	//   Agent 的令牌」，与令牌不存在不可区分，不泄露身份是否存在。
	var status string
	if err := db.DB.QueryRowContext(ctx,
		`SELECT status FROM agents WHERE agent_name = $1`, *agentName).Scan(&status); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return "", nil
		}
		return "", err
	}
	if status == "disabled" {
		return "", nil
	}
	return *agentName, nil
}

// ClaimAgentKey 在待绑定密钥首次注册时把它落定到该 Agent 名。
// 已绑定的密钥不受影响（WHERE agent_name IS NULL），因此不能借一把已绑定的密钥改注册别的 Agent。
func ClaimAgentKey(ctx context.Context, token, agentName string) error {
	_, err := db.DB.ExecContext(ctx,
		`UPDATE agent_keys SET agent_name = $2 WHERE key_token = $1 AND agent_name IS NULL`,
		token, agentName)
	return err
}

// ---------- 用户密钥（用户自助签发） ----------

// CreateUserKey 为用户签发一条客户端连接密钥。
func CreateUserKey(ctx context.Context, userID uuid.UUID, label, keyType string, expiresHours int) (*models.UserKey, error) {
	expires, err := expiryFor(keyType, expiresHours)
	if err != nil {
		return nil, err
	}
	token, err := newToken()
	if err != nil {
		return nil, err
	}

	k := &models.UserKey{
		Token:     token,
		TokenHint: models.TokenHint(token),
		UserID:    userID,
		Label:     label,
		KeyType:   keyType,
		ExpiresAt: expires,
	}
	err = db.DB.QueryRowContext(ctx, `
		INSERT INTO user_keys (key_token, user_id, label, key_type, expires_at)
		VALUES ($1, $2, $3, $4, $5)
		RETURNING key_id, created_at
	`, token, userID, label, keyType, expires).Scan(&k.ID, &k.CreatedAt)
	if err != nil {
		return nil, err
	}
	return k, nil
}

// ListUserKeys 列出某用户的连接密钥（不含 token 全文）。
func ListUserKeys(ctx context.Context, userID uuid.UUID) ([]models.UserKey, error) {
	rows, err := db.DB.QueryContext(ctx, `
		SELECT key_id, key_token, user_id, label, key_type, expires_at, used_at, created_at
		FROM user_keys WHERE user_id = $1 ORDER BY created_at DESC
	`, userID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	keys := []models.UserKey{}
	for rows.Next() {
		var k models.UserKey
		var token string
		if err := rows.Scan(&k.ID, &token, &k.UserID, &k.Label, &k.KeyType,
			&k.ExpiresAt, &k.UsedAt, &k.CreatedAt); err != nil {
			return nil, err
		}
		k.TokenHint = models.TokenHint(token)
		keys = append(keys, k)
	}
	return keys, rows.Err()
}

// DeleteUserKey 删除自己的一条密钥。带 user_id 条件，避免删掉别人的。
func DeleteUserKey(ctx context.Context, userID, keyID uuid.UUID) error {
	tag, err := db.DB.ExecContext(ctx,
		`DELETE FROM user_keys WHERE key_id = $1 AND user_id = $2`, keyID, userID)
	if err != nil {
		return err
	}
	if n, _ := tag.RowsAffected(); n == 0 {
		return ErrKeyNotFound
	}
	return nil
}

// VerifyUserKey 校验用户密钥并返回对应用户。
// 用户必须仍处于 active 状态——禁用账号后其密钥应当立即失效。
func VerifyUserKey(ctx context.Context, token string) (*models.User, error) {
	var (
		id        uuid.UUID
		userID    uuid.UUID
		keyType   string
		expiresAt *time.Time
		usedAt    *time.Time
	)
	err := db.DB.QueryRowContext(ctx, `
		SELECT key_id, user_id, key_type, expires_at, used_at
		FROM user_keys WHERE key_token = $1
	`, token).Scan(&id, &userID, &keyType, &expiresAt, &usedAt)
	if errors.Is(err, sql.ErrNoRows) {
		return nil, ErrKeyNotFound
	}
	if err != nil {
		return nil, err
	}
	if err := checkKeyUsable(keyType, expiresAt, usedAt); err != nil {
		return nil, err
	}

	if keyType == models.KeyOneTime {
		tag, err := db.DB.ExecContext(ctx,
			`UPDATE user_keys SET used_at = NOW() WHERE key_id = $1 AND used_at IS NULL`, id)
		if err != nil {
			return nil, err
		}
		if n, _ := tag.RowsAffected(); n == 0 {
			return nil, ErrKeyUsed
		}
	}

	u, err := GetUserByID(ctx, userID)
	if err != nil {
		return nil, err
	}
	if u.Status != "active" {
		return nil, ErrKeyNotFound // 账号已禁用，密钥一并失效
	}
	return u, nil
}

// ensureAgentRow 确保 agents 行存在，**不**改动已有行的任何字段。
//
// ★ 与 CreateOrUpdateAgent 的区别是刻意的：
//   - 那个把 status 设成 'online'、last_seen 设成 NOW()
//     —— 它服务于「插件刚启动，自报家门」，语义是「它现在活着」。
//   - 这个服务于「管理员预建一个身份」。它此刻**还没连上**，
//     写成 online 会让 WebUI 显示成在线，而实际上没有任何进程在跑
//     （症状：联系人里一个永远在线、却永远不回信的幽灵）。
//
// 所以这里只 INSERT，不 UPDATE；已存在就什么都不做。
// status 交给首次握手时的 TouchAgentLastSeen 去改。
func ensureAgentRow(ctx context.Context, agentName string) error {
	if _, err := db.DB.ExecContext(ctx, `
		INSERT INTO agents (agent_name, secret, workspaces, platform, status)
		VALUES ($1, '', '[]', $2, 'offline')
		ON CONFLICT (agent_name) DO NOTHING
	`, agentName, "mcp"); err != nil {
		if db.IsUniqueViolation(err) {
			return nil // 并发下另一个请求已建好
		}
		return fmt.Errorf("创建 agent 身份 %s: %w", agentName, err)
	}
	return nil
}