cmd.go

package app

import (
	"encoding/base64"
	"fmt"
	"os"
	"time"

	"github.com/rs/zerolog/log"
	"github.com/spf13/cobra"
	"gitlab.com/tensorsecurity-rd/waf-console/api"
	"gitlab.com/tensorsecurity-rd/waf-console/cmd/config"
	"gitlab.com/tensorsecurity-rd/waf-console/internal/service"
	es "gitlab.com/tensorsecurity-rd/waf-console/internal/store"
	"gitlab.com/tensorsecurity-rd/waf-console/internal/utils"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
)

func NewRootCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "waf-console",
		Short: "Start waf-console service.",
		Args:  cobra.ExactArgs(0),
		RunE: func(cmd *cobra.Command, args []string) error {
			config := config.LoadConfig()
			debugMode := os.Getenv("DEBUG_MODE")
			log.Info().Msgf("DEBUG_MODE: %s", debugMode)
			if debugMode == "true" {
				config.Debug = true
				// config.Conf.Debug = true
			}
			dbConfig := config.DBConfig
			var db *gorm.DB
			var err error
			if dbConfig != nil {
				dns := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=True&loc=Local", dbConfig.User, dbConfig.Password, dbConfig.Host, dbConfig.Port, dbConfig.Database)
				log.Info().Msgf("dns: %s", dns)
				db, err = gorm.Open(mysql.Open(dns), &gorm.Config{})
				if err != nil {
					panic(err)
				}
			} else {
				panic("dbConfig is nil")
			}

			clusterClientManager := utils.NewClusterClientManager()
			for _, regionConfig := range config.RegionConfigs {
				var caData []byte
				var clientCertData []byte
				var clientKeyData []byte
				if regionConfig.CAData != "" {
					caData, err = base64.StdEncoding.DecodeString(regionConfig.CAData)
					if err != nil {
						panic(err)
					}
				}
				if regionConfig.ClientCertData != "" {
					clientCertData, err = base64.StdEncoding.DecodeString(regionConfig.ClientCertData)
					if err != nil {
						panic(err)
					}
				}
				if regionConfig.ClientKeyData != "" {
					clientKeyData, err = base64.StdEncoding.DecodeString(regionConfig.ClientKeyData)
					if err != nil {
						panic(err)
					}
				}
				kubeClient := utils.NewKubeClient(regionConfig.ApiServer, caData, clientCertData, clientKeyData, regionConfig.Insecure)
				clusterClientManager.AddClient(regionConfig.RegionCode, kubeClient)
			}

			esClient, err := es.CreateEsClientFromConfig(config.ElasticsearchConfig)
			if err != nil {
				panic(err)
			}
			regionUrlMap := make(map[string]string)
			for _, regionConfig := range config.RegionConfigs {
				regionUrlMap[regionConfig.RegionCode] = regionConfig.ApiServer
			}
			e := api.SetRouters(db, clusterClientManager, config.GatewayUrl, config.SSOUrl, esClient, config.Debug, regionUrlMap)
			esStore := es.NewESStore(es.Config{
				ESBatchSize:   100,
				ESConcurrency: 10,
				ESTimeout:     10 * time.Second,
			}, esClient)
			esStore.Init()
			logConsumerService := service.NewLogConsumerService(db, esStore, config.KafkaConfig)
			go logConsumerService.Consume()
			return e.Run(":8080")
		},
	}
}