Prism/internal/api/handlers.go

package api

import (
	"context"
	"net/http"
	"strconv"
	"time"

	"github.com/gin-gonic/gin"

	"github.com/prism/proxy/internal/auth"

	"github.com/prism/proxy/internal/country"
	"github.com/prism/proxy/internal/health"
	"github.com/prism/proxy/internal/metrics"
	applog "github.com/prism/proxy/internal/log"
	"github.com/prism/proxy/internal/mihomoapi"
	"github.com/prism/proxy/internal/service"
	"github.com/prism/proxy/internal/store"
	"github.com/prism/proxy/internal/subscription"
	"github.com/prism/proxy/internal/trafficlog"
)

type Server struct {
	app          *service.App
	health       *health.Checker
	metrics      *metrics.Collector
	subscription *subscription.Service
	logs         *applog.Collector
	traffic      *trafficlog.Collector
	loginLimiter *auth.LoginLimiter
}

func NewServer(app *service.App, hc *health.Checker, mc *metrics.Collector, sub *subscription.Service, logs *applog.Collector, traffic *trafficlog.Collector) *Server {
	return &Server{
		app:          app,
		health:       hc,
		metrics:      mc,
		subscription: sub,
		logs:         logs,
		traffic:      traffic,
		loginLimiter: auth.NewLoginLimiter(app.Store),
	}
}

func (s *Server) syncMihomoClient(ctx context.Context) {
	base, secret, err := s.app.MihomoAPIBase(ctx)
	if err != nil {
		return
	}
	s.health.UpdateAPIBase(base, secret)
	s.metrics.UpdateAPIBase(base, secret)
	if s.traffic != nil {
		s.traffic.UpdateClient(mihomoapi.New(base, secret))
	}
}

func (s *Server) Router() *gin.Engine {
	gin.SetMode(gin.ReleaseMode)
	r := gin.New()
	r.Use(gin.Recovery(), s.corsMiddleware())

	r.GET("/api/v1/openapi.json", s.serveOpenAPIJSON)
	r.GET("/api/v1/openapi.yaml", s.serveOpenAPIYAML)
	r.GET("/api/docs", s.serveOpenAPIDocs)

	api := r.Group("/api/v1")
	api.POST("/auth/login", s.login)
	api.GET("/auth/status", s.authStatus)

	protected := api.Group("")
	protected.Use(s.authMiddleware())
	{
		protected.GET("/dashboard", s.dashboard)
		protected.GET("/connections", s.connections)

		protected.GET("/subscriptions", s.listSubscriptions)
		protected.POST("/subscriptions", s.createSubscription)
		protected.PUT("/subscriptions/:id", s.updateSubscription)
		protected.DELETE("/subscriptions/:id", s.deleteSubscription)
		protected.POST("/subscriptions/:id/refresh", s.refreshSubscription)

		protected.GET("/proxies", s.listProxies)
		protected.POST("/proxies", s.createProxy)
		protected.PUT("/proxies/:id", s.updateProxy)
		protected.DELETE("/proxies/:id", s.deleteProxy)
		protected.POST("/proxies/:id/test", s.testProxy)
		protected.POST("/proxies/test-all", s.testAllProxies)
		protected.POST("/proxies/rebuild-countries", s.rebuildCountries)
		protected.GET("/proxies/export", s.exportProxies)
		protected.POST("/proxies/import", s.importProxies)
		protected.GET("/countries", s.listCountries)

		protected.GET("/outbound", s.outboundOverview)
		protected.PUT("/outbound/groups/:name", s.selectOutboundGroup)
		protected.PUT("/outbound/default-policy", s.setDefaultPolicy)

		protected.GET("/traffic-logs", s.listTrafficLogs)

		protected.GET("/rules", s.listRules)
		protected.POST("/rules", s.createRule)
		protected.PUT("/rules/:id", s.updateRule)
		protected.DELETE("/rules/:id", s.deleteRule)
		protected.GET("/rules/merged", s.mergedRules)
		protected.POST("/rules/test", s.testRule)
		protected.GET("/rules/export", s.exportRules)
		protected.POST("/rules/import", s.importRules)

		protected.GET("/settings", s.getSettings)
		protected.PUT("/settings", s.putSettings)

		protected.GET("/api-keys", s.listAPIKeys)
		protected.POST("/api-keys", s.createAPIKey)
		protected.DELETE("/api-keys/:id", s.deleteAPIKey)

		protected.GET("/logs", s.listLogs)
		protected.GET("/audit-logs", s.listAuditLogs)

		protected.GET("/system/status", s.systemStatus)
		protected.POST("/system/reload", s.systemReload)
		protected.GET("/system/geo", s.geoStatus)
		protected.POST("/system/geo/update", s.geoUpdate)
	}

	r.NoRoute(s.serveSPA)
	return r
}

