modified code to align with new status flag code, fix minor issues

aquasecurity · Dec 4, 2024 · 21986f0 · 21986f0
1 parent 866bca3
commit 21986f0
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 65 deletions.
diff --git a/cmd/traceectl/cmd/event.go b/cmd/traceectl/cmd/event.go
@@ -15,9 +15,9 @@ var eventCmd = &cobra.Command{
 	Use:   "event [command]",
 	Short: "Event management for tracee",
 	Long: `Event Management for tracee 
-	Let you enable and disable events in tracee.
-	Get descriptions of events.
-	`,
+Let you enable and disable events in tracee.
+Get descriptions of events.
+`,
 	Args: cobra.MinimumNArgs(1),
 	Run: func(cmd *cobra.Command, args []string) {
 		if len(args) == 0 {
@@ -76,7 +76,7 @@ var disableEventCmd = &cobra.Command{
 
 func listEvents(cmd *cobra.Command, args []string) {
 	var traceeClient client.ServiceClient
-	if err := traceeClient.NewServiceClient(serverInfo); err != nil {
+	if err := traceeClient.NewServiceClient(server); err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
 		return
 	}
@@ -109,7 +109,7 @@ func listEvents(cmd *cobra.Command, args []string) {
 
 func getEventDescriptions(cmd *cobra.Command, args []string) {
 	var traceeClient client.ServiceClient
-	if err := traceeClient.NewServiceClient(serverInfo); err != nil {
+	if err := traceeClient.NewServiceClient(server); err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
 		return
 	}
@@ -152,7 +152,7 @@ func prepareDescription(event *pb.EventDefinition) []string {
 }
 func enableEvents(cmd *cobra.Command, eventName string) {
 	var traceeClient client.ServiceClient
-	if err := traceeClient.NewServiceClient(serverInfo); err != nil {
+	if err := traceeClient.NewServiceClient(server); err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
 		return
 	}
@@ -165,7 +165,7 @@ func enableEvents(cmd *cobra.Command, eventName string) {
 }
 func disableEvents(cmd *cobra.Command, eventName string) {
 	var traceeClient client.ServiceClient
-	if err := traceeClient.NewServiceClient(serverInfo); err != nil {
+	if err := traceeClient.NewServiceClient(server); err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
 		return
 	}

diff --git a/cmd/traceectl/cmd/root.go b/cmd/traceectl/cmd/root.go
@@ -11,20 +11,28 @@ import (
 	"github.com/spf13/cobra"
 )
 
-var formatFlag string
-var outputFlag string
 var (
-	serverInfo client.ServerInfo = client.ServerInfo{
+	formatFlag string
+	outputFlag string
+	server     client.ServerInfo = client.ServerInfo{
 		ConnectionType: client.Protocol_UNIX,
 		Addr:           client.Socket,
 	}
+)
 
+var (
 	rootCmd = &cobra.Command{
 		Use:   "traceectl [flags] [command]",
 		Short: "TraceeCtl is a CLI tool for tracee",
-		Long:  "TraceeCtl is the client for the tracee API server.",
+		Long: `TraceeCtl is a CLI toll for tracee:
+This tool allows you to mange event, stream events directly from tracee, and get info about tracee.
+`,
 		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
-			if err := flags.PrepareOutput(cmd, outputFlag); err != nil {
+			var err error
+			if err = flags.PrepareOutput(cmd, outputFlag); err != nil {
+				return err
+			}
+			if server, err = flags.PrepareServer(cmd, server); err != nil {
 				return err
 			}
 			return nil
@@ -41,7 +49,7 @@ func init() {
 	rootCmd.AddCommand(metricsCmd)
 	rootCmd.AddCommand(versionCmd)
 
-	rootCmd.PersistentFlags().StringVar(&serverInfo.Addr, "server", client.Socket, `Server connection path or address.
+	rootCmd.PersistentFlags().StringVar(&server.Addr, "server", client.Socket, `Server connection path or address.
 	for unix socket <socket_path> (default: /tmp/tracee.sock)
 	for tcp <IP:Port>`)
 	rootCmd.PersistentFlags().StringVarP(&outputFlag, "output", "o", "", "Specify the output format")
@@ -73,7 +81,7 @@ func Execute() {
 
 func displayMetrics(cmd *cobra.Command, _ []string) {
 	var traceeClient client.DiagnosticClient
-	if err := traceeClient.NewDiagnosticClient(serverInfo); err != nil {
+	if err := traceeClient.NewDiagnosticClient(server); err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
 		return
 	}
@@ -97,7 +105,7 @@ func displayMetrics(cmd *cobra.Command, _ []string) {
 
 func displayVersion(cmd *cobra.Command, _ []string) {
 	var traceeClient client.ServiceClient
-	if err := traceeClient.NewServiceClient(serverInfo); err != nil {
+	if err := traceeClient.NewServiceClient(server); err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
 		return
 	}

diff --git a/cmd/traceectl/cmd/stream.go b/cmd/traceectl/cmd/stream.go
@@ -14,8 +14,8 @@ var streamCmd = &cobra.Command{
 	Use:   "stream [policies...]",
 	Short: "Stream events from tracee",
 	Long: `Stream Management:
-	
-	`,
+stream event directly from tracee to the output preferred
+`,
 	Run: func(cmd *cobra.Command, args []string) {
 		streamEvents(cmd, args)
 	},
@@ -28,7 +28,7 @@ func init() {
 func streamEvents(cmd *cobra.Command, args []string) {
 	//connect to tracee
 	var traceeClient client.ServiceClient
-	err := traceeClient.NewServiceClient(serverInfo)
+	err := traceeClient.NewServiceClient(server)
 	if err != nil {
 		cmd.PrintErrln("Error creating client: ", err)
 		return
@@ -57,7 +57,7 @@ func streamEvents(cmd *cobra.Command, args []string) {
 				if err.Error() == "EOF" {
 					break
 				}
-				//Error receiving streamed event
+				cmd.PrintErrln("Error receiving streamed event")
 			}
 			format.PrintJson(res.Event.String())
 		}
@@ -70,9 +70,8 @@ func streamEvents(cmd *cobra.Command, args []string) {
 				if err.Error() == "EOF" {
 					break
 				}
-				//Error receiving streamed event
+				cmd.PrintErrln("Error receiving streamed event")
 			}
-
 			format.PrintTableRow(prepareEvent(res.Event))
 		}
 	default:

diff --git a/cmd/traceectl/pkg/client/client.go b/cmd/traceectl/pkg/client/client.go
@@ -1,10 +1,7 @@
 package client
 
 import (
-	"fmt"
 	"log"
-	"net"
-	"strings"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
@@ -21,53 +18,14 @@ type ServerInfo struct {
 	Addr           string
 }
 
-func connectToServer(serverInfo ServerInfo) (*grpc.ClientConn, error) {
+func connectToServer(server ServerInfo) (*grpc.ClientConn, error) {
 	var opts []grpc.DialOption
 	var conn *grpc.ClientConn
-	var err error
-	var address string
 	opts = append(opts, grpc.WithTransportCredentials(insecure.NewCredentials()))
-	err = determineConnectionType(serverInfo)
-	if err != nil {
-		return nil, err
-	}
-	switch serverInfo.ConnectionType {
-	case Protocol_UNIX:
-		address = fmt.Sprintf("unix://%s", serverInfo.Addr)
-	case Protocol_TCP:
-		address = fmt.Sprintf(serverInfo.Addr)
-	default:
-		return nil, fmt.Errorf("unsupported connection type: %s", serverInfo.ConnectionType)
-	}
-	conn, err = grpc.NewClient(address, opts...)
+	conn, err := grpc.NewClient(server.Addr, opts...)
 	if err != nil {
 		log.Fatalf("failed to connect to server: %v", err)
 		return nil, err
 	}
 	return conn, nil
 }
-
-func determineConnectionType(serverInfo ServerInfo) error {
-	if strings.Contains(serverInfo.Addr, ":") && isValidTCPAddress(serverInfo.Addr) {
-		serverInfo.ConnectionType = Protocol_TCP
-		return nil
-	}
-	if strings.HasPrefix(serverInfo.Addr, "/") {
-		serverInfo.ConnectionType = Protocol_UNIX
-		return nil
-	}
-
-	return fmt.Errorf("unsupported connection type: %s", serverInfo.Addr)
-
-}
-func isValidTCPAddress(addr string) bool {
-	host, port, err := net.SplitHostPort(addr)
-	if err != nil || host == "" || port == "" {
-		return false
-	}
-	if _, err := net.LookupPort("tcp", port); err != nil {
-		return false
-	}
-
-	return true
-}