router.go 2.38 KB
Newer Older
qiuqunfeng's avatar
qiuqunfeng committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102
package api

import (
	"fmt"
	"io"
	"net/http"

	"github.com/gin-gonic/gin"
	"gitlab.com/security-rd/go-pkg/logging"
	"gitlab.com/tensorsecurity-rd/waf-console/internal/config"
	"gitlab.com/tensorsecurity-rd/waf-console/internal/utils"
	"gitlab.com/tensorsecurity-rd/waf-console/pkg/generated/clientset/versioned"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"k8s.io/client-go/rest"
)

const (
	DB_USER     = "ivan"
	DB_PASSWORD = "Mysql-ha@123"
	DB_HOST     = "localhost"
	DB_PORT     = "3306"
	DB_NAME     = "waf"
)

func SetRouters() *gin.Engine {
	var engine *gin.Engine
	dns := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=True&loc=Local", DB_USER, DB_PASSWORD, DB_HOST, DB_PORT, DB_NAME)
	db, err := gorm.Open(mysql.Open(dns), &gorm.Config{})
	if err != nil {
		panic(err)
	}

	if !config.Conf.Debug {
		// 生产模式
		logging.Get().Info().Msg("release mode")
		engine = ReleaseRouter()
		engine.Use(
		// middleware.RequestCostHandler(),
		// middleware.CustomLogger(),
		// middleware.CustomRecovery(),
		// middleware.CorsHandler(),
		)
	} else {
		// 开发调试模式
		logging.Get().Info().Msg("debug mode")
		engine = gin.New()
		engine.Use(
			// middleware.RequestCostHandler(),
			gin.Logger(),
			// middleware.CustomRecovery(),
			// middleware.CorsHandler(),
		)
	}
	// set up trusted agents
	err = engine.SetTrustedProxies([]string{"127.0.0.1"})
	if err != nil {
		panic(err)
	}
	// ping
	engine.GET("/ping", func(c *gin.Context) {
		c.AbortWithStatusJSON(http.StatusOK, gin.H{
			"message": "pong!",
		})
	})

	// 设置 API 路由
	// SetIPInforRouter(engine)
	// loadkubeConfig()
	// clientcmd.LoadFromFile("kubeconfig.yaml")
	client := versioned.NewForConfigOrDie(&rest.Config{
		Host: "https://127.0.0.1:6443",
		TLSClientConfig: rest.TLSClientConfig{
			Insecure: false,
			CAData:   []byte(""),
			CertData: []byte(""),
			KeyData:  []byte(""),
		},
		// BearerToken: "1234567890",
	})
	SetWafRouter(engine, client, db)

	// 统一处理 404
	engine.NoRoute(func(c *gin.Context) {
		utils.AssembleResponse(c, nil, err)
		// response2.Resp().SetHttpCode(http.StatusNotFound).FailCode(c, errors.NotFound)
	})

	return engine
}

// ReleaseRouter 生产模式使用官方建议设置为 release 模式
func ReleaseRouter() *gin.Engine {
	// 切换到生产模式
	gin.SetMode(gin.ReleaseMode)
	// 禁用 gin 输出接口访问日志
	gin.DefaultWriter = io.Discard

	engine := gin.New()

	return engine
}