Merge pull request #4 from ermetic-research/feat/add-state-flag

Added support for state flag filter on commands

cmd/commands/common/command.go

@@ -18,6 +18,11 @@ func CommandBefore(c *cli.Context) error {
 			return err
 		}
 	}
+	if state := c.String("status"); state != "" {
+		if _, err := cnappgoat.StateFromString(state); err != nil {
+			return err
+		}
+	}
 
 	return nil
 }
@@ -69,5 +74,10 @@ func CommandFlags() []cli.Flag {
 			Usage:   "CNAPPgoat platform to operate on, e.g. 'AWS, 'AZURE', 'GCP'",
 			Aliases: []string{"p"},
 		},
+		&cli.StringFlag{
+			Name:    "status",
+			Usage:   "CNAPPgoat status to filter on, e.g. 'deployed', 'destroyed', 'error', 'not-deployed'",
+			Aliases: []string{"s"},
+		},
 	}
 }

cmd/commands/common/scenario.go

@@ -12,12 +12,13 @@ func GetScenarios(
 	args cli.Args,
 	reg *cnappgoat.Registry,
 	module string,
-	platform string) ([]*cnappgoat.Scenario, error) {
+	platform string,
+	state string) ([]*cnappgoat.Scenario, error) {
 	if scenarios := getScenarios(
 		args,
 		reg,
 		module,
-		platform); len(scenarios) > 0 {
+		platform, state); len(scenarios) > 0 {
 		return scenarios, nil
 	}
 
@@ -54,7 +55,8 @@ func getScenarios(
 	args cli.Args,
 	reg *cnappgoat.Registry,
 	module string,
-	platform string) (scenarios []*cnappgoat.Scenario) {
+	platform string,
+	state string) (scenarios []*cnappgoat.Scenario) {
 	if args.Len() > 0 {
 		for _, arg := range separateArgs(args.Slice()) {
 			scenario, ok := reg.GetScenario(arg)
@@ -75,33 +77,38 @@ func getScenarios(
 				continue
 			}
 
+			if state != "" &&
+				scenario.State.State != strings.ToUpper(state) {
+				logrus.Errorf("Skipping scenario %s because it is not in state %s", scenario.ScenarioParams.ID, state)
+				continue
+			}
+
 			scenarios = append(scenarios, scenario)
 		}
 
-		return
+		return scenarios
 	}
-
-	if module != "" && platform != "" {
-		logrus.Debugf("list all scenarios for %s/%s", module, platform)
-		scenarios = reg.ListScenariosByModuleAndPlatform(cnappgoat.Module(module), cnappgoat.Platform(platform))
-		return
+	var options []cnappgoat.ListScenariosOption
+	if module == "" && platform == "" && state == "" {
+		logrus.Debugf("list all scenarios for all modules")
+		scenarios = reg.ListScenarios()
+		return scenarios
+	}
+	if module != "" { // Assuming Module has a String method
+		options = append(options, cnappgoat.WithModule(cnappgoat.Module(module)))
 	}
 
-	if platform != "" {
-		logrus.Debugf("list all scenarios for %s", platform)
-		scenarios = reg.ListScenariosByPlatform(cnappgoat.Platform(platform))
-		return
+	if platform != "" { // Assuming Platform has a String method
+		options = append(options, cnappgoat.WithPlatform(cnappgoat.Platform(platform)))
 	}
 
-	if module != "" {
-		logrus.Debugf("list all scenarios for %s", module)
-		scenarios = reg.ListScenariosByModule(cnappgoat.Module(module))
-		return
+	if state != "" { // Assuming State has a String method
+		options = append(options, cnappgoat.WithState(cnappgoat.State{State: state}))
 	}
 
-	logrus.Debugf("list all scenarios for all modules")
-	scenarios = reg.ListScenarios()
-	return
+	logrus.Debugf("list scenarios for module %s, platform %s, state %s", module, platform, state)
+	scenarios = reg.ListScenariosWithOptions(options...)
+	return scenarios
 }
 
 func separateArgs(args []string) (positional []string) {

cmd/commands/describe.go

@@ -44,7 +44,8 @@ var DescribeCommand = &cli.Command{
 			c.Args(),
 			reg,
 			c.String("module"),
-			c.String("platform"))
+			c.String("platform"),
+			c.String("status"))
 		if err != nil {
 			return err
 		}

cmd/commands/destroy.go

@@ -23,7 +23,8 @@ var DestroyCommand = &cli.Command{
 				c.Args(),
 				reg,
 				c.String("module"),
-				c.String("platform"))
+				c.String("platform"),
+				c.String("status"))
 		if err != nil {
 			return err
 		}

cmd/commands/list.go

@@ -21,7 +21,8 @@ var ListCommand = &cli.Command{
 			c.Args(),
 			reg,
 			c.String("module"),
-			c.String("platform"))
+			c.String("platform"),
+			c.String("status"))
 		if err != nil {
 			return err
 		}

cmd/commands/provision.go

@@ -23,7 +23,8 @@ var ProvisionCommand = &cli.Command{
 				c.Args(),
 				reg,
 				c.String("module"),
-				c.String("platform"))
+				c.String("platform"),
+				c.String("status"))
 		if err != nil {
 			return err
 		}

registry.go

@@ -10,6 +10,14 @@ type Registry struct {
 	storage   *LocalStorage
 }
 
+type ListScenariosOptions struct {
+	Module   Module
+	Platform Platform
+	State    State
+}
+
+type ListScenariosOption func(options *ListScenariosOptions)
+
 func NewRegistry(s *LocalStorage) (*Registry, error) {
 	registry := Registry{
 		scenarios: make(map[string]*Scenario),
@@ -67,10 +75,17 @@ func (r *Registry) ListScenarios() []*Scenario {
 	return scenarios
 }
 
-func (r *Registry) ListScenariosByModule(module Module) []*Scenario {
+func (r *Registry) ListScenariosWithOptions(opts ...ListScenariosOption) []*Scenario {
+	options := ListScenariosOptions{}
+	for _, opt := range opts {
+		opt(&options)
+	}
+
 	var scenarios []*Scenario
 	for name := range r.scenarios {
-		if r.scenarios[name].ScenarioParams.Module.Equals(module) || module.Equals("") {
+		if (r.scenarios[name].ScenarioParams.Platform.Equals(options.Platform) || options.Platform.Equals("")) &&
+			(r.scenarios[name].ScenarioParams.Module.Equals(options.Module) || options.Module.Equals("")) &&
+			(r.scenarios[name].State.Equals(options.State) || options.State.Equals(State{})) {
 			scenarios = append(scenarios, r.scenarios[name])
 		}
 	}
@@ -79,29 +94,22 @@ func (r *Registry) ListScenariosByModule(module Module) []*Scenario {
 	return scenarios
 }
 
-func (r *Registry) ListScenariosByPlatform(platform Platform) []*Scenario {
-	var scenarios []*Scenario
-	for name := range r.scenarios {
-		if r.scenarios[name].ScenarioParams.Platform.Equals(platform) || platform.Equals("") {
-			scenarios = append(scenarios, r.scenarios[name])
-		}
+func WithModule(module Module) ListScenariosOption {
+	return func(opts *ListScenariosOptions) {
+		opts.Module = module
 	}
-
-	sortScenarios(scenarios)
-	return scenarios
 }
 
-func (r *Registry) ListScenariosByModuleAndPlatform(module Module, platform Platform) []*Scenario {
-	var scenarios []*Scenario
-	for name := range r.scenarios {
-		if r.scenarios[name].ScenarioParams.Module.Equals(module) &&
-			r.scenarios[name].ScenarioParams.Platform.Equals(platform) {
-			scenarios = append(scenarios, r.scenarios[name])
-		}
+func WithPlatform(platform Platform) ListScenariosOption {
+	return func(opts *ListScenariosOptions) {
+		opts.Platform = platform
 	}
+}
 
-	sortScenarios(scenarios)
-	return scenarios
+func WithState(state State) ListScenariosOption {
+	return func(opts *ListScenariosOptions) {
+		opts.State = state
+	}
 }
 
 func (r *Registry) SetState(scenario *Scenario, state State) error {

scenario.go

@@ -41,7 +41,7 @@ const (
 	Deployed    = "deployed"
 	Destroyed   = "destroyed"
 	Error       = "error"
-	NotDeployed = "not deployed"
+	NotDeployed = "not-deployed"
 )
 
 type Module string
@@ -149,6 +149,10 @@ func (s State) String() string {
 	return s.State
 }
 
+func (s State) Equals(state State) bool {
+	return strings.EqualFold(s.State, state.State)
+}
+
 func ModuleFromString(name string) (Module, error) {
 	switch strings.ToLower(name) {
 	case strings.ToLower(CSPM.String()):
@@ -176,3 +180,18 @@ func PlatformFromString(name string) (Platform, error) {
 		return "", errors.New("unknown platform name: " + name)
 	}
 }
+
+func StateFromString(state string) (State, error) {
+	switch strings.ToLower(state) {
+	case strings.ToLower(Deployed):
+		return State{State: Deployed}, nil
+	case strings.ToLower(Destroyed):
+		return State{State: Destroyed}, nil
+	case strings.ToLower(Error):
+		return State{State: Error}, nil
+	case strings.ToLower(NotDeployed):
+		return State{State: NotDeployed}, nil
+	default:
+		return State{}, errors.New("unknown state: " + state)
+	}
+}