feat: add node toleration option (#3823)

go.mod

@@ -100,6 +100,7 @@ require (
 	google.golang.org/protobuf v1.28.1
 	gopkg.in/yaml.v3 v3.0.1
 	gotest.tools v2.2.0+incompatible
+	k8s.io/api v0.26.2
 	k8s.io/utils v0.0.0-20230115233650-391b47cb4029
 	modernc.org/sqlite v1.20.3
 )
@@ -360,7 +361,6 @@ require (
 	gopkg.in/warnings.v0 v0.1.2 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	helm.sh/helm/v3 v3.11.1 // indirect
-	k8s.io/api v0.26.2 // indirect
 	k8s.io/apiextensions-apiserver v0.26.0 // indirect
 	k8s.io/apimachinery v0.26.2 // indirect
 	k8s.io/apiserver v0.26.0 // indirect

pkg/flag/kubernetes_flags.go

@@ -1,5 +1,14 @@
 package flag
 
+import (
+	"fmt"
+	"strconv"
+	"strings"
+
+	"github.com/samber/lo"
+	corev1 "k8s.io/api/core/v1"
+)
+
 var (
 	ClusterContextFlag = Flag{
 		Name:       "context",
@@ -38,6 +47,12 @@ var (
 		Value:      "",
 		Usage:      "specify k8s version to validate outdated api by it (example: 1.21.0)",
 	}
+	TolerationsFlag = Flag{
+		Name:       "tolerations",
+		ConfigName: "kubernetes.tolerations",
+		Value:      []string{},
+		Usage:      "specify node-collector job tolerations (example: key1=value1:NoExecute,key2=value2:NoSchedule)",
+	}
 )
 
 type K8sFlagGroup struct {
@@ -46,6 +61,7 @@ type K8sFlagGroup struct {
 	KubeConfig     *Flag
 	Components     *Flag
 	K8sVersion     *Flag
+	Tolerations    *Flag
 }
 
 type K8sOptions struct {
@@ -54,6 +70,7 @@ type K8sOptions struct {
 	KubeConfig     string
 	Components     []string
 	K8sVersion     string
+	Tolerations    []corev1.Toleration
 }
 
 func NewK8sFlagGroup() *K8sFlagGroup {
@@ -63,6 +80,7 @@ func NewK8sFlagGroup() *K8sFlagGroup {
 		KubeConfig:     &KubeConfigFlag,
 		Components:     &ComponentsFlag,
 		K8sVersion:     &K8sVersionFlag,
+		Tolerations:    &TolerationsFlag,
 	}
 }
 
@@ -77,15 +95,58 @@ func (f *K8sFlagGroup) Flags() []*Flag {
 		f.KubeConfig,
 		f.Components,
 		f.K8sVersion,
+		f.Tolerations,
 	}
 }
 
-func (f *K8sFlagGroup) ToOptions() K8sOptions {
+func (f *K8sFlagGroup) ToOptions() (K8sOptions, error) {
+	tolerations, err := optionToTolerations(getStringSlice(f.Tolerations))
+	if err != nil {
+		return K8sOptions{}, err
+	}
 	return K8sOptions{
 		ClusterContext: getString(f.ClusterContext),
 		Namespace:      getString(f.Namespace),
 		KubeConfig:     getString(f.KubeConfig),
 		Components:     getStringSlice(f.Components),
 		K8sVersion:     getString(f.K8sVersion),
+		Tolerations:    tolerations,
+	}, nil
+}
+
+func optionToTolerations(tolerationsOptions []string) ([]corev1.Toleration, error) {
+	tolerations := make([]corev1.Toleration, 0)
+	for _, toleration := range tolerationsOptions {
+		tolerationParts := strings.Split(toleration, ":")
+		if len(tolerationParts) < 2 {
+			return []corev1.Toleration{}, fmt.Errorf("toleration must include key and effect")
+		}
+		if corev1.TaintEffect(tolerationParts[1]) != corev1.TaintEffectNoSchedule &&
+			corev1.TaintEffect(tolerationParts[1]) != corev1.TaintEffectPreferNoSchedule &&
+			corev1.TaintEffect(tolerationParts[1]) != corev1.TaintEffectNoExecute {
+			return []corev1.Toleration{}, fmt.Errorf("toleration effect must be a valid value")
+		}
+		keyValue := strings.Split(tolerationParts[0], "=")
+		operator := corev1.TolerationOpEqual
+		if len(keyValue[1]) == 0 {
+			operator = corev1.TolerationOpExists
+		}
+		toleration := corev1.Toleration{
+			Key:      keyValue[0],
+			Value:    keyValue[1],
+			Operator: operator,
+			Effect:   corev1.TaintEffect(tolerationParts[1]),
+		}
+		var tolerationSec int
+		var err error
+		if len(tolerationParts) == 3 {
+			tolerationSec, err = strconv.Atoi(tolerationParts[2])
+			if err != nil {
+				return nil, fmt.Errorf("TolerationSeconds must must be a number")
+			}
+		}
+		toleration.TolerationSeconds = lo.ToPtr(int64(tolerationSec))
+		tolerations = append(tolerations, toleration)
 	}
+	return tolerations, nil
 }

pkg/flag/options.go

@@ -394,7 +394,10 @@ func (f *Flags) ToOptions(appVersion string, args []string, globalFlags *GlobalF
 	}
 
 	if f.K8sFlagGroup != nil {
-		opts.K8sOptions = f.K8sFlagGroup.ToOptions()
+		opts.K8sOptions, err = f.K8sFlagGroup.ToOptions()
+		if err != nil {
+			return Options{}, xerrors.Errorf("k8s flag error: %w", err)
+		}
 	}
 
 	if f.LicenseFlagGroup != nil {

pkg/k8s/commands/cluster.go

@@ -22,7 +22,7 @@ func clusterRun(ctx context.Context, opts flag.Options, cluster k8s.Cluster) err
 	var artifacts []*artifacts.Artifact
 	var err error
 	if opts.Scanners.AnyEnabled(types.MisconfigScanner) && slices.Contains(opts.Components, "infra") {
-		artifacts, err = trivyk8s.New(cluster, log.Logger).ListArtifactAndNodeInfo(ctx)
+		artifacts, err = trivyk8s.New(cluster, log.Logger).ListArtifactAndNodeInfo(ctx, opts.Tolerations...)
 		if err != nil {
 			return xerrors.Errorf("get k8s artifacts with node info error: %w", err)
 		}