From f1a24d3fea732314d6a6dedcc7533d57f63148f0 Mon Sep 17 00:00:00 2001 From: ShohamBit Date: Sun, 3 Nov 2024 21:45:20 +0000 Subject: [PATCH 1/6] add base mock servers for each command, fix creation of mock server, edit test files for stream and the files itself --- cmd/root.go | 2 +- cmd/stream.go | 1 + cmd/stream_test.go | 2 + pkg/mock/event_server.go | 14 +++++ pkg/mock/root_server.go | 12 ++++ pkg/mock/server.go | 63 +++++++++++++++++++ .../{service_server.go => stream_server.go} | 46 -------------- 7 files changed, 93 insertions(+), 47 deletions(-) create mode 100644 pkg/mock/event_server.go create mode 100644 pkg/mock/root_server.go create mode 100644 pkg/mock/server.go rename pkg/mock/{service_server.go => stream_server.go} (67%) diff --git a/cmd/root.go b/cmd/root.go index 6b41d07..9d9349f 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -56,7 +56,7 @@ func init() { //unix connection type flag rootCmd.PersistentFlags().StringVar(&serverInfo.UnixSocketPath, "socketPath", client.SOCKET, "Path of the unix socket") //tcp connection type flag - rootCmd.PersistentFlags().StringVarP(&serverInfo.ADDR, "server", "s", client.DefaultIP+":"+client.DefaultPort, "he address and port of the Kubernetes API server") + rootCmd.PersistentFlags().StringVarP(&serverInfo.ADDR, "server", "s", client.DefaultIP+":"+client.DefaultPort, "The address and port of the Kubernetes API server") } diff --git a/cmd/stream.go b/cmd/stream.go index cd3af87..ac7bc3c 100644 --- a/cmd/stream.go +++ b/cmd/stream.go @@ -13,6 +13,7 @@ var streamCmd = &cobra.Command{ Use: "stream [policies...]", Short: "Stream events from tracee", Long: `Stream Management: + - traceectl stream [POLICIES...] - stream event directly from tracee - traceectl stream create --name [--destination ] [--format ] [--fields ] [--parse-data] [--filter ] - traceectl stream describe - traceectl stream list diff --git a/cmd/stream_test.go b/cmd/stream_test.go index 75f7ff9..2d3e3ab 100644 --- a/cmd/stream_test.go +++ b/cmd/stream_test.go @@ -11,6 +11,8 @@ import ( pb "github.com/aquasecurity/tracee/api/v1beta1" ) +// test run the mock server and start the stream command +// currently stream can connect to the server and print the output of events to the stream var streamTests = []models.TestCase{ { Name: "No stream subcommand", diff --git a/pkg/mock/event_server.go b/pkg/mock/event_server.go new file mode 100644 index 0000000..6e842e3 --- /dev/null +++ b/pkg/mock/event_server.go @@ -0,0 +1,14 @@ +package mock + +import ( + "context" + + pb "github.com/aquasecurity/tracee/api/v1beta1" +) + +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 +} diff --git a/pkg/mock/root_server.go b/pkg/mock/root_server.go new file mode 100644 index 0000000..11e5994 --- /dev/null +++ b/pkg/mock/root_server.go @@ -0,0 +1,12 @@ +package mock + +import ( + "context" + + pb "github.com/aquasecurity/tracee/api/v1beta1" +) + +func (s *MockServiceServer) GetVersion(ctx context.Context, req *pb.GetVersionRequest) (*pb.GetVersionResponse, error) { + // Return a mock version response + return &pb.GetVersionResponse{Version: ExpectedVersion}, nil +} diff --git a/pkg/mock/server.go b/pkg/mock/server.go new file mode 100644 index 0000000..49d16bd --- /dev/null +++ b/pkg/mock/server.go @@ -0,0 +1,63 @@ +package mock + +import ( + "fmt" + "net" + "os" + + "github.com/ShohamBit/traceectl/pkg/client" + pb "github.com/aquasecurity/tracee/api/v1beta1" + + "google.golang.org/grpc" +) + +var ( + ExpectedVersion string = "v0.22.0-15-gd09d7fca0d" // Match the output format + serverInfo client.ServerInfo = client.ServerInfo{ + ADDR: client.DefaultIP + ":" + client.DefaultPort, + UnixSocketPath: client.SOCKET, + } +) + +// MockServiceServer implements the gRPC server interface for testing +type MockServiceServer struct { + pb.UnimplementedTraceeServiceServer // Embed the unimplemented server +} + +// CreateMockServer initializes the gRPC server and binds it to a Unix socket listener +func CreateMockServer() (*grpc.Server, net.Listener, error) { + //check for unix socket + if _, err := os.Stat(serverInfo.UnixSocketPath); err == nil { + err := os.Remove(serverInfo.UnixSocketPath) + if err != nil { + return nil, nil, fmt.Errorf("failed to cleanup gRPC listening address (%s): %v", serverInfo.UnixSocketPath, err) + } + } + + // Create the Unix socket listener + listener, err := net.Listen("unix", serverInfo.UnixSocketPath) + if err != nil { + return nil, nil, fmt.Errorf("failed to create Unix socket listener: %v", err) + } + + // Create a new gRPC server + server := grpc.NewServer() + + return server, listener, nil +} + +func StartMockServiceServer() (*grpc.Server, error) { + mockServer, listener, err := CreateMockServer() + if err != nil { + return nil, fmt.Errorf("failed to create mock server: %v", err) + } + pb.RegisterTraceeServiceServer(mockServer, &MockServiceServer{}) + + // Start serving in a goroutine + go func() { + if err := mockServer.Serve(listener); err != nil { + fmt.Printf("gRPC server failed: %v\n", err) + } + }() + return mockServer, nil +} diff --git a/pkg/mock/service_server.go b/pkg/mock/stream_server.go similarity index 67% rename from pkg/mock/service_server.go rename to pkg/mock/stream_server.go index fdc5fb1..485a194 100644 --- a/pkg/mock/service_server.go +++ b/pkg/mock/stream_server.go @@ -1,59 +1,13 @@ package mock import ( - "context" - "net" "sort" "strings" "time" - "github.com/ShohamBit/traceectl/pkg/client" pb "github.com/aquasecurity/tracee/api/v1beta1" - "google.golang.org/grpc" ) -var ( - ExpectedVersion string = "v0.22.0-15-gd09d7fca0d" // Match the output format - serverInfo client.ServerInfo = client.ServerInfo{ - ADDR: client.DefaultIP + ":" + client.DefaultPort, - } -) - -// MockServiceServer implements the gRPC server interface for testing -type MockServiceServer struct { - pb.UnimplementedTraceeServiceServer // Embed the unimplemented server -} - -// Start a mock gRPC server -func StartMockServiceServer() (*grpc.Server, error) { - lis, err := net.Listen("tcp", serverInfo.ADDR) - if err != nil { - return nil, err - } - - s := grpc.NewServer() - pb.RegisterTraceeServiceServer(s, &MockServiceServer{}) - - go func() { - if err := s.Serve(lis); err != nil { - // Handle the error (e.g., log it) - } - }() - - return s, nil -} - -func (s *MockServiceServer) GetVersion(ctx context.Context, req *pb.GetVersionRequest) (*pb.GetVersionResponse, error) { - // Return a mock version response - return &pb.GetVersionResponse{Version: ExpectedVersion}, nil -} -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 78bf131ecb0bb116ecefd942d3c660323285844d Mon Sep 17 00:00:00 2001 From: ShohamBit Date: Mon, 4 Nov 2024 10:00:41 +0000 Subject: [PATCH 2/6] change formattion and print design --- cmd/stream.go | 58 +-- pkg/cmd/formatter/JSONFormatter.go | 5 +- pkg/cmd/formatter/TableFormatter.go | 22 +- pkg/cmd/formatter/formatter.go | 44 ++ pkg/cmd/printer/printer.go | 705 +++------------------------- 5 files changed, 113 insertions(+), 721 deletions(-) create mode 100644 pkg/cmd/formatter/formatter.go diff --git a/cmd/stream.go b/cmd/stream.go index ac7bc3c..b17f86e 100644 --- a/cmd/stream.go +++ b/cmd/stream.go @@ -1,9 +1,10 @@ package cmd import ( + "github.com/ShohamBit/traceectl/pkg/cmd/formatter" + "github.com/ShohamBit/traceectl/pkg/cmd/printer" pb "github.com/aquasecurity/tracee/api/v1beta1" - "github.com/ShohamBit/traceectl/pkg/cmd/formatter" "github.com/spf13/cobra" ) @@ -118,6 +119,7 @@ var resumeStreamCmd = &cobra.Command{ }, } +// stream events directly from tracee func stream(cmd *cobra.Command, args []string) { // Create service client err := TCS.NewServiceClient(serverInfo) @@ -133,54 +135,12 @@ func stream(cmd *cobra.Command, args []string) { cmd.PrintErrln("Error calling Stream: ", err) } - //add check for the output flag - //TODO:support only table and json format for now - switch formatFlag { - case "json": - jsonStreamEvents(cmd, args, stream) - case "table": - tableStreamEvents(cmd, args, stream) - case "template": // go template - fallthrough - default: - cmd.PrintErrln("Error: output format not supported") - return - } -} - -// tableStreamEvents prints events in a table format -func tableStreamEvents(cmd *cobra.Command, _ []string, stream pb.TraceeService_StreamEventsClient) { - // Init table header before streaming starts - tbl := formatter.New(formatFlag, outputFlag, cmd) - tbl.PrintTableHeaders() - // Receive and process streamed responses - for { - res, err := stream.Recv() - if err != nil { - // Handle the error that occurs when the server closes the stream - if err.Error() == "EOF" { - break - } - cmd.PrintErrln("Error receiving streamed event: ", err) - } - tbl.PrintTableRow(res.Event) - + //create formatter for output + format, err := formatter.New(formatFlag, outputFlag, cmd) + if err != nil { + cmd.PrintErrln("Error creating formatter: ", err) } -} + //show events + printer.StreamEvents(format, args, stream) -// jsonStreamEvents prints events in json format -func jsonStreamEvents(cmd *cobra.Command, _ []string, stream pb.TraceeService_StreamEventsClient) { - // Receive and process streamed responses - for { - res, err := stream.Recv() - if err != nil { - // Handle the error that occurs when the server closes the stream - if err.Error() == "EOF" { - break - } - cmd.PrintErrln("Error receiving streamed event: ", err) - } - // Print each event as a row in json format - formatter.PrintJSON(cmd, res.Event, outputFlag) - } } diff --git a/pkg/cmd/formatter/JSONFormatter.go b/pkg/cmd/formatter/JSONFormatter.go index 5c8dd58..b948a33 100644 --- a/pkg/cmd/formatter/JSONFormatter.go +++ b/pkg/cmd/formatter/JSONFormatter.go @@ -2,11 +2,10 @@ package formatter import ( pb "github.com/aquasecurity/tracee/api/v1beta1" - "github.com/spf13/cobra" ) // PrintJSON prints an event in JSON format -func PrintJSON(cmd *cobra.Command, event *pb.Event, _ string) { +func (f *Formatter) PrintJSON(event *pb.Event) { //TODO: add more output formats - cmd.Printf("%s\n", event.String()) + f.CMD.Printf("%s\n", event.String()) } diff --git a/pkg/cmd/formatter/TableFormatter.go b/pkg/cmd/formatter/TableFormatter.go index 875b8b1..23d3507 100644 --- a/pkg/cmd/formatter/TableFormatter.go +++ b/pkg/cmd/formatter/TableFormatter.go @@ -5,24 +5,10 @@ import ( "strings" pb "github.com/aquasecurity/tracee/api/v1beta1" - "github.com/spf13/cobra" ) -type Formatter struct { - format string - output string - cmd *cobra.Command -} - -func New(format string, output string, cmd *cobra.Command) *Formatter { - return &Formatter{ - format: format, - output: output, - cmd: cmd, - } -} func (f *Formatter) PrintTableHeaders() { - f.cmd.Printf("%-15s %-25s %-15s %-15s %s\n", + f.CMD.Printf("%-15s %-25s %-15s %-15s %s\n", "TIME", "EVENT NAME", "POLICIES", @@ -33,7 +19,7 @@ func (f *Formatter) PrintTableHeaders() { func (f *Formatter) PrintTableRow(event *pb.Event) { timestamp := event.Timestamp.AsTime().Format("15:04:05.000") - f.cmd.Printf("%-15s %-25s %-15s %-15s %s\n", + f.CMD.Printf("%-15s %-25s %-15s %-15s %s\n", timestamp, event.Name, strings.Join(event.Policies.Matched, ","), @@ -43,10 +29,6 @@ func (f *Formatter) PrintTableRow(event *pb.Event) { } -// func getEventContext(context *pb.Context) string { -// return " " -// } - // generate event data func getEventData(data []*pb.EventValue) string { var result []string diff --git a/pkg/cmd/formatter/formatter.go b/pkg/cmd/formatter/formatter.go new file mode 100644 index 0000000..8ece362 --- /dev/null +++ b/pkg/cmd/formatter/formatter.go @@ -0,0 +1,44 @@ +package formatter + +import ( + "fmt" + + "github.com/spf13/cobra" +) + +const ( + FormatJSON = "json" + FormatTable = "table" + FormatGoTpl = "gotemplate" +) + +// SupportedFormats is a slice of all supported format types +var SupportedFormats = []string{FormatJSON, FormatTable, FormatGoTpl} + +type Formatter struct { + Format string + Output string + CMD *cobra.Command +} + +func New(format string, output string, cmd *cobra.Command) (*Formatter, error) { + if !containsFormat(format) { + return nil, fmt.Errorf("format %s is not supported", format) + + } + return &Formatter{ + Format: format, + Output: output, + CMD: cmd, + }, nil +} + +// containsFormat checks if a format is in the SupportedFormats slice +func containsFormat(format string) bool { + for _, f := range SupportedFormats { + if f == format { + return true + } + } + return false +} diff --git a/pkg/cmd/printer/printer.go b/pkg/cmd/printer/printer.go index b1e56ab..13c2bac 100644 --- a/pkg/cmd/printer/printer.go +++ b/pkg/cmd/printer/printer.go @@ -1,651 +1,58 @@ package printer -// import ( -// "bytes" -// "encoding/json" -// "fmt" -// "io" -// "net/http" -// "net/url" -// "path/filepath" -// "strconv" -// "strings" -// "text/template" -// "time" - -// forward "github.com/IBM/fluent-forward-go/fluent/client" -// "github.com/Masterminds/sprig/v3" - -// "github.com/aquasecurity/tracee/pkg/config" -// "github.com/aquasecurity/tracee/pkg/errfmt" -// "github.com/aquasecurity/tracee/pkg/logger" -// "github.com/aquasecurity/tracee/pkg/metrics" -// "github.com/aquasecurity/tracee/types/trace" -// ) - -// type EventPrinter interface { -// // Init serves as the initializer method for every event Printer type -// Init() error -// // Preamble prints something before event printing begins (one time) -// Preamble() -// // Epilogue prints something after event printing ends (one time) -// Epilogue(stats metrics.Stats) -// // Print prints a single event -// Print(event trace.Event) -// // dispose of resources -// Close() -// } - -// func New(cfg config.PrinterConfig) (EventPrinter, error) { -// var res EventPrinter -// kind := cfg.Kind - -// if cfg.OutFile == nil { -// return res, errfmt.Errorf("out file is not set") -// } - -// switch { -// case kind == "ignore": -// res = &ignoreEventPrinter{} -// case kind == "table": -// res = &tableEventPrinter{ -// out: cfg.OutFile, -// verbose: false, -// containerMode: cfg.ContainerMode, -// } -// case kind == "table-verbose": -// res = &tableEventPrinter{ -// out: cfg.OutFile, -// verbose: true, -// containerMode: cfg.ContainerMode, -// } -// case kind == "json": -// res = &jsonEventPrinter{ -// out: cfg.OutFile, -// } -// case kind == "forward": -// res = &forwardEventPrinter{ -// outPath: cfg.OutPath, -// } -// case kind == "webhook": -// res = &webhookEventPrinter{ -// outPath: cfg.OutPath, -// } -// case strings.HasPrefix(kind, "gotemplate="): -// res = &templateEventPrinter{ -// out: cfg.OutFile, -// templatePath: strings.Split(kind, "=")[1], -// } -// } -// err := res.Init() -// if err != nil { -// return nil, err -// } -// return res, nil -// } - -// type tableEventPrinter struct { -// out io.WriteCloser -// verbose bool -// containerMode config.ContainerMode -// relativeTS bool -// } - -// func (p tableEventPrinter) Init() error { return nil } - -// func (p tableEventPrinter) Preamble() { -// if p.verbose { -// switch p.containerMode { -// case config.ContainerModeDisabled: -// fmt.Fprintf(p.out, -// "%-16s %-17s %-13s %-12s %-12s %-6s %-16s %-7s %-7s %-7s %-16s %-25s %s", -// "TIME", -// "UTS_NAME", -// "CONTAINER_ID", -// "MNT_NS", -// "PID_NS", -// "UID", -// "COMM", -// "PID", -// "TID", -// "PPID", -// "RET", -// "EVENT", -// "ARGS", -// ) -// case config.ContainerModeEnabled: -// fmt.Fprintf(p.out, -// "%-16s %-17s %-13s %-12s %-12s %-6s %-16s %-15s %-15s %-15s %-16s %-25s %s", -// "TIME", -// "UTS_NAME", -// "CONTAINER_ID", -// "MNT_NS", -// "PID_NS", -// "UID", -// "COMM", -// "PID/host", -// "TID/host", -// "PPID/host", -// "RET", -// "EVENT", -// "ARGS", -// ) -// case config.ContainerModeEnriched: -// fmt.Fprintf(p.out, -// "%-16s %-17s %-13s %-16s %-12s %-12s %-6s %-16s %-15s %-15s %-15s %-16s %-25s %s", -// "TIME", -// "UTS_NAME", -// "CONTAINER_ID", -// "IMAGE", -// "MNT_NS", -// "PID_NS", -// "UID", -// "COMM", -// "PID/host", -// "TID/host", -// "PPID/host", -// "RET", -// "EVENT", -// "ARGS", -// ) -// } -// } else { -// switch p.containerMode { -// case config.ContainerModeDisabled: -// fmt.Fprintf(p.out, -// "%-16s %-6s %-16s %-7s %-7s %-16s %-25s %s", -// "TIME", -// "UID", -// "COMM", -// "PID", -// "TID", -// "RET", -// "EVENT", -// "ARGS", -// ) -// case config.ContainerModeEnabled: -// fmt.Fprintf(p.out, -// "%-16s %-13s %-6s %-16s %-15s %-15s %-16s %-25s %s", -// "TIME", -// "CONTAINER_ID", -// "UID", -// "COMM", -// "PID/host", -// "TID/host", -// "RET", -// "EVENT", -// "ARGS", -// ) -// case config.ContainerModeEnriched: -// fmt.Fprintf(p.out, -// "%-16s %-13s %-16s %-6s %-16s %-15s %-15s %-16s %-25s %s", -// "TIME", -// "CONTAINER_ID", -// "IMAGE", -// "UID", -// "COMM", -// "PID/host", -// "TID/host", -// "RET", -// "EVENT", -// "ARGS", -// ) -// } -// } -// fmt.Fprintln(p.out) -// } - -// func (p tableEventPrinter) Print(event trace.Event) { -// ut := time.Unix(0, int64(event.Timestamp)) -// if p.relativeTS { -// ut = ut.UTC() -// } -// timestamp := fmt.Sprintf("%02d:%02d:%02d:%06d", ut.Hour(), ut.Minute(), ut.Second(), ut.Nanosecond()/1000) - -// containerId := event.Container.ID -// if len(containerId) > 12 { -// containerId = containerId[:12] -// } -// containerImage := event.Container.ImageName -// if len(containerImage) > 16 { -// containerImage = containerImage[:16] -// } - -// eventName := event.EventName -// if len(eventName) > 25 { -// eventName = eventName[:22] + "..." -// } - -// if p.verbose { -// switch p.containerMode { -// case config.ContainerModeDisabled: -// fmt.Fprintf(p.out, -// "%-16s %-17s %-13s %-12d %-12d %-6d %-16s %-7d %-7d %-7d %-16d %-25s ", -// timestamp, -// event.HostName, -// containerId, -// event.MountNS, -// event.PIDNS, -// event.UserID, -// event.ProcessName, -// event.ProcessID, -// event.ThreadID, -// event.ParentProcessID, -// event.ReturnValue, -// event.EventName, -// ) -// case config.ContainerModeEnabled: -// fmt.Fprintf(p.out, -// "%-16s %-17s %-13s %-12d %-12d %-6d %-16s %-7d/%-7d %-7d/%-7d %-7d/%-7d %-16d %-25s ", -// timestamp, -// event.HostName, -// containerId, -// event.MountNS, -// event.PIDNS, -// event.UserID, -// event.ProcessName, -// event.ProcessID, -// event.HostProcessID, -// event.ThreadID, -// event.HostThreadID, -// event.ParentProcessID, -// event.HostParentProcessID, -// event.ReturnValue, -// event.EventName, -// ) -// case config.ContainerModeEnriched: -// fmt.Fprintf(p.out, -// "%-16s %-17s %-13s %-16s %-12d %-12d %-6d %-16s %-7d/%-7d %-7d/%-7d %-7d/%-7d %-16d %-25s ", -// timestamp, -// event.HostName, -// containerId, -// event.Container.ImageName, -// event.MountNS, -// event.PIDNS, -// event.UserID, -// event.ProcessName, -// event.ProcessID, -// event.HostProcessID, -// event.ThreadID, -// event.HostThreadID, -// event.ParentProcessID, -// event.HostParentProcessID, -// event.ReturnValue, -// event.EventName, -// ) -// } -// } else { -// switch p.containerMode { -// case config.ContainerModeDisabled: -// fmt.Fprintf(p.out, -// "%-16s %-6d %-16s %-7d %-7d %-16d %-25s ", -// timestamp, -// event.UserID, -// event.ProcessName, -// event.ProcessID, -// event.ThreadID, -// event.ReturnValue, -// eventName, -// ) -// case config.ContainerModeEnabled: -// fmt.Fprintf(p.out, -// "%-16s %-13s %-6d %-16s %-7d/%-7d %-7d/%-7d %-16d %-25s ", -// timestamp, -// containerId, -// event.UserID, -// event.ProcessName, -// event.ProcessID, -// event.HostProcessID, -// event.ThreadID, -// event.HostThreadID, -// event.ReturnValue, -// eventName, -// ) -// case config.ContainerModeEnriched: -// fmt.Fprintf(p.out, -// "%-16s %-13s %-16s %-6d %-16s %-7d/%-7d %-7d/%-7d %-16d %-25s ", -// timestamp, -// containerId, -// containerImage, -// event.UserID, -// event.ProcessName, -// event.ProcessID, -// event.HostProcessID, -// event.ThreadID, -// event.HostThreadID, -// event.ReturnValue, -// eventName, -// ) -// } -// } -// for i, arg := range event.Args { -// name := arg.Name -// value := arg.Value - -// // triggeredBy from pkg/ebpf/finding.go breaks the table output, -// // so we simplify it -// if name == "triggeredBy" { -// value = fmt.Sprintf("%s", value.(map[string]interface{})["name"]) -// } - -// if i == 0 { -// fmt.Fprintf(p.out, "%s: %v", name, value) -// } else { -// fmt.Fprintf(p.out, ", %s: %v", name, value) -// } -// } -// fmt.Fprintln(p.out) -// } - -// func (p tableEventPrinter) Epilogue(stats metrics.Stats) { -// fmt.Println() -// fmt.Fprintf(p.out, "End of events stream\n") -// fmt.Fprintf(p.out, "Stats: %+v\n", stats) -// } - -// func (p tableEventPrinter) Close() { -// } - -// type templateEventPrinter struct { -// out io.WriteCloser -// templatePath string -// templateObj **template.Template -// } - -// func (p *templateEventPrinter) Init() error { -// tmplPath := p.templatePath -// if tmplPath == "" { -// return errfmt.Errorf("please specify a gotemplate for event-based output") -// } -// tmpl, err := template.ParseFiles(tmplPath) -// if err != nil { -// return errfmt.WrapError(err) -// } -// p.templateObj = &tmpl - -// return nil -// } - -// func (p templateEventPrinter) Preamble() {} - -// func (p templateEventPrinter) Print(event trace.Event) { -// if p.templateObj != nil { -// err := (*p.templateObj).Execute(p.out, event) -// if err != nil { -// logger.Errorw("Error executing template", "error", err) -// } -// } else { -// fmt.Fprintf(p.out, "Template Obj is nil") -// } -// } - -// func (p templateEventPrinter) Epilogue(stats metrics.Stats) {} - -// func (p templateEventPrinter) Close() { -// } - -// type jsonEventPrinter struct { -// out io.WriteCloser -// } - -// func (p jsonEventPrinter) Init() error { return nil } - -// func (p jsonEventPrinter) Preamble() {} - -// func (p jsonEventPrinter) Print(event trace.Event) { -// eBytes, err := json.Marshal(event) -// if err != nil { -// logger.Errorw("Error marshaling event to json", "error", err) -// } -// fmt.Fprintln(p.out, string(eBytes)) -// } - -// func (p jsonEventPrinter) Epilogue(stats metrics.Stats) {} - -// func (p jsonEventPrinter) Close() { -// } - -// // ignoreEventPrinter ignores events -// type ignoreEventPrinter struct{} - -// func (p *ignoreEventPrinter) Init() error { -// return nil -// } - -// func (p *ignoreEventPrinter) Preamble() {} - -// func (p *ignoreEventPrinter) Print(event trace.Event) {} - -// func (p *ignoreEventPrinter) Epilogue(stats metrics.Stats) {} - -// func (p ignoreEventPrinter) Close() {} - -// // forwardEventPrinter sends events over the Fluent Forward protocol to a receiver -// type forwardEventPrinter struct { -// outPath string -// url *url.URL -// client *forward.Client -// // These parameters can be set up from the URL -// tag string `default:"tracee"` -// } - -// func getParameterValue(parameters url.Values, key string, defaultValue string) string { -// param, found := parameters[key] -// // Ensure we have a non-empty parameter set for this key -// if found && param[0] != "" { -// return param[0] -// } -// // Otherwise use the default value -// return defaultValue -// } - -// func (p *forwardEventPrinter) Init() error { -// // Now parse the optional parameters with defaults and some basic verification -// u, err := url.Parse(p.outPath) -// if err != nil { -// return fmt.Errorf("unable to parse URL %q: %w", p.url, err) -// } -// p.url = u - -// parameters, _ := url.ParseQuery(p.url.RawQuery) - -// // Check if we have a tag set or default it -// p.tag = getParameterValue(parameters, "tag", "tracee") - -// // Do we want to enable requireAck? -// requireAckString := getParameterValue(parameters, "requireAck", "false") -// requireAck, err := strconv.ParseBool(requireAckString) -// if err != nil { -// return errfmt.Errorf("unable to convert requireAck value %q: %v", requireAckString, err) -// } - -// // Timeout conversion from string -// timeoutValueString := getParameterValue(parameters, "connectionTimeout", "10s") -// connectionTimeout, err := time.ParseDuration(timeoutValueString) -// if err != nil { -// return errfmt.Errorf("unable to convert connectionTimeout value %q: %v", timeoutValueString, err) -// } - -// // We should have both username and password or neither for basic auth -// username := p.url.User.Username() -// password, isPasswordSet := p.url.User.Password() -// if username != "" && !isPasswordSet { -// return errfmt.Errorf("missing basic auth configuration for Forward destination") -// } - -// // Ensure we support tcp or udp protocols -// protocol := "tcp" -// if p.url.Scheme != "" { -// protocol = p.url.Scheme -// } -// if protocol != "tcp" && protocol != "udp" { -// return errfmt.Errorf("unsupported protocol for Forward destination: %s", protocol) -// } - -// // Extract the host (and port) -// address := p.url.Host -// logger.Infow("Attempting to connect to Forward destination", "url", address, "tag", p.tag) - -// // Create a TCP connection to the forward receiver -// p.client = forward.New(forward.ConnectionOptions{ -// Factory: &forward.ConnFactory{ -// Network: protocol, -// Address: address, -// }, -// RequireAck: requireAck, -// ConnectionTimeout: connectionTimeout, -// AuthInfo: forward.AuthInfo{ -// Username: username, -// Password: password, -// }, -// }) - -// err = p.client.Connect() -// if err != nil { -// // The destination may not be available but may appear later so do not return an error here and just connect later. -// logger.Errorw("Error connecting to Forward destination", "url", p.url.String(), "error", err) -// } -// return nil -// } - -// func (p *forwardEventPrinter) Preamble() {} - -// func (p *forwardEventPrinter) Print(event trace.Event) { -// if p.client == nil { -// logger.Errorw("Invalid Forward client") -// return -// } - -// // The actual event is marshalled as JSON then sent with the other information (tag, etc.) -// eBytes, err := json.Marshal(event) -// if err != nil { -// logger.Errorw("Error marshaling event to json", "error", err) -// } - -// record := map[string]interface{}{ -// "event": string(eBytes), -// } - -// err = p.client.SendMessage(p.tag, record) -// // Assuming all is well we continue but if the connection is dropped or some other error we retry -// if err != nil { -// logger.Errorw("Error writing to Forward destination", "destination", p.url.Host, "tag", p.tag, "error", err) -// // Try five times to reconnect and send before giving up -// // TODO: consider using go-kit for circuit break, retry, etc -// for attempts := 0; attempts < 5; attempts++ { -// // Attempt to reconnect (remote end may have dropped/restarted) -// err = p.client.Reconnect() -// if err == nil { -// // Re-attempt to send -// err = p.client.SendMessage(p.tag, record) -// if err == nil { -// break -// } -// } -// } -// } -// } - -// func (p *forwardEventPrinter) Epilogue(stats metrics.Stats) {} - -// func (p forwardEventPrinter) Close() { -// if p.client != nil { -// logger.Infow("Disconnecting from Forward destination", "url", p.url.Host, "tag", p.tag) -// if err := p.client.Disconnect(); err != nil { -// logger.Errorw("Disconnecting from Forward destination", "error", err) -// } -// } -// } - -// type webhookEventPrinter struct { -// outPath string -// url *url.URL -// timeout time.Duration -// templateObj *template.Template -// contentType string -// } - -// func (ws *webhookEventPrinter) Init() error { -// u, err := url.Parse(ws.outPath) -// if err != nil { -// return errfmt.Errorf("unable to parse URL %q: %v", ws.outPath, err) -// } -// ws.url = u - -// parameters, _ := url.ParseQuery(ws.url.RawQuery) - -// timeout := getParameterValue(parameters, "timeout", "10s") -// t, err := time.ParseDuration(timeout) -// if err != nil { -// return errfmt.Errorf("unable to convert timeout value %q: %v", timeout, err) -// } -// ws.timeout = t - -// gotemplate := getParameterValue(parameters, "gotemplate", "") -// if gotemplate != "" { -// tmpl, err := template.New(filepath.Base(gotemplate)). -// Funcs(sprig.TxtFuncMap()). -// ParseFiles(gotemplate) - -// if err != nil { -// return errfmt.WrapError(err) -// } -// ws.templateObj = tmpl -// } - -// contentType := getParameterValue(parameters, "contentType", "application/json") -// ws.contentType = contentType - -// return nil -// } - -// func (ws *webhookEventPrinter) Preamble() {} - -// func (ws *webhookEventPrinter) Print(event trace.Event) { -// var ( -// payload []byte -// err error -// ) - -// if ws.templateObj != nil { -// buf := bytes.Buffer{} -// if err := ws.templateObj.Execute(&buf, event); err != nil { -// logger.Errorw("error writing to the template", "error", err) -// return -// } -// payload = buf.Bytes() -// } else { -// payload, err = json.Marshal(event) -// if err != nil { -// logger.Errorw("Error marshalling event", "error", err) -// return -// } -// } - -// client := http.Client{Timeout: ws.timeout} - -// req, err := http.NewRequest(http.MethodPost, ws.url.String(), bytes.NewReader(payload)) -// if err != nil { -// logger.Errorw("Error creating request", "error", err) -// return -// } - -// req.Header.Set("Content-Type", ws.contentType) - -// resp, err := client.Do(req) -// if err != nil { -// logger.Errorw("Error sending webhook", "error", err) -// return -// } - -// if resp.StatusCode != http.StatusOK { -// logger.Errorw(fmt.Sprintf("Error sending webhook, http status: %d", resp.StatusCode)) -// } - -// _ = resp.Body.Close() -// } - -// func (ws *webhookEventPrinter) Epilogue(stats metrics.Stats) {} - -// func (ws *webhookEventPrinter) Close() { -// } +import ( + "github.com/ShohamBit/traceectl/pkg/cmd/formatter" + pb "github.com/aquasecurity/tracee/api/v1beta1" +) + +func StreamEvents(format *formatter.Formatter, args []string, stream pb.TraceeService_StreamEventsClient) { + + //add check for the output flag + //TODO:support only table and json format for now + switch format.Format { + case formatter.FormatJSON: + jsonStreamEvents(args, stream, format) + case formatter.FormatTable: + tableStreamEvents(args, stream, format) + case formatter.FormatGoTpl: // gotemplate + fallthrough + default: + format.CMD.PrintErrln("Error: output format not supported") + return + } +} + +// tableStreamEvents prints events in a table format +func tableStreamEvents(_ []string, stream pb.TraceeService_StreamEventsClient, tbl *formatter.Formatter) { + // Init table header before streaming starts + tbl.PrintTableHeaders() + // Receive and process streamed responses + for { + res, err := stream.Recv() + if err != nil { + // Handle the error that occurs when the server closes the stream + if err.Error() == "EOF" { + break + } + tbl.CMD.PrintErrln("Error receiving streamed event: ", err) + } + tbl.PrintTableRow(res.Event) + + } +} + +// jsonStreamEvents prints events in json format +func jsonStreamEvents(_ []string, stream pb.TraceeService_StreamEventsClient, tbl *formatter.Formatter) { // Receive and process streamed responses + for { + res, err := stream.Recv() + if err != nil { + // Handle the error that occurs when the server closes the stream + if err.Error() == "EOF" { + break + } + tbl.CMD.PrintErrln("Error receiving streamed event: ", err) + } + // Print each event as a row in json format + tbl.PrintJSON(res.Event) + } +} From 9ecd5ac7d4399839e895d52f495a5e38906fb578 Mon Sep 17 00:00:00 2001 From: ShohamBit Date: Mon, 4 Nov 2024 10:19:53 +0000 Subject: [PATCH 3/6] fix error handeling --- cmd/root.go | 11 +++++++---- cmd/stream.go | 4 ++++ cmd/stream_test.go | 9 +++++---- pkg/client/client.go | 14 ++++++++++++-- pkg/client/diagnostic.go | 1 + pkg/client/service.go | 1 + 6 files changed, 30 insertions(+), 10 deletions(-) diff --git a/cmd/root.go b/cmd/root.go index 9d9349f..5b36a2d 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -140,7 +140,8 @@ func displayMetrics(cmd *cobra.Command, _ []string) { //get metrics response, err := TCD.GetMetrics(context.Background(), &pb.GetMetricsRequest{}) if err != nil { - cmd.PrintErrln("Error getting version: ", err) + cmd.PrintErrln("Error getting metrics: ", err) + return } // Display the metrics @@ -156,7 +157,6 @@ func displayMetrics(cmd *cobra.Command, _ []string) { } func displayVersion(cmd *cobra.Command, _ []string) { - //create service client if err := TCS.NewServiceClient(serverInfo); err != nil { cmd.PrintErrln("Error creating client: ", err) @@ -164,9 +164,12 @@ func displayVersion(cmd *cobra.Command, _ []string) { defer TCS.CloseConnection() //get version response, err := TCS.GetVersion(context.Background(), &pb.GetVersionRequest{}) + if err != nil { cmd.PrintErrln("Error getting version: ", err) + return + } else { + //display version + cmd.Println("Version: ", response.Version) } - //display version - cmd.Println("Version: ", response.Version) } diff --git a/cmd/stream.go b/cmd/stream.go index b17f86e..02d4955 100644 --- a/cmd/stream.go +++ b/cmd/stream.go @@ -125,6 +125,8 @@ func stream(cmd *cobra.Command, args []string) { err := TCS.NewServiceClient(serverInfo) if err != nil { cmd.PrintErrln("Error creating client: ", err) + TCS.CloseConnection() + return } defer TCS.CloseConnection() @@ -133,12 +135,14 @@ func stream(cmd *cobra.Command, args []string) { stream, err := TCS.StreamEvents(cmd.Context(), req) if err != nil { cmd.PrintErrln("Error calling Stream: ", err) + return } //create formatter for output format, err := formatter.New(formatFlag, outputFlag, cmd) if err != nil { cmd.PrintErrln("Error creating formatter: ", err) + return } //show events printer.StreamEvents(format, args, stream) diff --git a/cmd/stream_test.go b/cmd/stream_test.go index 2d3e3ab..e569f18 100644 --- a/cmd/stream_test.go +++ b/cmd/stream_test.go @@ -2,6 +2,7 @@ package cmd import ( "bytes" + "fmt" "strings" "testing" "time" @@ -15,9 +16,9 @@ import ( // currently stream can connect to the server and print the output of events to the stream var streamTests = []models.TestCase{ { - Name: "No stream subcommand", + Name: "No subcommand", Args: []string{"stream"}, - ExpectedOutput: "", //TODO: Update expected output + ExpectedOutput: mock.CreateEventsFromPolicies([]string{""}), //TODO: Update expected output }, } @@ -51,7 +52,7 @@ func TestStreamEvent(t *testing.T) { if expectedEvents, ok := test.ExpectedOutput.([]*pb.StreamEventsResponse); ok { // Split the actual output by newlines actualEvents := strings.Split(strings.TrimSpace(buf.String()), "\n") - + fmt.Println(buf.String()) // Check if the number of events match if len(actualEvents) != len(expectedEvents) { t.Errorf("Expected %d events, got %d", len(expectedEvents), len(actualEvents)) @@ -65,7 +66,7 @@ func TestStreamEvent(t *testing.T) { } } } else { - t.Errorf("Type assertion failed, expected output is not []*pb.StreamEventsResponse") + t.Errorf("Type assertion failed, expected output is not []*pb.StreamEventsResponse: %v", test.ExpectedOutput) } }) } diff --git a/pkg/client/client.go b/pkg/client/client.go index 3e32801..c11669b 100644 --- a/pkg/client/client.go +++ b/pkg/client/client.go @@ -33,15 +33,25 @@ func connectToServer(serverInfo ServerInfo) (*grpc.ClientConn, error) { // Use switch case to determine connection type var conn *grpc.ClientConn var err error - switch serverInfo.ConnectionType { case PROTOCOL_UNIX: // Dial a Unix socket - conn, err = grpc.NewClient(fmt.Sprintf("unix://%s", serverInfo.UnixSocketPath), opts...) + address := fmt.Sprintf("unix://%s", serverInfo.UnixSocketPath) + conn, err = grpc.NewClient(address, opts...) + + if err != nil { + log.Fatalf("failed to connect to server: %v", err) + return nil, err + } case PROTOCOL_TCP: // Dial a TCP address address := fmt.Sprintf(serverInfo.ADDR) conn, err = grpc.NewClient(address, opts...) + + if err != nil { + log.Fatalf("failed to connect to server: %v", err) + return nil, err + } default: return nil, fmt.Errorf("unsupported connection type: %s", serverInfo.ConnectionType) } diff --git a/pkg/client/diagnostic.go b/pkg/client/diagnostic.go index 71909fe..384b13e 100644 --- a/pkg/client/diagnostic.go +++ b/pkg/client/diagnostic.go @@ -33,6 +33,7 @@ func (tc *DiagnosticClient) NewDiagnosticClient(serverInfo ServerInfo) error { func (tc *DiagnosticClient) CloseConnection() { if err := tc.conn.Close(); err != nil { log.Printf("Failed to close connection: %v", err) + return } } diff --git a/pkg/client/service.go b/pkg/client/service.go index 68c792f..d73d819 100644 --- a/pkg/client/service.go +++ b/pkg/client/service.go @@ -32,6 +32,7 @@ func (tc *ServiceClient) NewServiceClient(serverInfo ServerInfo) error { func (tc *ServiceClient) CloseConnection() { if err := tc.conn.Close(); err != nil { log.Printf("Failed to close connection: %v", err) + return } } From b8cd2d50d40f985d22b39afe91be3daf89120dc3 Mon Sep 17 00:00:00 2001 From: ShohamBit Date: Mon, 4 Nov 2024 14:55:15 +0000 Subject: [PATCH 4/6] added test for root cmd --- cmd/root.go | 5 ++--- cmd/root_test.go | 23 +++++++++++++++++------ 2 files changed, 19 insertions(+), 9 deletions(-) diff --git a/cmd/root.go b/cmd/root.go index 5b36a2d..d0b03bd 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -2,7 +2,6 @@ package cmd import ( "context" - "fmt" "os" "github.com/ShohamBit/traceectl/pkg/client" @@ -104,7 +103,6 @@ var configCmd = &cobra.Command{ Short: "View or modify the Tracee Daemon configuration at runtime.", Long: `View or modify the Tracee Daemon configuration at runtime.`, Run: func(cmd *cobra.Command, args []string) { - fmt.Println("config called") }, } @@ -133,8 +131,9 @@ func GetRootCmd() *cobra.Command { func displayMetrics(cmd *cobra.Command, _ []string) { //create service client - if err := TCD.NewDiagnosticClient(serverInfo); err == nil { + if err := TCD.NewDiagnosticClient(serverInfo); err != nil { cmd.PrintErrln("Error creating client: ", err) + return } defer TCD.CloseConnection() //get metrics diff --git a/cmd/root_test.go b/cmd/root_test.go index cbedb2f..255460b 100644 --- a/cmd/root_test.go +++ b/cmd/root_test.go @@ -2,6 +2,7 @@ package cmd import ( "bytes" + "strings" "testing" "time" @@ -11,15 +12,25 @@ import ( var rootTests = []models.TestCase{ { - Name: "No root subcommand", - Args: []string{"root"}, - ExpectedOutput: rootCmd.Help(), // Update expected output + Name: "version", + Args: []string{"version"}, + ExpectedOutput: mock.ExpectedVersion, + }, + { + Name: "metrics", + Args: []string{"metrics"}, + ExpectedOutput: (func() string { + str := mock.ExpectedMetrics.String() // convert to string + str = strings.ReplaceAll(str, " ", "\n") // add newlines + str = strings.ReplaceAll(str, ":", ": ") // add spaces + return str + })(), }, } func TestRootCmd(t *testing.T) { // Start the mock server - mockServer, err := mock.StartMockServiceServer() + mockServer, err := mock.StartMockServer() if err != nil { t.Fatalf("Failed to start mock server: %v", err) } @@ -47,8 +58,8 @@ func TestRootCmd(t *testing.T) { // Validate output and error (if any) output := buf.String() - if output != test.ExpectedOutput { - t.Errorf("Expected output: %s, got: %s", test.ExpectedOutput, output) + if !strings.Contains(output, test.ExpectedOutput.(string)) { + t.Errorf("Expected output:\n%s\ngot:\n%s", test.ExpectedOutput, output) } }) } From 78f843695bd729d251fa531d765bae2d77edf561 Mon Sep 17 00:00:00 2001 From: ShohamBit Date: Mon, 4 Nov 2024 14:56:28 +0000 Subject: [PATCH 5/6] move and marge diagnostic and service to the same file --- cmd/event_test.go | 2 +- cmd/plugin_test.go | 2 +- cmd/policy_test.go | 2 +- cmd/stream_test.go | 7 +++---- pkg/mock/diagnostic_server.go | 26 -------------------------- pkg/mock/server.go | 27 ++++++++++++++++++++++----- 6 files changed, 28 insertions(+), 38 deletions(-) diff --git a/cmd/event_test.go b/cmd/event_test.go index e215610..a32c3ac 100644 --- a/cmd/event_test.go +++ b/cmd/event_test.go @@ -45,7 +45,7 @@ var eventTests = []models.TestCase{ func TestEvent(t *testing.T) { // Start the mock server - mockServer, err := mock.StartMockServiceServer() + mockServer, err := mock.StartMockServer() if err != nil { t.Fatalf("Failed to start mock server: %v", err) } diff --git a/cmd/plugin_test.go b/cmd/plugin_test.go index 0154074..c00a265 100644 --- a/cmd/plugin_test.go +++ b/cmd/plugin_test.go @@ -19,7 +19,7 @@ var pluginTests = []models.TestCase{ func TestPluginCmd(t *testing.T) { // Start the mock server - mockServer, err := mock.StartMockServiceServer() + mockServer, err := mock.StartMockServer() if err != nil { t.Fatalf("Failed to start mock server: %v", err) } diff --git a/cmd/policy_test.go b/cmd/policy_test.go index 5bc1e8f..8905d20 100644 --- a/cmd/policy_test.go +++ b/cmd/policy_test.go @@ -19,7 +19,7 @@ var policyTests = []models.TestCase{ func TestPolicyCmd(t *testing.T) { // Start the mock server - mockServer, err := mock.StartMockServiceServer() + mockServer, err := mock.StartMockServer() if err != nil { t.Fatalf("Failed to start mock server: %v", err) } diff --git a/cmd/stream_test.go b/cmd/stream_test.go index e569f18..11e6ea8 100644 --- a/cmd/stream_test.go +++ b/cmd/stream_test.go @@ -2,7 +2,6 @@ package cmd import ( "bytes" - "fmt" "strings" "testing" "time" @@ -18,15 +17,16 @@ var streamTests = []models.TestCase{ { Name: "No subcommand", Args: []string{"stream"}, - ExpectedOutput: mock.CreateEventsFromPolicies([]string{""}), //TODO: Update expected output + ExpectedOutput: mock.CreateEventsFromPolicies([]string{""}), }, + //TODO: add tests for subcommands } func TestStreamEvent(t *testing.T) { for _, test := range streamTests { t.Run(test.Name, func(t *testing.T) { // Start the mock server - mockServer, err := mock.StartMockServiceServer() + mockServer, err := mock.StartMockServer() if err != nil { t.Fatalf("Failed to start mock server: %v", err) } @@ -52,7 +52,6 @@ func TestStreamEvent(t *testing.T) { if expectedEvents, ok := test.ExpectedOutput.([]*pb.StreamEventsResponse); ok { // Split the actual output by newlines actualEvents := strings.Split(strings.TrimSpace(buf.String()), "\n") - fmt.Println(buf.String()) // Check if the number of events match if len(actualEvents) != len(expectedEvents) { t.Errorf("Expected %d events, got %d", len(expectedEvents), len(actualEvents)) diff --git a/pkg/mock/diagnostic_server.go b/pkg/mock/diagnostic_server.go index 7a0bec9..1d4fcff 100644 --- a/pkg/mock/diagnostic_server.go +++ b/pkg/mock/diagnostic_server.go @@ -2,10 +2,8 @@ package mock import ( "context" - "net" pb "github.com/aquasecurity/tracee/api/v1beta1" - "google.golang.org/grpc" ) var ( @@ -14,30 +12,6 @@ var ( LostWrCount: 7, LostNtCapCount: 8, LostBPFLogsCount: 9} ) -// MockDiagnosticServer implements the gRPC server interface for testing -type MockDiagnosticServer struct { - pb.UnimplementedDiagnosticServiceServer // Embed the unimplemented server -} - -// Start a mock gRPC server -func StartMockDiagnosticServer() (*grpc.Server, error) { - lis, err := net.Listen("tcp", serverInfo.ADDR) - if err != nil { - return nil, err - } - - s := grpc.NewServer() - pb.RegisterDiagnosticServiceServer(s, &MockDiagnosticServer{}) - - go func() { - if err := s.Serve(lis); err != nil { - // Handle the error (e.g., log it) - } - }() - - return s, nil -} - func (s *MockDiagnosticServer) GetMetrics(ctx context.Context, req *pb.GetMetricsRequest) (*pb.GetMetricsResponse, error) { return &ExpectedMetrics, nil } diff --git a/pkg/mock/server.go b/pkg/mock/server.go index 49d16bd..2014d0c 100644 --- a/pkg/mock/server.go +++ b/pkg/mock/server.go @@ -24,12 +24,16 @@ type MockServiceServer struct { pb.UnimplementedTraceeServiceServer // Embed the unimplemented server } +// MockDiagnosticServer implements the gRPC server interface for testing +type MockDiagnosticServer struct { + pb.UnimplementedDiagnosticServiceServer // Embed the unimplemented server +} + // CreateMockServer initializes the gRPC server and binds it to a Unix socket listener func CreateMockServer() (*grpc.Server, net.Listener, error) { - //check for unix socket + // Check for existing Unix socket and remove it if necessary if _, err := os.Stat(serverInfo.UnixSocketPath); err == nil { - err := os.Remove(serverInfo.UnixSocketPath) - if err != nil { + if err := os.Remove(serverInfo.UnixSocketPath); err != nil { return nil, nil, fmt.Errorf("failed to cleanup gRPC listening address (%s): %v", serverInfo.UnixSocketPath, err) } } @@ -43,15 +47,19 @@ func CreateMockServer() (*grpc.Server, net.Listener, error) { // Create a new gRPC server server := grpc.NewServer() + // Register both TraceeService and DiagnosticService with the server + pb.RegisterTraceeServiceServer(server, &MockServiceServer{}) + pb.RegisterDiagnosticServiceServer(server, &MockDiagnosticServer{}) + return server, listener, nil } -func StartMockServiceServer() (*grpc.Server, error) { +// StartMockServer starts the gRPC server with both services registered +func StartMockServer() (*grpc.Server, error) { mockServer, listener, err := CreateMockServer() if err != nil { return nil, fmt.Errorf("failed to create mock server: %v", err) } - pb.RegisterTraceeServiceServer(mockServer, &MockServiceServer{}) // Start serving in a goroutine go func() { @@ -59,5 +67,14 @@ func StartMockServiceServer() (*grpc.Server, error) { fmt.Printf("gRPC server failed: %v\n", err) } }() + return mockServer, nil } + +// StopMockServer stops the server and removes the Unix socket +func StopMockServer(server *grpc.Server) { + server.GracefulStop() + if err := os.Remove(serverInfo.UnixSocketPath); err != nil { + fmt.Printf("failed to remove Unix socket: %v\n", err) + } +} From d06b99aff7b6bd87a4251446b9e4836961fce8d7 Mon Sep 17 00:00:00 2001 From: ShohamBit Date: Tue, 5 Nov 2024 14:56:45 +0000 Subject: [PATCH 6/6] add support for traceectl event describe --- cmd/event.go | 42 +++++++++++++++++++--- cmd/root.go | 20 +++++------ cmd/stream.go | 10 +++--- go.mod | 4 +++ go.sum | 8 +++++ pkg/client/service.go | 3 ++ pkg/cmd/formatter/TableFormatter.go | 54 +++++++++++++++++++++++++++-- pkg/cmd/printer/printer.go | 14 ++++++-- traceeCTLdesign.md | 26 +++++++------- 9 files changed, 144 insertions(+), 37 deletions(-) diff --git a/cmd/event.go b/cmd/event.go index 809721a..da7c31b 100644 --- a/cmd/event.go +++ b/cmd/event.go @@ -3,6 +3,9 @@ package cmd import ( "context" + "github.com/ShohamBit/traceectl/pkg/client" + "github.com/ShohamBit/traceectl/pkg/cmd/formatter" + "github.com/ShohamBit/traceectl/pkg/cmd/printer" pb "github.com/aquasecurity/tracee/api/v1beta1" "github.com/spf13/cobra" ) @@ -57,7 +60,7 @@ var describeEventCmd = &cobra.Command{ Long: `Retrieves the detailed definition of a specific event, including its fields, types, and other metadata.`, Args: cobra.MinimumNArgs(1), Run: func(cmd *cobra.Command, args []string) { - + getEventDescriptions(cmd, args) }, } @@ -93,14 +96,16 @@ var runEventCmd = &cobra.Command{ func enableEvents(cmd *cobra.Command, eventNames []string) { // Create Tracee gRPC client - if err := TCS.NewServiceClient(serverInfo); err != nil { + var traceeClient client.ServiceClient // tracee client + + if err := traceeClient.NewServiceClient(serverInfo); err != nil { cmd.PrintErrln("Error creating client: ", err) return // Exit on error } // Iterate over event names and enable each one for _, eventName := range eventNames { - _, err := TCS.EnableEvent(context.Background(), &pb.EnableEventRequest{Name: eventName}) + _, err := traceeClient.EnableEvent(context.Background(), &pb.EnableEventRequest{Name: eventName}) if err != nil { cmd.PrintErrln("Error enabling event:", err) continue // Continue on error with the next event @@ -111,14 +116,15 @@ func enableEvents(cmd *cobra.Command, eventNames []string) { func disableEvents(cmd *cobra.Command, eventNames []string) { // Create Tracee gRPC client - if err := TCS.NewServiceClient(serverInfo); err != nil { + var traceeClient client.ServiceClient + if err := traceeClient.NewServiceClient(serverInfo); 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 := TCS.DisableEvent(context.Background(), &pb.DisableEventRequest{Name: eventName}) + _, err := traceeClient.DisableEvent(context.Background(), &pb.DisableEventRequest{Name: eventName}) if err != nil { cmd.PrintErrln("Error disabling event:", err) continue // Continue on error with the next event @@ -126,3 +132,29 @@ func disableEvents(cmd *cobra.Command, eventNames []string) { cmd.Println("Disabled event:", eventName) } } + +func getEventDescriptions(cmd *cobra.Command, args []string) { + //create service client + var traceeClient client.ServiceClient + if err := traceeClient.NewServiceClient(serverInfo); err != nil { + cmd.PrintErrln("Error creating client: ", err) + } + defer traceeClient.CloseConnection() + response, err := traceeClient.GetEventDefinitions(context.Background(), &pb.GetEventDefinitionsRequest{EventNames: args}) + + if err != nil { + cmd.PrintErrln("Error getting event definitions: ", err) + return + + } + //display event definitions + //don't support different outputs and formats + format, err := formatter.New("table", "", cmd) + if err != nil { + cmd.PrintErrln("Error creating formatter: ", err) + return + } + //show events + printer.DescribeEvent(format, args, response) + +} diff --git a/cmd/root.go b/cmd/root.go index d0b03bd..5d5ff5d 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -11,9 +11,7 @@ import ( ) var ( - TCS client.ServiceClient // tracee service client - TCD client.DiagnosticClient // tracee diagnostic client - serverInfo client.ServerInfo = client.ServerInfo{ + serverInfo client.ServerInfo = client.ServerInfo{ ConnectionType: client.PROTOCOL_UNIX, UnixSocketPath: client.SOCKET, ADDR: client.DefaultIP + ":" + client.DefaultPort, @@ -64,7 +62,6 @@ var connectCmd = &cobra.Command{ Short: "Connect to the server", Long: "Connects to a stream and displays events in real time.", Run: func(cmd *cobra.Command, args []string) { - }, } var metricsCmd = &cobra.Command{ @@ -129,15 +126,15 @@ func GetRootCmd() *cobra.Command { // displayMetrics fetches and prints Tracee metrics func displayMetrics(cmd *cobra.Command, _ []string) { - + var traceeClient client.DiagnosticClient //create service client - if err := TCD.NewDiagnosticClient(serverInfo); err != nil { + if err := traceeClient.NewDiagnosticClient(serverInfo); err != nil { cmd.PrintErrln("Error creating client: ", err) return } - defer TCD.CloseConnection() + defer traceeClient.CloseConnection() //get metrics - response, err := TCD.GetMetrics(context.Background(), &pb.GetMetricsRequest{}) + response, err := traceeClient.GetMetrics(context.Background(), &pb.GetMetricsRequest{}) if err != nil { cmd.PrintErrln("Error getting metrics: ", err) return @@ -157,12 +154,13 @@ func displayMetrics(cmd *cobra.Command, _ []string) { func displayVersion(cmd *cobra.Command, _ []string) { //create service client - if err := TCS.NewServiceClient(serverInfo); err != nil { + var traceeClient client.ServiceClient + if err := traceeClient.NewServiceClient(serverInfo); err != nil { cmd.PrintErrln("Error creating client: ", err) } - defer TCS.CloseConnection() + defer traceeClient.CloseConnection() //get version - response, err := TCS.GetVersion(context.Background(), &pb.GetVersionRequest{}) + response, err := traceeClient.GetVersion(context.Background(), &pb.GetVersionRequest{}) if err != nil { cmd.PrintErrln("Error getting version: ", err) diff --git a/cmd/stream.go b/cmd/stream.go index 02d4955..5514daa 100644 --- a/cmd/stream.go +++ b/cmd/stream.go @@ -1,6 +1,7 @@ package cmd import ( + "github.com/ShohamBit/traceectl/pkg/client" "github.com/ShohamBit/traceectl/pkg/cmd/formatter" "github.com/ShohamBit/traceectl/pkg/cmd/printer" pb "github.com/aquasecurity/tracee/api/v1beta1" @@ -122,17 +123,18 @@ var resumeStreamCmd = &cobra.Command{ // stream events directly from tracee func stream(cmd *cobra.Command, args []string) { // Create service client - err := TCS.NewServiceClient(serverInfo) + var traceeClient client.ServiceClient + err := traceeClient.NewServiceClient(serverInfo) if err != nil { cmd.PrintErrln("Error creating client: ", err) - TCS.CloseConnection() + traceeClient.CloseConnection() return } - defer TCS.CloseConnection() + defer traceeClient.CloseConnection() // create stream from client req := &pb.StreamEventsRequest{Policies: args} - stream, err := TCS.StreamEvents(cmd.Context(), req) + stream, err := traceeClient.StreamEvents(cmd.Context(), req) if err != nil { cmd.PrintErrln("Error calling Stream: ", err) return diff --git a/go.mod b/go.mod index bc8fd89..3c26268 100644 --- a/go.mod +++ b/go.mod @@ -9,10 +9,14 @@ require ( ) require ( + github.com/aquasecurity/table v1.8.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect + github.com/mattn/go-runewidth v0.0.13 // indirect + github.com/rivo/uniseg v0.2.0 // indirect github.com/spf13/pflag v1.0.5 // indirect golang.org/x/net v0.28.0 // indirect golang.org/x/sys v0.24.0 // indirect + golang.org/x/term v0.23.0 // indirect golang.org/x/text v0.17.0 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142 // indirect google.golang.org/protobuf v1.34.2 // indirect diff --git a/go.sum b/go.sum index 3707558..10b9fed 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,5 @@ +github.com/aquasecurity/table v1.8.0 h1:9ntpSwrUfjrM6/YviArlx/ZBGd6ix8W+MtojQcM7tv0= +github.com/aquasecurity/table v1.8.0/go.mod h1:eqOmvjjB7AhXFgFqpJUEE/ietg7RrMSJZXyTN8E/wZw= github.com/aquasecurity/tracee/api v0.0.0-20240918153521-1b3f9e8657e0 h1:3pt2Ceg6urRTUdOaX++LheYObFE4DbS63chktnQu24I= github.com/aquasecurity/tracee/api v0.0.0-20240918153521-1b3f9e8657e0/go.mod h1:Gn6xVkaBkVe1pOQ0++uuHl+lMMClv0TPY8mCQ6j88aA= github.com/cpuguy83/go-md2man/v2 v2.0.4/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o= @@ -5,6 +7,10 @@ github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= +github.com/mattn/go-runewidth v0.0.13 h1:lTGmDsbAYt5DmK6OnoV7EuIF1wEIFAcxld6ypU4OSgU= +github.com/mattn/go-runewidth v0.0.13/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= +github.com/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY= +github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= github.com/spf13/cobra v1.8.1 h1:e5/vxKd/rZsfSJMUX1agtjeTDf+qv1/JdBF8gg5k9ZM= github.com/spf13/cobra v1.8.1/go.mod h1:wHxEcudfqmLYa8iTfL+OuZPbBZkmvliBWKIezN3kD9Y= @@ -14,6 +20,8 @@ golang.org/x/net v0.28.0 h1:a9JDOJc5GMUJ0+UDqmLT86WiEy7iWyIhz8gz8E4e5hE= golang.org/x/net v0.28.0/go.mod h1:yqtgsTWOOnlGLG9GFRrK3++bGOUEkNBoHZc8MEDWPNg= golang.org/x/sys v0.24.0 h1:Twjiwq9dn6R1fQcyiK+wQyHWfaz/BJB+YIpzU/Cv3Xg= golang.org/x/sys v0.24.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/term v0.23.0 h1:F6D4vR+EHoL9/sWAWgAR1H2DcHr4PareCbAaCo1RpuU= +golang.org/x/term v0.23.0/go.mod h1:DgV24QBUrK6jhZXl+20l6UWznPlwAHm1Q1mGHtydmSk= golang.org/x/text v0.17.0 h1:XtiM5bkSOt+ewxlOE/aE/AKEHibwj/6gvWMl9Rsh0Qc= golang.org/x/text v0.17.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY= google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142 h1:e7S5W7MGGLaSu8j3YjdezkZ+m1/Nm0uRVRMEMGk26Xs= diff --git a/pkg/client/service.go b/pkg/client/service.go index d73d819..316d27a 100644 --- a/pkg/client/service.go +++ b/pkg/client/service.go @@ -56,3 +56,6 @@ func (tc *ServiceClient) DisableEvent(ctx context.Context, req *pb.DisableEventR func (tc *ServiceClient) StreamEvents(ctx context.Context, req *pb.StreamEventsRequest) (pb.TraceeService_StreamEventsClient, error) { return tc.client.StreamEvents(ctx, req) } +func (tc *ServiceClient) GetEventDefinitions(ctx context.Context, req *pb.GetEventDefinitionsRequest) (*pb.GetEventDefinitionsResponse, error) { + return tc.client.GetEventDefinitions(ctx, req) +} diff --git a/pkg/cmd/formatter/TableFormatter.go b/pkg/cmd/formatter/TableFormatter.go index 23d3507..70857d7 100644 --- a/pkg/cmd/formatter/TableFormatter.go +++ b/pkg/cmd/formatter/TableFormatter.go @@ -2,12 +2,14 @@ package formatter import ( "fmt" + "os" "strings" + "github.com/aquasecurity/table" pb "github.com/aquasecurity/tracee/api/v1beta1" ) -func (f *Formatter) PrintTableHeaders() { +func (f *Formatter) PrintSteamTableHeaders() { f.CMD.Printf("%-15s %-25s %-15s %-15s %s\n", "TIME", "EVENT NAME", @@ -16,7 +18,7 @@ func (f *Formatter) PrintTableHeaders() { "DATA", ) } -func (f *Formatter) PrintTableRow(event *pb.Event) { +func (f *Formatter) PrintStreamTableRow(event *pb.Event) { timestamp := event.Timestamp.AsTime().Format("15:04:05.000") f.CMD.Printf("%-15s %-25s %-15s %-15s %s\n", @@ -70,3 +72,51 @@ func getEventValue(ev *pb.EventValue) string { return "unknown" } } + +func (f *Formatter) PrintEventDescription(response *pb.GetEventDefinitionsResponse) *table.Table { + tbl := table.New(os.Stdout) + tbl.SetHeaders("ID", "Name", "Version", "Description", "Tags", "Threat") + tbl.AddHeaders("ID", "Name", "Version", "Description", "Tags", "description", "mitre", "severity", "name", "properties") + tbl.SetHeaderColSpans(0, 1, 1, 1, 1, 1, 5) + tbl.SetAutoMergeHeaders(true) + for _, event := range response.Definitions { + // Check if the optional field Threat is set (non-nil) + + if event.Threat != nil { + tbl.AddRow( + fmt.Sprintf("%d", event.Id), + event.Name, + fmt.Sprintf("%d.%d.%d", event.Version.Major, event.Version.Minor, event.Version.Patch), + event.Description, + strings.Join(event.Tags, ", "), + event.Threat.Description, + event.Threat.Mitre.String(), + event.Threat.Severity.String(), + event.Threat.Name, + mapToString(event.Threat.Properties), + ) + } else { + + tbl.AddRow( + fmt.Sprintf("%d", event.Id), + event.Name, + fmt.Sprintf("%d.%d.%d", event.Version.Major, event.Version.Minor, event.Version.Patch), + event.Description, + strings.Join(event.Tags, ", "), + ) + } + } + return tbl +} + +func mapToString(m map[string]string) string { + var builder strings.Builder + for key, value := range m { + builder.WriteString(fmt.Sprintf("%s: %s, ", key, value)) + } + result := builder.String() + if len(result) > 0 { + result = result[:len(result)-2] // Remove the trailing ", " + } + return result +} diff --git a/pkg/cmd/printer/printer.go b/pkg/cmd/printer/printer.go index 13c2bac..74b17e1 100644 --- a/pkg/cmd/printer/printer.go +++ b/pkg/cmd/printer/printer.go @@ -25,7 +25,7 @@ func StreamEvents(format *formatter.Formatter, args []string, stream pb.TraceeSe // tableStreamEvents prints events in a table format func tableStreamEvents(_ []string, stream pb.TraceeService_StreamEventsClient, tbl *formatter.Formatter) { // Init table header before streaming starts - tbl.PrintTableHeaders() + tbl.PrintSteamTableHeaders() // Receive and process streamed responses for { res, err := stream.Recv() @@ -36,7 +36,7 @@ func tableStreamEvents(_ []string, stream pb.TraceeService_StreamEventsClient, t } tbl.CMD.PrintErrln("Error receiving streamed event: ", err) } - tbl.PrintTableRow(res.Event) + tbl.PrintStreamTableRow(res.Event) } } @@ -56,3 +56,13 @@ func jsonStreamEvents(_ []string, stream pb.TraceeService_StreamEventsClient, tb tbl.PrintJSON(res.Event) } } + +func DescribeEvent(format *formatter.Formatter, args []string, response *pb.GetEventDefinitionsResponse) { + //this can add support for other output formats + tableDescribeEvent(format, args, response) + +} +func tableDescribeEvent(format *formatter.Formatter, _ []string, response *pb.GetEventDefinitionsResponse) { + tbl := format.PrintEventDescription(response) + tbl.Render() +} diff --git a/traceeCTLdesign.md b/traceeCTLdesign.md index 9a81291..9e6d2ac 100644 --- a/traceeCTLdesign.md +++ b/traceeCTLdesign.md @@ -1,3 +1,4 @@ +# TraceeCTL Policy Management: traceectl policy create @@ -13,13 +14,13 @@ Event Management: traceectl event describe traceectl event enable traceectl event disable - traceectl event run [--args ] + traceectl event run --args {\} Stream Management: - traceectl stream create --name [--destination ] [--format ] [--fields ] [--parse-data] [--filter ] + traceectl stream create --name [--destination \] [--format \] [--fields \] [--parse-data] [--filter \] traceectl stream describe traceectl stream list - traceectl stream update [--destination ] [--format ] [--fields ] [--parse-data] [--filter ] + traceectl stream update [--destination \] [--format \] [--fields \] [--parse-data] [--filter \] traceectl stream delete traceectl stream connect traceectl stream set-default @@ -33,11 +34,11 @@ Plugin Management: Additional Commands (Potential): traceectl connect [] - traceectl metrics [--output ] + traceectl metrics [--output \] traceectl diagnose [--component ] - traceectl logs [--filter ] + traceectl logs [--filter \] traceectl status - traceectl config [set|get|update] [