func (s *Server) corsMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Header("Access-Control-Allow-Origin", "*")
		c.Header("Access-Control-Allow-Headers", "Authorization, Content-Type, X-API-Key")
		c.Header("Access-Control-Allow-Methods", "GET,POST,PUT,DELETE,OPTIONS")
		if c.Request.Method == http.MethodOptions {
			c.AbortWithStatus(204)
			return
		}
		c.Next()
	}
}

func (s *Server) dashboard(c *gin.Context) {
	ctx := c.Request.Context()
	s.metrics.RefreshAsync(10 * time.Second)
	s.app.RefreshEngineStatusAsync()
	history, _ := s.app.Store.ListTrafficSnapshots(ctx, 60)
	conns := s.metrics.CachedConnections()
	policyGroups, _ := s.app.PolicyGroupRuntimes(ctx)
	names, _ := s.app.ProxyDisplayNames(ctx)
	liveDomainsAll, liveNodesAll := metrics.AggregateConnections(conns, policyGroups)
	liveDomainsAll = metrics.ApplyDomainLabels(liveDomainsAll, names)
	liveNodesAll = metrics.ApplyTrafficLabels(liveNodesAll, names)

	liveDomainLimit, liveDomainOffset := parseLimitOffset(c.DefaultQuery("live_domain_limit", "50"), c.DefaultQuery("live_domain_offset", "0"))
	liveNodeLimit, liveNodeOffset := parseLimitOffset(c.DefaultQuery("live_node_limit", "50"), c.DefaultQuery("live_node_offset", "0"))
	periodDomainLimit, periodDomainOffset := parseLimitOffset(c.DefaultQuery("period_domain_limit", "50"), c.DefaultQuery("period_domain_offset", "0"))
	periodNodeLimit, periodNodeOffset := parseLimitOffset(c.DefaultQuery("period_node_limit", "50"), c.DefaultQuery("period_node_offset", "0"))

	liveDomains, liveDomainTotal := paginateSlice(liveDomainsAll, liveDomainLimit, liveDomainOffset)
	liveNodes, liveNodeTotal := paginateSlice(liveNodesAll, liveNodeLimit, liveNodeOffset)

	hostNodeTotals, _ := s.app.Store.AggregateTrafficByHostNode(ctx, 24, periodDomainLimit, periodDomainOffset)
	nodeTotals, _ := s.app.Store.AggregateTrafficByNode(ctx, 24, periodNodeLimit, periodNodeOffset)
	periodDomainTotal, _ := s.app.Store.CountAggregateTrafficByHostNode(ctx, 24)
	periodNodeTotal, _ := s.app.Store.CountAggregateTrafficByNode(ctx, 24)
	periodDomains := metrics.TrafficHostNodeAggToItems(hostNodeTotals, names)
	periodNodes := metrics.TrafficAggToItems(nodeTotals, names)

	c.JSON(http.StatusOK, gin.H{
		"snapshot":           s.metrics.Snapshot(ctx),
		"history":            history,
		"connections":        conns,
		"engine_running":     s.app.EngineStatusCached(),
		"live_domains":       liveDomains,
		"live_nodes":         liveNodes,
		"live_domain_total":  liveDomainTotal,
		"live_node_total":    liveNodeTotal,
		"period_domains":     periodDomains,
		"period_nodes":       periodNodes,
		"period_domain_total": periodDomainTotal,
		"period_node_total":  periodNodeTotal,
		"period_hours":       24,
	})
}

