Use positionnal argument for directory/file input/output path (#232)

* init * fix unit tests * fix linter * Update pkg/config/collector.go Co-authored-by: Edouard Schweisguth <edouard.schweisguth@datadoghq.com> * PR comment --------- Co-authored-by: Edouard Schweisguth <edouard.schweisguth@datadoghq.com>
DataDog · Jul 30, 2024 · 3db007e · 3db007e
1 parent e31957f
commit 3db007e
Show file tree

Hide file tree

Showing 8 changed files with 35 additions and 28 deletions.
diff --git a/cmd/kubehound/dumper.go b/cmd/kubehound/dumper.go
@@ -34,7 +34,7 @@ var (
 		},
 		RunE: func(cobraCmd *cobra.Command, args []string) error {
 			// using compress feature
-			viper.Set(config.CollectorFileArchiveFormat, true)
+			viper.Set(config.CollectorFileArchiveNoCompress, false)
 
 			// Create a temporary directory
 			tmpDir, err := os.MkdirTemp("", "kubehound")
@@ -64,10 +64,13 @@ var (
 		},
 	}
 	dumpLocalCmd = &cobra.Command{
-		Use:   "local",
+		Use:   "local [directory to dump the data]",
 		Short: "Dump locally the k8s resources of a targeted cluster",
+		Args:  cobra.ExactArgs(1),
 		Long:  `Collect all Kubernetes resources needed to build the attack path in an offline format locally (compressed or flat)`,
 		PreRunE: func(cobraCmd *cobra.Command, args []string) error {
+			viper.Set(config.CollectorFileDirectory, args[0])
+
 			return cmd.InitializeKubehoundConfig(cobraCmd.Context(), "", true, true)
 		},
 		RunE: func(cobraCmd *cobra.Command, args []string) error {

diff --git a/cmd/kubehound/ingest.go b/cmd/kubehound/ingest.go
@@ -10,10 +10,6 @@ import (
 	"github.com/spf13/viper"
 )
 
-var (
-	inputFilePath string
-)
-
 var (
 	ingestCmd = &cobra.Command{
 		Use:   "ingest",
@@ -22,9 +18,10 @@ var (
 	}
 
 	localIngestCmd = &cobra.Command{
-		Use:   "local",
+		Use:   "local [directory or tar.gz path]",
 		Short: "Ingest data locally from a KubeHound dump",
 		Long:  `Run an ingestion locally using a previous dump (directory or tar.gz)`,
+		Args:  cobra.ExactArgs(1),
 		PreRunE: func(cobraCmd *cobra.Command, args []string) error {
 			return cmd.InitializeKubehoundConfig(cobraCmd.Context(), "", true, true)
 		},
@@ -35,7 +32,7 @@ var (
 				return fmt.Errorf("get config: %w", err)
 			}
 
-			return core.CoreLocalIngest(cobraCmd.Context(), khCfg, inputFilePath)
+			return core.CoreLocalIngest(cobraCmd.Context(), khCfg, args[0])
 		},
 	}
 
@@ -82,7 +79,6 @@ func init() {
 
 	ingestCmd.AddCommand(localIngestCmd)
 	cmd.InitLocalIngestCmd(localIngestCmd)
-	localIngestCmd.Flags().StringVar(&inputFilePath, "data", "", "Filepath for the data to process (directory or tar.gz path)")
 
 	ingestCmd.AddCommand(remoteIngestCmd)
 	cmd.InitRemoteIngestCmd(remoteIngestCmd, true)

diff --git a/pkg/cmd/dump.go b/pkg/cmd/dump.go
@@ -44,12 +44,8 @@ func InitDumpCmd(cmd *cobra.Command) {
 }
 
 func InitLocalDumpCmd(cmd *cobra.Command) {
-	cmd.Flags().Bool("compress", false, "Enable compression for the dumped data (generates a tar.gz file)")
-	viper.BindPFlag(config.CollectorFileArchiveFormat, cmd.Flags().Lookup("compress")) //nolint: errcheck
-
-	cmd.Flags().String("output-dir", "", "Directory to dump the data")
-	viper.BindPFlag(config.CollectorFileDirectory, cmd.Flags().Lookup("output-dir")) //nolint: errcheck
-	cmd.MarkFlagRequired("output-dir")                                               //nolint: errcheck
+	cmd.Flags().Bool("no-compress", false, "Disable compression for the dumped data (generates a directory)")
+	viper.BindPFlag(config.CollectorFileArchiveNoCompress, cmd.Flags().Lookup("no-compress")) //nolint: errcheck
 }
 
 func InitRemoteDumpCmd(cmd *cobra.Command) {

diff --git a/pkg/config/collector.go b/pkg/config/collector.go
@@ -10,16 +10,17 @@ const (
 	DefaultK8sAPIPageBufferSize     int32 = 10
 	DefaultK8sAPIRateLimitPerSecond int   = 100
 	DefaultK8sAPINonInteractive     bool  = false
-
-	CollectorLiveRate           = "collector.live.rate_limit_per_second"
-	CollectorLivePageSize       = "collector.live.page_size"
-	CollectorLivePageBufferSize = "collector.live.page_buffer_size"
-	CollectorNonInteractive     = "collector.non_interactive"
-	CollectorFileArchiveFormat  = "collector.file.archive.format"
-	CollectorFileDirectory      = "collector.file.directory"
-	CollectorFileClusterName    = "collector.file.cluster_name"
-	CollectorFileBlobRegion     = "collector.file.blob.region"
-	CollectorFileBlobBucket     = "collector.file.blob.bucket"
+	DefaultArchiveNoCompress        bool  = false
+
+	CollectorLiveRate              = "collector.live.rate_limit_per_second"
+	CollectorLivePageSize          = "collector.live.page_size"
+	CollectorLivePageBufferSize    = "collector.live.page_buffer_size"
+	CollectorNonInteractive        = "collector.non_interactive"
+	CollectorFileArchiveNoCompress = "collector.file.archive.no_compress"
+	CollectorFileDirectory         = "collector.file.directory"
+	CollectorFileClusterName       = "collector.file.cluster_name"
+	CollectorFileBlobRegion        = "collector.file.blob.region"
+	CollectorFileBlobBucket        = "collector.file.blob.bucket"
 )
 
 // CollectorConfig configures collector specific parameters.
@@ -47,7 +48,7 @@ type FileCollectorConfig struct {
 
 type FileArchiveConfig struct {
 	ArchiveName string `mapstructure:"archive_name"` // Name of the output archive
-	Format      bool   `mapstructure:"format"`       // Enable compression for the dumped data (generates a tar.gz file)
+	NoCompress  bool   `mapstructure:"no_compress"`  // Disable compression for the dumped data (generates a tar.gz file)
 }
 
 type BlobConfig struct {

diff --git a/pkg/config/config.go b/pkg/config/config.go
@@ -95,6 +95,9 @@ func SetDefaultValues(c *viper.Viper) {
 	c.SetDefault(CollectorLiveRate, DefaultK8sAPIRateLimitPerSecond)
 	c.SetDefault(CollectorNonInteractive, DefaultK8sAPINonInteractive)
 
+	// File collector module
+	c.SetDefault(CollectorFileArchiveNoCompress, DefaultArchiveNoCompress)
+
 	// Default values for storage provider
 	c.SetDefault("storage.wipe", true)
 	c.SetDefault("storage.retry", DefaultRetry)

diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go
@@ -34,6 +34,9 @@ func TestMustLoadConfig(t *testing.T) {
 					File: &FileCollectorConfig{
 						Directory:   "cluster-data/",
 						ClusterName: "test-cluster",
+						Archive: &FileArchiveConfig{
+							NoCompress: DefaultArchiveNoCompress,
+						},
 					},
 					// This is always set as the default value
 					Live: &K8SAPICollectorConfig{
@@ -102,6 +105,11 @@ func TestMustLoadConfig(t *testing.T) {
 				},
 				Collector: CollectorConfig{
 					Type: CollectorTypeK8sAPI,
+					File: &FileCollectorConfig{
+						Archive: &FileArchiveConfig{
+							NoCompress: DefaultArchiveNoCompress,
+						},
+					},
 					Live: &K8SAPICollectorConfig{
 						PageSize:           500,
 						PageBufferSize:     10,

diff --git a/pkg/kubehound/core/core_dump.go b/pkg/kubehound/core/core_dump.go
@@ -95,7 +95,7 @@ func runLocalDump(ctx context.Context, khCfg *config.KubehoundConfig) (string, e
 
 	// Create the dumper instance
 	collectorLocalOutputDir := khCfg.Collector.File.Directory
-	collectorLocalCompress := khCfg.Collector.File.Archive.Format
+	collectorLocalCompress := !khCfg.Collector.File.Archive.NoCompress
 	log.I.Infof("Dumping %q to %q", khCfg.Dynamic.ClusterName, collectorLocalOutputDir)
 	dumpIngestor, err := dump.NewDumpIngestor(ctx, collect, collectorLocalCompress, collectorLocalOutputDir, khCfg.Dynamic.RunID)
 	if err != nil {

diff --git a/test/system/setup_test.go b/test/system/setup_test.go
@@ -90,7 +90,7 @@ func Dump(ctx context.Context, compress bool) (*config.KubehoundConfig, string)
 	}
 	cmd.InitDumpCmd(dumpCmd)
 
-	viper.Set(config.CollectorFileArchiveFormat, compress)
+	viper.Set(config.CollectorFileArchiveNoCompress, !compress)
 
 	tmpDir, err := os.MkdirTemp("/tmp/", "kh-system-tests-*")
 	if err != nil {