[Dumper] Rework of the object ingestion

pkg/dump/ingestor.go

@@ -1,27 +1,19 @@
 package dump
 
 import (
 	"context"
-	"encoding/json"
 	"fmt"
 	"path"
 
 	"time"
 
 	"github.com/DataDog/KubeHound/pkg/collector"
 	"github.com/DataDog/KubeHound/pkg/dump/writer"
-	"github.com/DataDog/KubeHound/pkg/globals/types"
-	"github.com/DataDog/KubeHound/pkg/kubehound/ingestor/preflight"
-	"github.com/DataDog/KubeHound/pkg/telemetry/log"
-	"github.com/DataDog/KubeHound/pkg/telemetry/metric"
 	"github.com/DataDog/KubeHound/pkg/telemetry/span"
-	"github.com/DataDog/KubeHound/pkg/telemetry/statsd"
 	"github.com/DataDog/KubeHound/pkg/telemetry/tag"
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
 )
 
-var _ collector.GenericIngestor = (*DumpIngestor)(nil)
-
 type DumpIngestor struct {
 	directoryOutput string
 	ResultName      string
@@ -81,158 +73,23 @@
 	return d.writer.OutputPath()
 }
 
-func (d *DumpIngestor) IngestNode(ctx context.Context, node types.NodeType) error {
-	if ok, err := preflight.CheckNode(node); !ok {
-
-		return err
-	}
-
-	return d.processObject(ctx, node, collector.NodePath)
-}
-
-func ingestPodPath(pod types.PodType) string {
-	return path.Join(pod.Namespace, collector.PodPath)
-}
-
-func (d *DumpIngestor) IngestPod(ctx context.Context, pod types.PodType) error {
-	if ok, err := preflight.CheckPod(pod); !ok {
-		entity := tag.Entity(tag.EntityPods)
-		_ = statsd.Incr(metric.CollectorSkip, append(d.collector.Tags(ctx), entity), 1)
-
-		return err
-	}
-
-	filePath := ingestPodPath(pod)
-
-	return d.processObject(ctx, pod, filePath)
-}
-
-func ingestRolePath(roleBinding types.RoleType) string {
-	return path.Join(roleBinding.Namespace, collector.RolesPath)
-}
-
-func (d *DumpIngestor) IngestRole(ctx context.Context, role types.RoleType) error {
-	if ok, err := preflight.CheckRole(role); !ok {
-		entity := tag.Entity(tag.EntityRoles)
-		_ = statsd.Incr(metric.CollectorSkip, append(d.collector.Tags(ctx), entity), 1)
-
-		return err
-	}
-
-	filePath := ingestRolePath(role)
-
-	return d.processObject(ctx, role, filePath)
-}
-
-func (d *DumpIngestor) IngestClusterRole(ctx context.Context, clusterRole types.ClusterRoleType) error {
-	if ok, err := preflight.CheckClusterRole(clusterRole); !ok {
-		entity := tag.Entity(tag.EntityClusterRoles)
-		_ = statsd.Incr(metric.CollectorSkip, append(d.collector.Tags(ctx), entity), 1)
-
-		return err
-	}
-
-	return d.processObject(ctx, clusterRole, collector.ClusterRolesPath)
-}
-
-func ingestRoleBindingPath(roleBinding types.RoleBindingType) string {
-	return path.Join(roleBinding.Namespace, collector.RoleBindingsPath)
-}
-
-func (d *DumpIngestor) IngestRoleBinding(ctx context.Context, roleBinding types.RoleBindingType) error {
-	if ok, err := preflight.CheckRoleBinding(roleBinding); !ok {
-		entity := tag.Entity(tag.EntityRolebindings)
-		_ = statsd.Incr(metric.CollectorSkip, append(d.collector.Tags(ctx), entity), 1)
-
-		return err
-	}
-
-	filePath := ingestRoleBindingPath(roleBinding)
-
-	return d.processObject(ctx, roleBinding, filePath)
-}
-
-func (d *DumpIngestor) IngestClusterRoleBinding(ctx context.Context, clusterRoleBinding types.ClusterRoleBindingType) error {
-	if ok, err := preflight.CheckClusterRoleBinding(clusterRoleBinding); !ok {
-		entity := tag.Entity(tag.EntityClusterRolebindings)
-		_ = statsd.Incr(metric.CollectorSkip, append(d.collector.Tags(ctx), entity), 1)
-
-		return err
-	}
-
-	return d.processObject(ctx, clusterRoleBinding, collector.ClusterRoleBindingsPath)
-}
-
-func ingestEndpointPath(endpoint types.EndpointType) string {
-	return path.Join(endpoint.Namespace, collector.EndpointPath)
-}
-
-func (d *DumpIngestor) IngestEndpoint(ctx context.Context, endpoint types.EndpointType) error {
-	if ok, err := preflight.CheckEndpoint(endpoint); !ok {
-		entity := tag.Entity(tag.EntityEndpoints)
-		_ = statsd.Incr(metric.CollectorSkip, append(d.collector.Tags(ctx), entity), 1)
-
-		return err
-	}
-
-	filePath := ingestEndpointPath(endpoint)
-
-	return d.processObject(ctx, endpoint, filePath)
-}
-
-// Static wrapper to dump k8s object dynamically (streams Kubernetes objects to the collector writer).
-func dumpK8sObjs(ctx context.Context, operationName string, entity string, streamFunc StreamFunc) (context.Context, error) {
-	log.I.Infof("Dumping %s", entity)
-	span, ctx := tracer.StartSpanFromContext(ctx, operationName, tracer.Measured())
-	span.SetTag(tag.EntityTag, entity)
-	defer span.Finish()
-	err := streamFunc(ctx)
-
-	return ctx, err
-}
-
 func (d *DumpIngestor) DumpK8sObjects(ctx context.Context) error {
 	spanDump, ctx := tracer.StartSpanFromContext(ctx, span.CollectorDump, tracer.Measured())
 	defer spanDump.Finish()
 
-	ctx, pipeline, err := newPipelineDumpIngestor(ctx, d)
+	// ctx, pipeline, err := pipeline.NewPipelineDumpIngestor(ctx, d.collector, d.writer)
 	if err != nil {
 		return fmt.Errorf("create pipeline ingestor: %w", err)
 	}
-	spanDump.SetTag(tag.DumperWorkerNumberTag, pipeline.WorkerNumber)
-
-	err = pipeline.Run(ctx)
-	if err != nil {
-		return fmt.Errorf("run pipeline ingestor: %w", err)
-	}
-
-	return pipeline.Wait(ctx)
-}
-
-func marshalK8sObj(obj interface{}) ([]byte, error) {
-	jsonData, err := json.Marshal(obj)
-	if err != nil {
-		return nil, fmt.Errorf("failed to marshal Kubernetes object: %w", err)
-	}
-
-	return jsonData, nil
-}
-
-func (d *DumpIngestor) processObject(ctx context.Context, obj interface{}, filePath string) error {
-	jsonData, err := marshalK8sObj(obj)
-	if err != nil {
-		return err
-	}
-
-	return d.writer.Write(ctx, jsonData, filePath)
-}
+	return nil
+	// spanDump.SetTag(tag.DumperWorkerNumberTag, pipeline.WorkerNumber)
 
