feat: node-collector image ref (#258)

* feat: node-collector image ref Signed-off-by: chenk <[email protected]> * feat: node-collector image ref Signed-off-by: chenk <[email protected]> * feat: node-collector image ref Signed-off-by: chenk <[email protected]> * feat: node-collector image ref Signed-off-by: chenk <[email protected]> --------- Signed-off-by: chenk <[email protected]>
aquasecurity · Dec 3, 2023 · 50a0691 · 50a0691
1 parent 978b21a
commit 50a0691
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 15 deletions.
diff --git a/examples/trivy.go b/examples/trivy.go
@@ -7,7 +7,7 @@ import (
 
 	"github.com/aquasecurity/trivy-kubernetes/pkg/artifacts"
 	"github.com/aquasecurity/trivy-kubernetes/pkg/k8s"
-	"github.com/aquasecurity/trivy-kubernetes/pkg/trivyk8s"
+	tk "github.com/aquasecurity/trivy-kubernetes/pkg/trivyk8s"
 
 	"go.uber.org/zap"
 	corev1 "k8s.io/api/core/v1"
@@ -30,9 +30,7 @@ func main() {
 
 	fmt.Println("Current namespace:", cluster.GetCurrentNamespace())
 
-	trivyk8sCopy := trivyk8s.New(cluster, logger.Sugar(), trivyk8s.WithExcludeOwned(true))
-	trivyk8s := trivyk8s.New(cluster, logger.Sugar(), trivyk8s.WithExcludeOwned(true))
-
+	trivyk8s := tk.New(cluster, logger.Sugar(), tk.WithExcludeOwned(true))
 	fmt.Println("Scanning cluster")
 
 	//trivy k8s #cluster
@@ -51,13 +49,13 @@ func main() {
 
 	fmt.Println("Scanning namespace 'default'")
 	//trivy k8s --namespace default
-	artifacts, err = trivyk8sCopy.Namespace("default").ListArtifacts(ctx)
+	artifacts, err = trivyk8s.Namespace("default").ListArtifacts(ctx)
 	if err != nil {
 		log.Fatal(err)
 	}
 	printArtifacts(artifacts)
 	fmt.Println("Scanning all namespaces ")
-	artifacts, err = trivyk8sCopy.AllNamespaces().ListArtifacts(ctx)
+	artifacts, err = trivyk8s.AllNamespaces().ListArtifacts(ctx)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -66,7 +64,7 @@ func main() {
 	fmt.Println("Scanning namespace 'default', resource 'deployment/orion'")
 
 	//trivy k8s --namespace default deployment/orion
-	artifact, err := trivyk8sCopy.Namespace("default").GetArtifact(ctx, "deploy", "orion")
+	artifact, err := trivyk8s.Namespace("default").GetArtifact(ctx, "deploy", "orion")
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -75,15 +73,15 @@ func main() {
 	fmt.Println("Scanning 'deployments'")
 
 	//trivy k8s deployment
-	artifacts, err = trivyk8sCopy.Namespace("default").Resources("deployment").ListArtifacts(ctx)
+	artifacts, err = trivyk8s.Namespace("default").Resources("deployment").ListArtifacts(ctx)
 	if err != nil {
 		log.Fatal(err)
 	}
 	printArtifacts(artifacts)
 
 	fmt.Println("Scanning 'cm,pods'")
 	//trivy k8s clusterroles,pods
-	artifacts, err = trivyk8sCopy.Namespace("default").Resources("cm,pods").ListArtifacts(ctx)
+	artifacts, err = trivyk8s.Namespace("default").Resources("cm,pods").ListArtifacts(ctx)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -113,7 +111,10 @@ func main() {
 	}
 
 	// collect node info
-	ar, err := trivyk8sCopy.ListArtifactAndNodeInfo(ctx, "trivy-temp", map[string]string{"chen": "test"}, tolerations...)
+	ar, err := trivyk8s.ListArtifactAndNodeInfo(ctx, []tk.NodeCollectorOption{
+		tk.WithScanJobNamespace("trivy-temp"),
+		tk.WithIgnoreLabels(map[string]string{"chen": "test"}),
+		tk.WithTolerations(tolerations)}...)
 	if err != nil {
 		log.Fatal(err)
 	}

diff --git a/pkg/trivyk8s/trivyk8s.go b/pkg/trivyk8s/trivyk8s.go
@@ -40,7 +40,7 @@ type ArtifactsK8S interface {
 	// GetArtifact return kubernete scanable artifact
 	GetArtifact(context.Context, string, string) (*artifacts.Artifact, error)
 	// ListArtifactAndNodeInfo return kubernete scanable artifact and node info
-	ListArtifactAndNodeInfo(context.Context, string, map[string]string, ...corev1.Toleration) ([]*artifacts.Artifact, error)
+	ListArtifactAndNodeInfo(context.Context, ...NodeCollectorOption) ([]*artifacts.Artifact, error)
 	// ListClusterBomInfo returns kubernetes Bom (node,core components) information.
 	ListClusterBomInfo(context.Context) ([]*artifacts.Artifact, error)
 }
@@ -52,6 +52,7 @@ type client struct {
 	allNamespaces bool
 	logger        *zap.SugaredLogger
 	excludeOwned  bool
+	scanJobParams scanJobParams
 }
 
 type K8sOption func(*client)
@@ -166,8 +167,45 @@ func (c *client) ListArtifacts(ctx context.Context) ([]*artifacts.Artifact, erro
 	return artifactList, nil
 }
 
+type scanJobParams struct {
+	toleration       []corev1.Toleration
+	ignoreLabels     map[string]string
+	scanJobNamespace string
+	imageRef         string
+}
+
+type NodeCollectorOption func(*client)
+
+func WithTolerations(toleration []corev1.Toleration) NodeCollectorOption {
+	return func(c *client) {
+		c.scanJobParams.toleration = toleration
+	}
+}
+
+func WithIgnoreLabels(ignoreLabels map[string]string) NodeCollectorOption {
+	return func(c *client) {
+		c.scanJobParams.ignoreLabels = ignoreLabels
+	}
+}
+
+func WithScanJobNamespace(namespace string) NodeCollectorOption {
+	return func(c *client) {
+		c.scanJobParams.scanJobNamespace = namespace
+	}
+}
+
+func WithScanJobImageRef(imageRef string) NodeCollectorOption {
+	return func(c *client) {
+		c.scanJobParams.imageRef = imageRef
+	}
+}
+
 // ListArtifacts returns kubernetes scannable artifacs.
-func (c *client) ListArtifactAndNodeInfo(ctx context.Context, namespace string, ignoreLabels map[string]string, tolerations ...corev1.Toleration) ([]*artifacts.Artifact, error) {
+func (c *client) ListArtifactAndNodeInfo(ctx context.Context,
+	opts ...NodeCollectorOption) ([]*artifacts.Artifact, error) {
+	for _, opt := range opts {
+		opt(c)
+	}
 	artifactList, err := c.ListArtifacts(ctx)
 	if err != nil {
 		return nil, err
@@ -181,9 +219,10 @@ func (c *client) ListArtifactAndNodeInfo(ctx context.Context, namespace string,
 		c.cluster,
 		jobs.WithTimetout(time.Minute*5),
 		jobs.WithJobTemplateName(jobs.NodeCollectorName),
-		jobs.WithJobNamespace(namespace),
+		jobs.WithJobNamespace(c.scanJobParams.scanJobNamespace),
 		jobs.WithJobLabels(labels),
-		jobs.WithJobTolerations(tolerations),
+		jobs.WithImageRef(c.scanJobParams.imageRef),
+		jobs.WithJobTolerations(c.scanJobParams.toleration),
 	)
 	// delete trivy namespace
 	defer jc.Cleanup(ctx)
@@ -193,7 +232,7 @@ func (c *client) ListArtifactAndNodeInfo(ctx context.Context, namespace string,
 		if resource.Kind != "Node" {
 			continue
 		}
-		if ignoreNodeByLabel(resource, ignoreLabels) {
+		if ignoreNodeByLabel(resource, c.scanJobParams.ignoreLabels) {
 			continue
 		}