MailUI4Agents/gateway/internal/handler/keys.go

package handler

import (
	"net/http"
	"strings"

	"github.com/agentmail/gateway/internal/middleware"
	"github.com/agentmail/gateway/internal/models"
	"github.com/agentmail/gateway/internal/repo"
)

// ---------- 密钥管理 ----------
//
// 两套接口，权限边界不同：
//   /admin/agent-keys —— 管理员签发 Agent 接入密钥
//   /me/keys          —— 用户自助签发客户端连接密钥（不能注册 Agent）
//
// 密钥全文只在创建响应里出现一次，列表接口只给前 8 位 hint。

type createKeyRequest struct {
	// AgentName 仅 Agent 密钥使用；留空表示「待绑定」，首次注册时按注册请求的 name 落定
	AgentName string `json:"agent_name"`
	// Label 人类可读备注（如「我的笔记本」「CI 机器」）
	Label string `json:"label"`
	// KeyType permanent / one_time / timed
	KeyType string `json:"key_type"`
	// ExpiresHours 仅 timed 使用，必须为正
	ExpiresHours int `json:"expires_hours"`
	// KeyToken 仅 Agent 密钥使用：登记一把客户端已在本地生成的密钥。
	// 插件首次安装时自己生成密钥并打印出来，管理员把它填到这里完成登记，
	// 密钥全文因此不需要从服务器往客户端传。留空则由服务器生成。
	KeyToken string `json:"key_token"`
}

// normalizeKeyType 默认给 permanent，避免调用方漏填时落到非法值
func normalizeKeyType(t string) string {
	t = strings.TrimSpace(t)
	if t == "" {
		return models.KeyPermanent
	}
	return t
}

// POST /api/v1/admin/agent-keys
func CreateAgentKey(w http.ResponseWriter, r *http.Request) {
	admin := middleware.GetUser(r)
	if admin == nil {
		Error(w, http.StatusUnauthorized, "not authenticated")
		return
	}

	var req createKeyRequest
	if !DecodeBody(w, r, &req) {
		return
	}

	key, err := repo.CreateAgentKey(r.Context(),
		strings.TrimSpace(req.AgentName), normalizeKeyType(req.KeyType),
		strings.TrimSpace(req.Label), req.ExpiresHours, admin.ID,
		strings.TrimSpace(req.KeyToken))
	if err != nil {
		writeKeyErr(w, err)
		return
	}

	// 唯一一次回传全文
	JSON(w, http.StatusOK, map[string]any{"key": key})
}

// GET /api/v1/admin/agent-keys?agent_name=xxx
func ListAgentKeys(w http.ResponseWriter, r *http.Request) {
	keys, err := repo.ListAgentKeys(r.Context(), r.URL.Query().Get("agent_name"))
	if err != nil {
		Error(w, http.StatusInternalServerError, "Failed to list keys")
		return
	}
	JSON(w, http.StatusOK, map[string]any{"keys": keys})
}

// DELETE /api/v1/admin/agent-keys/{id}
func DeleteAgentKey(w http.ResponseWriter, r *http.Request) {
	id, ok := pathUUID(w, r, "id")
	if !ok {
		return
	}
	if err := repo.DeleteAgentKey(r.Context(), id); err != nil {
		writeKeyErr(w, err)
		return
	}
	JSON(w, http.StatusOK, map[string]string{"status": "deleted"})
}

type bindKeyRequest struct {
	AgentName string `json:"agent_name"`
}

// POST /api/v1/admin/agent-keys/{id}/bind
func BindAgentKey(w http.ResponseWriter, r *http.Request) {
	id, ok := pathUUID(w, r, "id")
	if !ok {
		return
	}
	var req bindKeyRequest
	if !DecodeBody(w, r, &req) {
		return
	}
	name := strings.TrimSpace(req.AgentName)
	if name == "" {
		Error(w, http.StatusBadRequest, "Missing agent_name")
		return
	}
	if err := repo.BindAgentKey(r.Context(), id, name); err != nil {
		writeKeyErr(w, err)
		return
	}
	JSON(w, http.StatusOK, map[string]string{"status": "bound", "agent_name": name})
}

// ---------- 用户连接密钥 ----------

// POST /api/v1/me/keys
func CreateMyKey(w http.ResponseWriter, r *http.Request) {
	user := middleware.GetUser(r)
	if user == nil {
		Error(w, http.StatusUnauthorized, "not authenticated")
		return
	}

	var req createKeyRequest
	if !DecodeBody(w, r, &req) {
		return
	}

	key, err := repo.CreateUserKey(r.Context(), user.ID,
		strings.TrimSpace(req.Label), normalizeKeyType(req.KeyType), req.ExpiresHours)
	if err != nil {
		writeKeyErr(w, err)
		return
	}
	JSON(w, http.StatusOK, map[string]any{"key": key})
}

// GET /api/v1/me/keys
func ListMyKeys(w http.ResponseWriter, r *http.Request) {
	user := middleware.GetUser(r)
	if user == nil {
		Error(w, http.StatusUnauthorized, "not authenticated")
		return
	}
	keys, err := repo.ListUserKeys(r.Context(), user.ID)
	if err != nil {
		Error(w, http.StatusInternalServerError, "Failed to list keys")
		return
	}
	JSON(w, http.StatusOK, map[string]any{"keys": keys})
}

// DELETE /api/v1/me/keys/{id}
func DeleteMyKey(w http.ResponseWriter, r *http.Request) {
	user := middleware.GetUser(r)
	if user == nil {
		Error(w, http.StatusUnauthorized, "not authenticated")
		return
	}
	id, ok := pathUUID(w, r, "id")
	if !ok {
		return
	}
	// repo 层带 user_id 条件，删不到就是不属于自己或不存在，统一 404
	if err := repo.DeleteUserKey(r.Context(), user.ID, id); err != nil {
		writeKeyErr(w, err)
		return
	}
	JSON(w, http.StatusOK, map[string]string{"status": "deleted"})
}