-// Complete() is invoked by the collector when all k8s assets have been streamed.
-// The function flushes all writers and waits for completion.
-func (d *DumpIngestor) Complete(ctx context.Context) error {
-	d.writer.Flush(ctx)
+	// err = pipeline.Run(ctx)
+	// if err != nil {
+	// 	return fmt.Errorf("run pipeline ingestor: %w", err)
+	// }
 
-	return nil
+	// return pipeline.Wait(ctx)
 }
 
 // Close() is invoked by the collector to close all handlers used to dump k8s objects.

pkg/dump/pipeline/cluster_role_binding_ingest.go

@@ -0,0 +1,37 @@
+package pipeline
+
+import (
+	"context"
+
+	"github.com/DataDog/KubeHound/pkg/collector"
+	"github.com/DataDog/KubeHound/pkg/dump/writer"
+	"github.com/DataDog/KubeHound/pkg/globals/types"
+	"github.com/DataDog/KubeHound/pkg/kubehound/ingestor/preflight"
+	rbacv1 "k8s.io/api/rbac/v1"
+)
+
+type ClusterRoleBindingIngestor struct {
+	buffer map[string]*rbacv1.ClusterRoleBindingList
+	writer writer.DumperWriter
+}
+
+func NewClusterRoleBindingIngestor(ctx context.Context, dumpWriter writer.DumperWriter) *ClusterRoleBindingIngestor {
+	return &ClusterRoleBindingIngestor{
+		buffer: make(map[string]*rbacv1.ClusterRoleBindingList),
+		writer: dumpWriter,
+	}
+}
+
+func (d *ClusterRoleBindingIngestor) IngestClusterRoleBinding(ctx context.Context, clusterRoleBinding types.ClusterRoleBindingType) error {
+	if ok, err := preflight.CheckClusterRoleBinding(clusterRoleBinding); !ok {
+		return err
+	}
+
+	return bufferObject[rbacv1.ClusterRoleBindingList, types.ClusterRoleBindingType](ctx, collector.ClusterRoleBindingsPath, d.buffer, clusterRoleBinding)
+}
+
+// Complete() is invoked by the collector when all k8s assets have been streamed.
+// The function flushes all writers and waits for completion.
+func (d *ClusterRoleBindingIngestor) Complete(ctx context.Context) error {
+	return dumpObj[*rbacv1.ClusterRoleBindingList](ctx, d.buffer, d.writer)
+}