test(flag): enhance TestCustomFlagGroups to include environment variable support

- Added a test case to verify the behavior of the `TRIVY_FOO` environment variable in the `TestCustomFlagGroups` function.
- Ensured that the flag system correctly retrieves values from both environment variables and custom flags.
- Improved test coverage for the flag handling functionality.

pkg/commands/app.go

@@ -19,6 +19,7 @@ import (
 	"github.com/aquasecurity/trivy/pkg/commands/clean"
 	"github.com/aquasecurity/trivy/pkg/commands/convert"
 	"github.com/aquasecurity/trivy/pkg/commands/server"
+	"github.com/aquasecurity/trivy/pkg/extension"
 	"github.com/aquasecurity/trivy/pkg/fanal/analyzer"
 	"github.com/aquasecurity/trivy/pkg/flag"
 	k8scommands "github.com/aquasecurity/trivy/pkg/k8s/commands"
@@ -280,6 +281,7 @@ func NewImageCommand(globalFlags *flag.GlobalFlagGroup) *cobra.Command {
 		flag.NewSecretFlagGroup(),
 		flag.NewVulnerabilityFlagGroup(),
 	}
+	imageFlags = append(imageFlags, extension.CustomFlagGroups("image")...)
 
 	cmd := &cobra.Command{
 		Use:     "image [flags] IMAGE_NAME",
@@ -362,6 +364,7 @@ func NewFilesystemCommand(globalFlags *flag.GlobalFlagGroup) *cobra.Command {
 		flag.NewSecretFlagGroup(),
 		flag.NewVulnerabilityFlagGroup(),
 	}
+	fsFlags = append(fsFlags, extension.CustomFlagGroups("filesystem")...)
 
 	cmd := &cobra.Command{
 		Use:     "filesystem [flags] PATH",
@@ -428,6 +431,7 @@ func NewRootfsCommand(globalFlags *flag.GlobalFlagGroup) *cobra.Command {
 		flag.NewSecretFlagGroup(),
 		flag.NewVulnerabilityFlagGroup(),
 	}
+	rootfsFlags = append(rootfsFlags, extension.CustomFlagGroups("rootfs")...)
 
 	cmd := &cobra.Command{
 		Use:     "rootfs [flags] ROOTDIR",
@@ -493,6 +497,7 @@ func NewRepositoryCommand(globalFlags *flag.GlobalFlagGroup) *cobra.Command {
 		flag.NewVulnerabilityFlagGroup(),
 		flag.NewRepoFlagGroup(),
 	}
+	repoFlags = append(repoFlags, extension.CustomFlagGroups("repository")...)
 
 	cmd := &cobra.Command{
 		Use:     "repository [flags] (REPO_PATH | REPO_URL)",
@@ -605,6 +610,7 @@ func NewClientCommand(globalFlags *flag.GlobalFlagGroup) *cobra.Command {
 		flag.NewScanFlagGroup(),
 		flag.NewVulnerabilityFlagGroup(),
 	}
+	clientFlags = append(clientFlags, extension.CustomFlagGroups("client")...)
 
 	cmd := &cobra.Command{
 		Use:     "client [flags] IMAGE_NAME",
@@ -653,6 +659,7 @@ func NewServerCommand(globalFlags *flag.GlobalFlagGroup) *cobra.Command {
 		flag.NewServerFlags(),
 		flag.NewRegistryFlagGroup(),
 	}
+	serverFlags = append(serverFlags, extension.CustomFlagGroups("server")...)
 
 	cmd := &cobra.Command{
 		Use:     "server [flags]",
@@ -704,7 +711,7 @@ func NewConfigCommand(globalFlags *flag.GlobalFlagGroup) *cobra.Command {
 	cacheFlagGroup := flag.NewCacheFlagGroup()
 	cacheFlagGroup.CacheBackend.Default = string(cache.TypeMemory)
 
-	configFlags := &flag.Flags{
+	configFlags := flag.Flags{
 		globalFlags,
 		cacheFlagGroup,
 		flag.NewMisconfFlagGroup(),
@@ -718,6 +725,7 @@ func NewConfigCommand(globalFlags *flag.GlobalFlagGroup) *cobra.Command {
 		reportFlagGroup,
 		scanFlags,
 	}
+	configFlags = append(configFlags, extension.CustomFlagGroups("config")...)
 
 	cmd := &cobra.Command{
 		Use:     "config [flags] DIR",
@@ -1024,7 +1032,7 @@ func NewKubernetesCommand(globalFlags *flag.GlobalFlagGroup) *cobra.Command {
 	packageFlagGroup := flag.NewPackageFlagGroup()
 	packageFlagGroup.IncludeDevDeps = nil // disable '--include-dev-deps'
 
-	k8sFlags := &flag.Flags{
+	k8sFlags := flag.Flags{
 		globalFlags,
 		flag.NewCacheFlagGroup(),
 		flag.NewDBFlagGroup(),
@@ -1039,6 +1047,7 @@ func NewKubernetesCommand(globalFlags *flag.GlobalFlagGroup) *cobra.Command {
 		flag.NewRegistryFlagGroup(),
 		flag.NewVulnerabilityFlagGroup(),
 	}
+	k8sFlags = append(k8sFlags, extension.CustomFlagGroups("kubernetes")...)
 
 	cmd := &cobra.Command{
 		Use:     "kubernetes [flags] [CONTEXT]",
@@ -1095,7 +1104,7 @@ func NewVMCommand(globalFlags *flag.GlobalFlagGroup) *cobra.Command {
 	misconfFlagGroup.CloudformationParamVars = nil // disable '--cf-params'
 	misconfFlagGroup.TerraformTFVars = nil         // disable '--tf-vars'
 
-	vmFlags := &flag.Flags{
+	vmFlags := flag.Flags{
 		globalFlags,
 		flag.NewCacheFlagGroup(),
 		flag.NewDBFlagGroup(),
@@ -1115,6 +1124,7 @@ func NewVMCommand(globalFlags *flag.GlobalFlagGroup) *cobra.Command {
 			},
 		},
 	}
+	vmFlags = append(vmFlags, extension.CustomFlagGroups("vm")...)
 
 	cmd := &cobra.Command{
 		Use:     "vm [flags] VM_IMAGE",
@@ -1185,7 +1195,7 @@ func NewSBOMCommand(globalFlags *flag.GlobalFlagGroup) *cobra.Command {
 	packageFlagGroup := flag.NewPackageFlagGroup()
 	packageFlagGroup.IncludeDevDeps = nil // disable '--include-dev-deps'
 
-	sbomFlags := &flag.Flags{
+	sbomFlags := flag.Flags{
 		globalFlags,
 		cacheFlagGroup,
 		flag.NewDBFlagGroup(),
@@ -1197,6 +1207,7 @@ func NewSBOMCommand(globalFlags *flag.GlobalFlagGroup) *cobra.Command {
 		flag.NewVulnerabilityFlagGroup(),
 		licenseFlagGroup,
 	}
+	sbomFlags = append(sbomFlags, extension.CustomFlagGroups("sbom")...)
 
 	cmd := &cobra.Command{
 		Use:     "sbom [flags] SBOM_PATH",

pkg/extension/flag.go

@@ -0,0 +1,40 @@
+package extension
+
+import (
+	"github.com/samber/lo"
+
+	"github.com/aquasecurity/trivy/pkg/flag"
+)
+
+var flagExtensions = make(map[string]FlagExtension)
+
+func RegisterFlagExtension(extension FlagExtension) {
+	flagExtensions[extension.Name()] = extension
+}
+
+func DeregisterFlagExtension(name string) {
+	delete(flagExtensions, name)
+}
+
+// FlagExtension is an extension that allows adding custom CLI flags.
+type FlagExtension interface {
+	Name() string
+
+	// CustomFlagGroup returns custom flag group to be added to Trivy CLI.
+	// The command parameter specifies which command the flags are for.
+	// If the command is empty, the flags will be applied to all commands.
+	CustomFlagGroup(command string) flag.FlagGroup
+}
+
+// CustomFlagGroups collects all flag groups from registered extensions for a specific command.
+func CustomFlagGroups(command string) []flag.FlagGroup {
+	var flagGroups []flag.FlagGroup
+	for _, e := range flagExtensions {
+		group := e.CustomFlagGroup(command)
+		if lo.IsNil(group) {
+			continue
+		}
+		flagGroups = append(flagGroups, group)
+	}
+	return flagGroups
+}

pkg/extension/flag_test.go

@@ -0,0 +1,134 @@
+package extension_test
+
+import (
+	"testing"
+
+	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/aquasecurity/trivy/pkg/extension"
+	"github.com/aquasecurity/trivy/pkg/flag"
+)
+
+type testOptionKey struct{}
+
+var foo = flag.Flag[string]{
+	Name:       "foo",
+	ConfigName: "foo",
+	Usage:      "foo",
+	Default:    "default-value",
+}
+
+// testFlagGroup is a flag group for testing
+type testFlagGroup struct {
+	Foo *flag.Flag[string]
+}
+
+type testOptions struct {
+	Foo string
+}
+
+func (fg *testFlagGroup) Name() string {
+	return "TestFlagGroup"
+}
+
+func (fg *testFlagGroup) Flags() []flag.Flagger {
+	return []flag.Flagger{
+		fg.Foo,
+	}
+}
+
+func (fg *testFlagGroup) ToOptions(opts *flag.Options) error {
+	if opts.CustomOptions == nil {
+		opts.CustomOptions = make(map[any]any)
+	}
+	opts.CustomOptions[testOptionKey{}] = testOptions{
+		Foo: fg.Foo.Value(),
+	}
+	return nil
+}
+
+// testExtension implements the FlagExtension interface for testing
+type testExtension struct{}
+
+func (e *testExtension) Name() string {
+	return "TestExtension"
+}
+
+func (e *testExtension) CustomFlagGroup(command string) flag.FlagGroup {
+	if command != "image" {
+		return nil
+	}
+
+	return &testFlagGroup{
+		Foo: foo.Clone(),
+	}
+}
+
+func TestCustomFlagGroups(t *testing.T) {
+	// Set up
+	te := &testExtension{}
+	extension.RegisterFlagExtension(te)
+	t.Cleanup(func() {
+		extension.DeregisterFlagExtension(te.Name())
+	})
+
+	t.Run("flag group is set", func(t *testing.T) {
+		t.Cleanup(viper.Reset)
+		flags := flag.Flags(extension.CustomFlagGroups("image"))
+		cmd := &cobra.Command{}
+		flags.AddFlags(cmd)
+		flags.Bind(cmd)
+
+		// Test with no custom value
+		opts, err := flags.ToOptions(nil)
+		require.NoError(t, err)
+
+		// Verify CustomOptions has the default value
+		testOpts := extractTestOptions(t, opts)
+		assert.Equal(t, "default-value", testOpts.Foo)
+
+		// Test with environment variable
+		t.Setenv("TRIVY_FOO", "env-value")
+		opts, err = flags.ToOptions(nil)
+		require.NoError(t, err)
+		testOpts = extractTestOptions(t, opts)
+		assert.Equal(t, "env-value", testOpts.Foo)
+
+		// Test with flag
+		viper.Set(foo.ConfigName, "custom-value")
+		opts, err = flags.ToOptions(nil)
+		require.NoError(t, err)
+
+		// Verify CustomOptions has the custom value
+		testOpts = extractTestOptions(t, opts)
+		assert.Equal(t, "custom-value", testOpts.Foo)
+	})
+
+	t.Run("flag group is not set", func(t *testing.T) {
+		t.Cleanup(viper.Reset)
+		flags := flag.Flags(extension.CustomFlagGroups("other"))
+		cmd := &cobra.Command{}
+		flags.AddFlags(cmd)
+		flags.Bind(cmd)
+
+		// Test
+		viper.Set(foo.ConfigName, "custom-value")
+		opts, err := flags.ToOptions(nil)
+		require.NoError(t, err)
+
+		// Verify CustomOptions is not set
+		require.Nil(t, opts.CustomOptions)
+	})
+}
+
+func extractTestOptions(t *testing.T, opts flag.Options) testOptions {
+	value, ok := opts.CustomOptions[testOptionKey{}]
+	require.True(t, ok)
+
+	testOpts, ok := value.(testOptions)
+	require.True(t, ok)
+	return testOpts
+}

pkg/flag/options.go

@@ -383,6 +383,9 @@ type Options struct {
 	SecretOptions
 	VulnerabilityOptions
 
+	// CustomOptions is a map of custom options.
+	CustomOptions map[any]any
+
 	// Trivy's version, not populated via CLI flags
 	AppVersion string