From af116f5d746f4cfbd5f29b9093485dbc7724df32 Mon Sep 17 00:00:00 2001 From: ShohamBit Date: Mon, 7 Oct 2024 10:46:48 +0000 Subject: [PATCH 1/2] added disable event command --- client/service.go | 4 +++ cmd/disableEvent.go | 44 +++++++++++++++++++++++++ cmd/disableEvent_test.go | 71 ++++++++++++++++++++++++++++++++++++++++ cmd/root.go | 1 + mock/service_server.go | 3 ++ 5 files changed, 123 insertions(+) create mode 100644 cmd/disableEvent.go create mode 100644 cmd/disableEvent_test.go diff --git a/client/service.go b/client/service.go index 9db7aff..f021093 100644 --- a/client/service.go +++ b/client/service.go @@ -42,6 +42,10 @@ func (tc *serviceClient) EnableEvent(ctx context.Context, req *pb.EnableEventReq return tc.client.EnableEvent(ctx, req) } +func (tc *serviceClient) DisableEvent(ctx context.Context, req *pb.DisableEventRequest) (*pb.DisableEventResponse, error) { + return tc.client.DisableEvent(ctx, req) +} + func (tc *serviceClient) StreamEvents(ctx context.Context, req *pb.StreamEventsRequest) (pb.TraceeService_StreamEventsClient, error) { return tc.client.StreamEvents(ctx, req) } diff --git a/cmd/disableEvent.go b/cmd/disableEvent.go new file mode 100644 index 0000000..78f7eed --- /dev/null +++ b/cmd/disableEvent.go @@ -0,0 +1,44 @@ +package cmd + +import ( + "context" + + "github.com/ShohamBit/TraceeClient/client" + + pb "github.com/aquasecurity/tracee/api/v1beta1" + "github.com/spf13/cobra" +) + +var disableEventCmd = &cobra.Command{ + Use: "disableEvent [event names...]", + Short: "disable specified events on the server", + Long: "long about the use", + Args: cobra.MinimumNArgs(1), + Run: func(cmd *cobra.Command, args []string) { + // Check if args are provided + if len(args) == 0 { + cmd.PrintErrln("Error: no event names provided. Please specify at least one event to disable.") + return // Exit if no arguments + } + disableEvents(cmd, args) + }, +} + +func disableEvents(cmd *cobra.Command, eventNames []string) { + // Create Tracee gRPC client + client, err := client.NewServiceClient(serverInfo) + if err != nil { + cmd.PrintErrln("Error creating client: ", err) + return // Exit on error + } + + // Iterate over event names and disable each one + for _, eventName := range eventNames { + _, err := client.DisableEvent(context.Background(), &pb.DisableEventRequest{Name: eventName}) + if err != nil { + cmd.PrintErrln("Error enabling event:", err) + continue // Continue on error with the next event + } + cmd.Println("Disabled event:", eventName) + } +} diff --git a/cmd/disableEvent_test.go b/cmd/disableEvent_test.go new file mode 100644 index 0000000..2edddc4 --- /dev/null +++ b/cmd/disableEvent_test.go @@ -0,0 +1,71 @@ +package cmd + +import ( + "bytes" + "testing" + "time" + + "github.com/ShohamBit/TraceeClient/mock" + "github.com/ShohamBit/TraceeClient/models" + "github.com/stretchr/testify/assert" +) + +var ( + DisableEventTests = []models.TestCase{ + { + Name: "No events", + Args: []string{"disableEvent"}, + ExpectedOutput: "Error: requires at least 1 arg(s), only received 0\n", // Update expected output + }, + { + Name: "Single event", + Args: []string{"disableEvent", "event1"}, + ExpectedOutput: "Disabled event: event1\n", + }, + { + Name: "Multiple events", + Args: []string{"disableEvent", "event1", "event2"}, + ExpectedOutput: "Disabled event: event1\nDisabled event: event2\n", + }, + } +) + +func TestDisableEvent(t *testing.T) { + // Start the mock server + mockServer, err := mock.StartMockServiceServer() + if err != nil { + t.Fatalf("Failed to start mock server: %v", err) + } + defer mockServer.Stop() // Ensure the server is stopped after the test + + // Wait for the server to start + time.Sleep(100 * time.Millisecond) + + for _, test := range DisableEventTests { + t.Run(test.Name, func(t *testing.T) { + // Capture output + var buf bytes.Buffer + rootCmd.SetOut(&buf) + rootCmd.SetErr(&buf) + + // Set arguments for the test + rootCmd.SetArgs(test.Args) + + // Execute the command + err := rootCmd.Execute() + + // Validate output and error (if any) + output := buf.String() + + // If no arguments provided, we expect an error + if test.Name == "No events" { + assert.Error(t, err) + } else { + assert.NoError(t, err) + } + + // Check if output matches expected output + assert.Contains(t, output, test.ExpectedOutput) + }) + } +} diff --git a/cmd/root.go b/cmd/root.go index 6ffab18..40f33fd 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -27,6 +27,7 @@ func init() { rootCmd.AddCommand(versionCmd) rootCmd.AddCommand(metricsCmd) rootCmd.AddCommand(enableEventCmd) + rootCmd.AddCommand(disableEventCmd) rootCmd.AddCommand(streamEventsCmd) //flags diff --git a/mock/service_server.go b/mock/service_server.go index d90cc40..314e4e6 100644 --- a/mock/service_server.go +++ b/mock/service_server.go @@ -47,6 +47,9 @@ func (s *MockServiceServer) GetVersion(ctx context.Context, req *pb.GetVersionRe func (s *MockServiceServer) EnableEvent(ctx context.Context, req *pb.EnableEventRequest) (*pb.EnableEventResponse, error) { return &pb.EnableEventResponse{}, nil } +func (s *MockServiceServer) DisableEvent(ctx context.Context, req *pb.DisableEventRequest) (*pb.DisableEventResponse, error) { + return &pb.DisableEventResponse{}, nil +} /* \stream events From 3ad8282d634c3e8f44fc4a66f27bcd4d374d54ef Mon Sep 17 00:00:00 2001 From: Shoham Yosef Bitton <141840572+ShohamBit@users.noreply.github.com> Date: Mon, 7 Oct 2024 14:04:02 +0300 Subject: [PATCH 2/2] Update cmd/disableEvent.go --- cmd/disableEvent.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmd/disableEvent.go b/cmd/disableEvent.go index 78f7eed..a47548c 100644 --- a/cmd/disableEvent.go +++ b/cmd/disableEvent.go @@ -36,7 +36,7 @@ func disableEvents(cmd *cobra.Command, eventNames []string) { for _, eventName := range eventNames { _, err := client.DisableEvent(context.Background(), &pb.DisableEventRequest{Name: eventName}) if err != nil { - cmd.PrintErrln("Error enabling event:", err) + cmd.PrintErrln("Error disabling event:", err) continue // Continue on error with the next event } cmd.Println("Disabled event:", eventName)