Wormhole/internal/api/extra_handlers.go

package api

import (
	"net/http"
	"strconv"

	"github.com/gin-gonic/gin"
	"github.com/wormhole/wormhole/internal/auth"
	"github.com/wormhole/wormhole/internal/store"
)

func (s *Server) dashboardRankings(c *gin.Context) {
	rankType := c.DefaultQuery("type", "tunnels")
	page, _ := strconv.Atoi(c.DefaultQuery("page", "1"))
	pageSize, _ := strconv.Atoi(c.DefaultQuery("page_size", "50"))
	if page < 1 {
		page = 1
	}
	if pageSize < 1 {
		pageSize = 50
	}
	if pageSize > 500 {
		pageSize = 500
	}
	offset := (page - 1) * pageSize
	q := c.Query("q")

	switch rankType {
	case "tunnels":
		items, total, err := s.store.ListTunnelsByFlow(pageSize, offset)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}
		c.JSON(http.StatusOK, gin.H{"items": items, "total": total})
	case "hosts":
		items, total, err := s.store.ListHostsByFlow(pageSize, offset)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}
		c.JSON(http.StatusOK, gin.H{"items": items, "total": total})
	case "ips":
		items, total, err := s.store.ListAllRequestIPs(pageSize, offset, q)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}
		c.JSON(http.StatusOK, gin.H{"items": requestIPViews(s.store, items), "total": total})
	default:
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid type"})
	}
}

func (s *Server) listAPIKeys(c *gin.Context) {
	claims := getClaims(c)
	keys, err := s.store.ListAPIKeys(claims.UserID, auth.IsAdmin(claims.Role))
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, keys)
}

func (s *Server) createAPIKey(c *gin.Context) {
	claims := getClaims(c)
	var req struct {
		Name string `json:"name"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "bad request"})
		return
	}
	key, plain, err := s.store.CreateAPIKey(req.Name, claims.UserID)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{
		"key":     key,
		"api_key": plain,
		"note":    "请妥善保存 API Key，仅显示一次",
	})
}

func (s *Server) deleteAPIKey(c *gin.Context) {
	claims := getClaims(c)
	id, _ := strconv.ParseUint(c.Param("id"), 10, 64)
	if err := s.store.DeleteAPIKey(uint(id), claims.UserID, auth.IsAdmin(claims.Role)); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"ok": true})
}

func (s *Server) pauseTunnel(c *gin.Context) {
	id, _ := strconv.ParseUint(c.Param("id"), 10, 64)
	claims := getClaims(c)
	existing, err := s.store.GetTunnelByID(uint(id))
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	if err := store.ClientOwnerCheck(s.store, existing.ClientID, claims.UserID, auth.IsAdmin(claims.Role)); err != nil {
		c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
		return
	}
	_ = s.proxy.StopTunnel(uint(id))
	existing.Status = false
	existing.RunStatus = false
	_ = s.store.UpdateTunnel(existing)
	s.cache.UpsertTunnel(existing)
	c.JSON(http.StatusOK, gin.H{"ok": true})
}

func (s *Server) resumeTunnel(c *gin.Context) {
	id, _ := strconv.ParseUint(c.Param("id"), 10, 64)
	claims := getClaims(c)
	existing, err := s.store.GetTunnelByID(uint(id))
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	if err := store.ClientOwnerCheck(s.store, existing.ClientID, claims.UserID, auth.IsAdmin(claims.Role)); err != nil {
		c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
		return
	}
	existing.Status = true
	_ = s.store.UpdateTunnel(existing)
	s.cache.UpsertTunnel(existing)
	_ = s.proxy.StartTunnel(uint(id))
	c.JSON(http.StatusOK, gin.H{"ok": true})
}

func (s *Server) pauseHost(c *gin.Context) {
	id, _ := strconv.ParseUint(c.Param("id"), 10, 64)
	claims := getClaims(c)
	existing, err := s.store.GetHostByID(uint(id))
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	if err := store.ClientOwnerCheck(s.store, existing.ClientID, claims.UserID, auth.IsAdmin(claims.Role)); err != nil {
		c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
		return
	}
	existing.Status = false
	_ = s.store.UpdateHost(existing)
	s.cache.UpsertHost(existing)
	c.JSON(http.StatusOK, gin.H{"ok": true})
}

func (s *Server) resumeHost(c *gin.Context) {
	id, _ := strconv.ParseUint(c.Param("id"), 10, 64)
	claims := getClaims(c)
	existing, err := s.store.GetHostByID(uint(id))
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	if err := store.ClientOwnerCheck(s.store, existing.ClientID, claims.UserID, auth.IsAdmin(claims.Role)); err != nil {
		c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
		return
	}
	existing.Status = true
	_ = s.store.UpdateHost(existing)
	s.cache.UpsertHost(existing)
	c.JSON(http.StatusOK, gin.H{"ok": true})
}

func (s *Server) openapiDoc(c *gin.Context) {
	c.Header("Content-Type", "application/json")
	c.String(http.StatusOK, openAPISpec)
}