func (s *Server) connections(c *gin.Context) {
	conns, err := s.metrics.GetConnectionsCached(8 * time.Second)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"connections": conns})
}

func (s *Server) listSubscriptions(c *gin.Context) {
	items, err := s.app.Store.ListSubscriptions(c.Request.Context())
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"items": items})
}

func (s *Server) createSubscription(c *gin.Context) {
	var sub store.Subscription
	if err := c.ShouldBindJSON(&sub); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	if sub.Type == "" {
		sub.Type = "auto"
	}
	if sub.IntervalSec <= 0 {
		sub.IntervalSec = 3600
	}
	sub.Enabled = true
	if err := s.app.Store.CreateSubscription(c.Request.Context(), &sub); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	_ = s.app.Store.AddAuditLog(c.Request.Context(), "create", "subscription", sub.Name)
	go func(id int64) {
		ctx := context.Background()
		_ = s.subscription.FetchAndStore(ctx, id)
		_ = s.app.ReloadEngine(ctx)
	}(sub.ID)
	c.JSON(http.StatusCreated, sub)
}

func (s *Server) updateSubscription(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	sub, err := s.app.Store.GetSubscription(c.Request.Context(), id)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	if err := c.ShouldBindJSON(sub); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	sub.ID = id
	if err := s.app.Store.UpdateSubscription(c.Request.Context(), sub); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, sub)
}

func (s *Server) deleteSubscription(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	if err := s.app.Store.DeleteSubscription(c.Request.Context(), id); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	_ = s.app.ReloadEngine(c.Request.Context())
	c.JSON(http.StatusOK, gin.H{"ok": true})
}

func (s *Server) refreshSubscription(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	if err := s.subscription.FetchAndStore(c.Request.Context(), id); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	if err := s.app.ReloadEngine(c.Request.Context()); err != nil {
		c.JSON(http.StatusOK, gin.H{"ok": true, "warning": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"ok": true})
}

func (s *Server) listProxies(c *gin.Context) {
	source := c.Query("source")
	country := c.Query("country")
	var subID *int64
	if v := c.Query("subscription_id"); v != "" {
		id, _ := strconv.ParseInt(v, 10, 64)
		subID = &id
	}
	limit, offset := parseLimitOffset(c.DefaultQuery("limit", "50"), c.DefaultQuery("offset", "0"))
	items, err := s.app.Store.ListProxyNodesPage(c.Request.Context(), source, subID, country, limit, offset)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	total, err := s.app.Store.CountProxyNodes(c.Request.Context(), source, subID, country)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"items": items, "total": total})
}

type proxyReq struct {
	Name      string `json:"name"`
	Type      string `json:"type"`
	RawConfig string `json:"raw_config"`
	Enabled   *bool  `json:"enabled"`
}

func (s *Server) createProxy(c *gin.Context) {
	var req proxyReq
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	n := store.ProxyNode{
		Name:        req.Name,
		RuntimeName: store.RuntimeName("manual", 0, req.Name),
		Type:        req.Type,
		RawConfig:   req.RawConfig,
		Source:      "manual",
		Country:     country.Detect(req.Name),
		Enabled:     true,
	}
	if req.Enabled != nil {
		n.Enabled = *req.Enabled
	}
	if err := s.app.Store.CreateProxyNode(c.Request.Context(), &n); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	_ = s.app.ReloadEngine(c.Request.Context())
	c.JSON(http.StatusCreated, n)
}

