add metric flags to clis

cmd/tracee-ebpf/main.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"io/ioutil"
 	"log"
+	"net/http"
 	"os"
 	"os/signal"
 	"path"
@@ -19,13 +20,17 @@ import (
 	"github.com/aquasecurity/tracee/cmd/tracee-ebpf/internal/flags"
 	"github.com/aquasecurity/tracee/pkg/capabilities"
 	"github.com/aquasecurity/tracee/pkg/external"
+	"github.com/aquasecurity/tracee/tracee-ebpf/metrics"
 	"github.com/aquasecurity/tracee/tracee-ebpf/tracee"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/syndtr/gocapability/capability"
 	cli "github.com/urfave/cli/v2"
 )
 
 var debug bool
 var traceeInstallPath string
+var listenMetrics bool
+var metricsAddr string
 
 var version string
 
@@ -159,6 +164,24 @@ func main() {
 				return fmt.Errorf("error creating Tracee: %v", err)
 			}
 
+			if listenMetrics {
+				err := metrics.RegisterPrometheus(t.Stats())
+				if err != nil {
+					fmt.Fprintf(os.Stderr, "Error registering prometheus metrics: %v\n", err)
+				} else {
+					mux := http.NewServeMux()
+					mux.Handle("/metrics", promhttp.Handler())
+
+					go func() {
+						fmt.Fprintf(os.Stdout, "Serving metrics endpoint at %s\n", metricsAddr)
+						if err := http.ListenAndServe(metricsAddr, mux); err != http.ErrServerClosed {
+							fmt.Fprintf(os.Stderr, "Error serving metrics endpoint: %v\n", err)
+						}
+					}()
+				}
+
+			}
+
 			if err := os.MkdirAll(cfg.Capture.OutputPath, 0755); err != nil {
 				t.Close()
 				return fmt.Errorf("error creating output path: %v", err)
@@ -220,8 +243,8 @@ func main() {
 
 			// always print stats before exiting
 			defer func() {
-				stats := t.GetStats()
-				printer.Epilogue(stats)
+				stats := t.Stats()
+				printer.Epilogue(*stats)
 				printer.Close()
 			}()
 
@@ -276,6 +299,18 @@ func main() {
 				Usage:       "path where tracee will install or lookup it's resources",
 				Destination: &traceeInstallPath,
 			},
+			&cli.BoolFlag{
+				Name:        "metrics",
+				Usage:       "enable metrics endpoint",
+				Destination: &listenMetrics,
+				Value:       false,
+			},
+			&cli.StringFlag{
+				Name:        "metrics-addr",
+				Usage:       "listening address of the metrics endpoint server",
+				Value:       ":3366",
+				Destination: &metricsAddr,
+			},
 		},
 	}
 

cmd/tracee-rules/main.go

@@ -14,13 +14,17 @@ import (
 	"syscall"
 
 	"github.com/aquasecurity/tracee/tracee-rules/engine"
+	"github.com/aquasecurity/tracee/tracee-rules/metrics"
 	"github.com/aquasecurity/tracee/types"
 	"github.com/open-policy-agent/opa/compile"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/urfave/cli/v2"
 )
 
 const (
 	signatureBufferFlag = "sig-buffer"
+	metricsFlag         = "metrics"
+	metricsAddrFlag     = "metrics-addr"
 )
 
 func main() {
@@ -117,6 +121,26 @@ func main() {
 			if err != nil {
 				return fmt.Errorf("constructing engine: %w", err)
 			}
+
+			if c.Bool(metricsFlag) {
+				err := metrics.RegisterPrometheus(e.Stats())
+				if err != nil {
+					fmt.Fprintf(os.Stderr, "Error registering prometheus metrics: %v\n", err)
+				} else {
+					mux := http.NewServeMux()
+					mux.Handle("/metrics", promhttp.Handler())
+
+					go func() {
+						metricsAddr := c.String(metricsAddrFlag)
+						fmt.Fprintf(os.Stdout, "Serving metrics endpoint at %s\n", metricsAddr)
+						if err := http.ListenAndServe(metricsAddr, mux); err != http.ErrServerClosed {
+							fmt.Fprintf(os.Stderr, "Error serving metrics endpoint: %v\n", err)
+						}
+					}()
+				}
+
+			}
+
 			e.Start(sigHandler())
 			return nil
 		},
@@ -188,6 +212,16 @@ func main() {
 				Usage: "size of the event channel's buffer consumed by signatures",
 				Value: 1000,
 			},
+			&cli.BoolFlag{
+				Name:  metricsFlag,
+				Usage: "enable metrics endpoint",
+				Value: false,
+			},
+			&cli.StringFlag{
+				Name:  metricsAddrFlag,
+				Usage: "listening address of the metrics endpoint server",
+				Value: ":4466",
+			},
 		},
 	}
 	err := app.Run(os.Args)