package middleware

import (
	"github.com/gin-gonic/gin"
	"net/http"
)

var (
	whitelistIP = map[string]bool{
		"192.168.1.100": true,
		"10.0.0.1":      true,
		// 添加更多的白名单IP地址
	}
	whitelistRoutes = map[string]bool{
		"/yuexin-pay/swagger": true,
		// 添加更多的白名单路由
	}
)

// WhitelistMiddleware 白名单中间件
func WhitelistMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 检查IP地址
		clientIP := c.ClientIP()
		if !isIPInWhitelist(clientIP) {
			c.JSON(http.StatusForbidden, gin.H{"error": "IP地址不在白名单中"})
			c.Abort()
			return
		}

		// 检查路由
		if !isRouteInWhitelist(c.FullPath()) {
			c.JSON(http.StatusForbidden, gin.H{"error": "路由不在白名单中"})
			c.Abort()
			return
		}

		c.Next()
	}
}

// isIPInWhitelist 检查IP地址是否在白名单中
func isIPInWhitelist(ip string) bool {
	_, ok := whitelistIP[ip]
	return ok
}

// isRouteInWhitelist 检查路由是否在白名单中
func isRouteInWhitelist(route string) bool {
	_, ok := whitelistRoutes[route]
	return ok
}