added simple implantation of server flag

aquasecurity · Jan 9, 2025 · a645543 · a645543
1 parent 22fa7ba
commit a645543
Show file tree

Hide file tree

Showing 8 changed files with 342 additions and 157 deletions.
diff --git a/cmd/tracee-ebpf/main.go b/cmd/tracee-ebpf/main.go
@@ -10,7 +10,8 @@ import (
 
 	"github.com/aquasecurity/tracee/pkg/cmd"
 	"github.com/aquasecurity/tracee/pkg/cmd/flags"
-	"github.com/aquasecurity/tracee/pkg/cmd/flags/server"
+
+	// "github.com/aquasecurity/tracee/pkg/cmd/flags/server"
 	"github.com/aquasecurity/tracee/pkg/cmd/initialize"
 	"github.com/aquasecurity/tracee/pkg/cmd/urfave"
 	"github.com/aquasecurity/tracee/pkg/logger"
@@ -128,31 +129,31 @@ func main() {
 				Value: "/tmp/tracee",
 				Usage: "path where tracee will install or lookup it's resources",
 			},
-			&cli.BoolFlag{
-				Name:  server.MetricsEndpointFlag,
-				Usage: "enable metrics endpoint",
-				Value: false,
-			},
-			&cli.BoolFlag{
-				Name:  server.HealthzEndpointFlag,
-				Usage: "enable healthz endpoint",
-				Value: false,
-			},
-			&cli.BoolFlag{
-				Name:  server.PProfEndpointFlag,
-				Usage: "enable pprof endpoints",
-				Value: false,
-			},
-			&cli.BoolFlag{
-				Name:  server.PyroscopeAgentFlag,
-				Usage: "enable pyroscope agent",
-				Value: false,
-			},
-			&cli.StringFlag{
-				Name:  server.HTTPListenEndpointFlag,
-				Usage: "listening address of the metrics endpoint server",
-				Value: ":3366",
-			},
+			// &cli.BoolFlag{
+			// 	Name:  server.MetricsEndpointFlag,
+			// 	Usage: "enable metrics endpoint",
+			// 	Value: false,
+			// },
+			// &cli.BoolFlag{
+			// 	Name:  server.HealthzEndpointFlag,
+			// 	Usage: "enable healthz endpoint",
+			// 	Value: false,
+			// },
+			// &cli.BoolFlag{
+			// 	Name:  server.PProfEndpointFlag,
+			// 	Usage: "enable pprof endpoints",
+			// 	Value: false,
+			// },
+			// &cli.BoolFlag{
+			// 	Name:  server.PyroscopeAgentFlag,
+			// 	Usage: "enable pyroscope agent",
+			// 	Value: false,
+			// },
+			// &cli.StringFlag{
+			// 	Name:  server.HTTPListenEndpointFlag,
+			// 	Usage: "listening address of the metrics endpoint server",
+			// 	Value: ":3366",
+			// },
 			&cli.BoolFlag{
 				Name:  "no-containers",
 				Usage: "disable container info enrichment to events. safeguard option.",

diff --git a/cmd/tracee-rules/main.go b/cmd/tracee-rules/main.go
@@ -139,13 +139,13 @@ func main() {
 				return fmt.Errorf("constructing engine: %w", err)
 			}
 
-			httpServer, err := server.PrepareHTTPServer(
-				c.String(server.HTTPListenEndpointFlag),
-				c.Bool(server.MetricsEndpointFlag),
-				c.Bool(server.HealthzEndpointFlag),
-				c.Bool(server.PProfEndpointFlag),
-				c.Bool(server.PyroscopeAgentFlag),
-			)
+			// httpServer, err := server.PrepareHTTPServer(
+			// 	c.String(server.HTTPListenEndpointFlag),
+			// 	c.Bool(server.MetricsEndpointFlag),
+			// 	c.Bool(server.HealthzEndpointFlag),
+			// 	c.Bool(server.PProfEndpointFlag),
+			// 	c.Bool(server.PyroscopeAgentFlag),
+			// )
 			if err != nil {
 				return err
 			}
@@ -158,9 +158,9 @@ func main() {
 			ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
 			defer stop()
 
-			if httpServer != nil {
-				go httpServer.Start(ctx)
-			}
+			// if httpServer != nil {
+			// 	go httpServer.Start(ctx)
+			// }
 
 			e.Start(ctx)
 
@@ -218,21 +218,21 @@ func main() {
 				Usage: "size of the event channel's buffer consumed by signatures",
 				Value: 1000,
 			},
-			&cli.BoolFlag{
-				Name:  server.MetricsEndpointFlag,
-				Usage: "enable metrics endpoint",
-				Value: false,
-			},
-			&cli.BoolFlag{
-				Name:  server.HealthzEndpointFlag,
-				Usage: "enable healthz endpoint",
-				Value: false,
-			},
-			&cli.StringFlag{
-				Name:  server.HTTPListenEndpointFlag,
-				Usage: "listening address of the metrics endpoint server",
-				Value: ":4466",
-			},
+			// &cli.BoolFlag{
+			// 	Name:  server.MetricsEndpointFlag,
+			// 	Usage: "enable metrics endpoint",
+			// 	Value: false,
+			// },
+			// &cli.BoolFlag{
+			// 	Name:  server.HealthzEndpointFlag,
+			// 	Usage: "enable healthz endpoint",
+			// 	Value: false,
+			// },
+			// &cli.StringFlag{
+			// 	Name:  server.HTTPListenEndpointFlag,
+			// 	Usage: "listening address of the metrics endpoint server",
+			// 	Value: ":4466",
+			// },
 			&cli.BoolFlag{
 				Name:  "allcaps",
 				Value: false,

diff --git a/cmd/tracee/cmd/root.go b/cmd/tracee/cmd/root.go
@@ -12,7 +12,6 @@ import (
 	"github.com/spf13/viper"
 
 	cmdcobra "github.com/aquasecurity/tracee/pkg/cmd/cobra"
-	"github.com/aquasecurity/tracee/pkg/cmd/flags/server"
 	"github.com/aquasecurity/tracee/pkg/cmd/initialize"
 	"github.com/aquasecurity/tracee/pkg/errfmt"
 	"github.com/aquasecurity/tracee/pkg/logger"
@@ -250,62 +249,13 @@ func initCmd() error {
 
 	// Server flags
 
-	rootCmd.Flags().Bool(
-		server.MetricsEndpointFlag,
-		false,
-		"\t\t\t\t\tEnable metrics endpoint",
-	)
-	err = viper.BindPFlag(server.MetricsEndpointFlag, rootCmd.Flags().Lookup(server.MetricsEndpointFlag))
-	if err != nil {
-		return errfmt.WrapError(err)
-	}
-
-	rootCmd.Flags().Bool(
-		server.HealthzEndpointFlag,
-		false,
-		"\t\t\t\t\tEnable healthz endpoint",
-	)
-	err = viper.BindPFlag(server.HealthzEndpointFlag, rootCmd.Flags().Lookup(server.HealthzEndpointFlag))
-	if err != nil {
-		return errfmt.WrapError(err)
-	}
-
-	rootCmd.Flags().Bool(
-		server.PProfEndpointFlag,
-		false,
-		"\t\t\t\t\tEnable pprof endpoints",
-	)
-	err = viper.BindPFlag(server.PProfEndpointFlag, rootCmd.Flags().Lookup(server.PProfEndpointFlag))
-	if err != nil {
-		return errfmt.WrapError(err)
-	}
-
-	rootCmd.Flags().Bool(
-		server.PyroscopeAgentFlag,
-		false,
-		"\t\t\t\t\tEnable pyroscope agent",
-	)
-	err = viper.BindPFlag(server.PyroscopeAgentFlag, rootCmd.Flags().Lookup(server.PyroscopeAgentFlag))
-	if err != nil {
-		return errfmt.WrapError(err)
-	}
-
-	rootCmd.Flags().String(
-		server.HTTPListenEndpointFlag,
-		":3366",
-		"<url:port>\t\t\t\tListening address of the metrics endpoint server",
+	rootCmd.Flags().StringArray(
+		"server",
+		[]string{""},
+		"Server configuration options (e.g., --server http.metrics=true)", //add help comments
 	)
-	err = viper.BindPFlag(server.HTTPListenEndpointFlag, rootCmd.Flags().Lookup(server.HTTPListenEndpointFlag))
-	if err != nil {
-		return errfmt.WrapError(err)
-	}
 
-	rootCmd.Flags().String(
-		server.GRPCListenEndpointFlag,
-		"", // disabled by default
-		"<protocol:addr>\t\t\tListening address of the grpc server eg: tcp:4466, unix:/tmp/tracee.sock (default: disabled)",
-	)
-	err = viper.BindPFlag(server.GRPCListenEndpointFlag, rootCmd.Flags().Lookup(server.GRPCListenEndpointFlag))
+	err = viper.BindPFlag("server", rootCmd.Flags().Lookup("server"))
 	if err != nil {
 		return errfmt.WrapError(err)
 	}

diff --git a/pkg/cmd/cobra/cobra.go b/pkg/cmd/cobra/cobra.go
@@ -292,24 +292,17 @@ func GetTraceeRunner(c *cobra.Command, version string) (cmd.Runner, error) {
 
 	// Prepare the server
 
-	httpServer, err := server.PrepareHTTPServer(
-		viper.GetString(server.HTTPListenEndpointFlag),
-		viper.GetBool(server.MetricsEndpointFlag),
-		viper.GetBool(server.HealthzEndpointFlag),
-		viper.GetBool(server.PProfEndpointFlag),
-		viper.GetBool(server.PyroscopeAgentFlag),
-	)
+	serverFlag, err := GetFlagsFromViper("server")
 	if err != nil {
 		return runner, err
 	}
-
-	grpcServer, err := flags.PrepareGRPCServer(viper.GetString(server.GRPCListenEndpointFlag))
+	server, err := server.PrepareServer(serverFlag)
 	if err != nil {
 		return runner, err
 	}
 
-	runner.HTTPServer = httpServer
-	runner.GRPCServer = grpcServer
+	runner.HTTPServer = server.HTTPServer
+	runner.GRPCServer = server.GRPCServer
 	runner.TraceeConfig = cfg
 	runner.Printer = p
 	runner.InstallPath = traceeInstallPath

diff --git a/pkg/cmd/cobra/config.go b/pkg/cmd/cobra/config.go
@@ -22,6 +22,8 @@ func GetFlagsFromViper(key string) ([]string, error) {
 	rawValue := viper.Get(key)
 
 	switch key {
+	case "server":
+		flagger = &ServerConfig{}
 	case "cache":
 		flagger = &CacheConfig{}
 	case "proctree":
@@ -122,6 +124,50 @@ func getCRIConfigFlags(rawValue interface{}) ([]string, error) {
 	return flags, nil
 }
 
+//
+// server flag
+//
+
+type ServerConfig struct {
+	Http HttpConfig `mapstructure:"http"`
+	Grpc GrpcConfig `mapstructure:"grpc"`
+}
+type HttpConfig struct {
+	Metrics   bool   `mapstructure:"metrics"`
+	Pprof     bool   `mapstructure:"pprof"`
+	Healthz   bool   `mapstructure:"healthz"`
+	Pyroscope bool   `mapstructure:"pyroscope"`
+	Address   string `mapstructure:"address"`
+}
+
+type GrpcConfig struct {
+	Address string `mapstructure:"address"`
+}
+
+func (s *ServerConfig) flags() []string {
+	flags := make([]string, 0)
+
+	if s.Grpc.Address != "" {
+		flags = append(flags, fmt.Sprintf("grpc.address=%s", s.Grpc.Address))
+	}
+	if s.Http.Address != "" {
+		flags = append(flags, fmt.Sprintf("http.address=%s", s.Http.Address))
+	}
+	if s.Http.Metrics {
+		flags = append(flags, "http.metrics=true")
+	}
+	if s.Http.Pprof {
+		flags = append(flags, "http.pprof=true")
+	}
+	if s.Http.Healthz {
+		flags = append(flags, "http.healthz=true")
+	}
+	if s.Http.Pyroscope {
+		flags = append(flags, "http.pyroscope=true")
+	}
+	return flags
+}
+
 //
 // config flag
 //