func (s *Server) updateProxy(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	n, err := s.app.Store.GetProxyNode(c.Request.Context(), id)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	var req proxyReq
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	if req.Name != "" {
		n.Name = req.Name
		n.Country = country.Detect(req.Name)
	}
	if req.Type != "" {
		n.Type = req.Type
	}
	if req.RawConfig != "" {
		n.RawConfig = req.RawConfig
	}
	if req.Enabled != nil {
		n.Enabled = *req.Enabled
	}
	if err := s.app.Store.UpdateProxyNode(c.Request.Context(), n); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	_ = s.app.ReloadEngine(c.Request.Context())
	c.JSON(http.StatusOK, n)
}

func (s *Server) deleteProxy(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	n, err := s.app.Store.GetProxyNode(c.Request.Context(), id)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	if n.Source != "manual" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "only manual proxies can be deleted"})
		return
	}
	if err := s.app.Store.DeleteProxyNode(c.Request.Context(), id); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	_ = s.app.ReloadEngine(c.Request.Context())
	c.JSON(http.StatusOK, gin.H{"ok": true})
}

func (s *Server) testProxy(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	latency, err := s.health.TestNode(c.Request.Context(), id, s.app.ReloadEngine)
	if err != nil {
		c.JSON(http.StatusOK, gin.H{"latency_ms": -1, "alive": false, "error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"latency_ms": latency, "alive": latency >= 0})
}

func (s *Server) testAllProxies(c *gin.Context) {
	if err := s.health.TestAll(c.Request.Context(), s.app.ReloadEngine); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"ok": true})
}

func (s *Server) listRules(c *gin.Context) {
	source := c.DefaultQuery("source", "manual")
	items, err := s.app.Store.ListRules(c.Request.Context(), source)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"items": items})
}

func (s *Server) createRule(c *gin.Context) {
	var r store.Rule
	if err := c.ShouldBindJSON(&r); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	r.Source = "manual"
	r.Enabled = true
	if r.Priority <= 0 {
		p, _ := s.app.Store.NextManualRulePriority(c.Request.Context())
		r.Priority = p
	}
	if err := s.app.Store.CreateRule(c.Request.Context(), &r); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	_ = s.app.ReloadEngine(c.Request.Context())
	c.JSON(http.StatusCreated, r)
}

func (s *Server) updateRule(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	r, err := s.app.Store.GetRule(c.Request.Context(), id)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	if err := c.ShouldBindJSON(r); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	r.ID = id
	if err := s.app.Store.UpdateRule(c.Request.Context(), r); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	_ = s.app.ReloadEngine(c.Request.Context())
	c.JSON(http.StatusOK, r)
}

func (s *Server) deleteRule(c *gin.Context) {
	id, _ := strconv.ParseInt(c.Param("id"), 10, 64)
	if err := s.app.Store.DeleteRule(c.Request.Context(), id); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	_ = s.app.ReloadEngine(c.Request.Context())
	c.JSON(http.StatusOK, gin.H{"ok": true})
}

func (s *Server) listLogs(c *gin.Context) {
	level := c.Query("level")
	keyword := c.Query("keyword")
	limit, offset := parseLimitOffset(c.DefaultQuery("limit", "50"), c.DefaultQuery("offset", "0"))
	items, total, err := s.app.Store.ListRequestLogs(c.Request.Context(), level, keyword, limit, offset)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"items": items, "total": total})
}

func (s *Server) listAuditLogs(c *gin.Context) {
	limit, offset := parseLimitOffset(c.DefaultQuery("limit", "50"), c.DefaultQuery("offset", "0"))
	items, err := s.app.Store.ListAuditLogs(c.Request.Context(), limit, offset)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	total, err := s.app.Store.CountAuditLogs(c.Request.Context())
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"items": items, "total": total})
}

func (s *Server) systemStatus(c *gin.Context) {
	c.JSON(http.StatusOK, gin.H{
		"engine_running":    s.app.EngineReachable(c.Request.Context()),
		"has_config":        len(s.app.Cache.GetLastYAML()) > 0,
		"last_reload_error": s.app.LastReloadError,
	})
}

func (s *Server) systemReload(c *gin.Context) {
	if err := s.app.ReloadEngine(c.Request.Context()); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"ok": true})
}