Browse Source

refact "cscli bouncers" (#2776)

mmetc 1 year ago
parent
commit
4192af30d5

+ 187 - 167
cmd/crowdsec-cli/bouncers.go

@@ -4,7 +4,8 @@ import (
 	"encoding/csv"
 	"encoding/json"
 	"fmt"
-	"io"
+	"os"
+	"slices"
 	"strings"
 	"time"
 
@@ -12,61 +13,41 @@ import (
 	"github.com/fatih/color"
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
-	"slices"
 
 	"github.com/crowdsecurity/crowdsec/cmd/crowdsec-cli/require"
 	middlewares "github.com/crowdsecurity/crowdsec/pkg/apiserver/middlewares/v1"
+	"github.com/crowdsecurity/crowdsec/pkg/csconfig"
 	"github.com/crowdsecurity/crowdsec/pkg/database"
 	"github.com/crowdsecurity/crowdsec/pkg/types"
 )
 
-func getBouncers(out io.Writer, dbClient *database.Client) error {
-	bouncers, err := dbClient.ListBouncers()
-	if err != nil {
-		return fmt.Errorf("unable to list bouncers: %s", err)
+func askYesNo(message string, defaultAnswer bool) (bool, error) {
+	var answer bool
+
+	prompt := &survey.Confirm{
+		Message: message,
+		Default: defaultAnswer,
 	}
 
-	switch csConfig.Cscli.Output {
-	case "human":
-		getBouncersTable(out, bouncers)
-	case "json":
-		enc := json.NewEncoder(out)
-		enc.SetIndent("", "  ")
-		if err := enc.Encode(bouncers); err != nil {
-			return fmt.Errorf("failed to unmarshal: %w", err)
-		}
-		return nil
-	case "raw":
-		csvwriter := csv.NewWriter(out)
-		err := csvwriter.Write([]string{"name", "ip", "revoked", "last_pull", "type", "version", "auth_type"})
-		if err != nil {
-			return fmt.Errorf("failed to write raw header: %w", err)
-		}
-		for _, b := range bouncers {
-			var revoked string
-			if !b.Revoked {
-				revoked = "validated"
-			} else {
-				revoked = "pending"
-			}
-			err := csvwriter.Write([]string{b.Name, b.IPAddress, revoked, b.LastPull.Format(time.RFC3339), b.Type, b.Version, b.AuthType})
-			if err != nil {
-				return fmt.Errorf("failed to write raw: %w", err)
-			}
-		}
-		csvwriter.Flush()
+	if err := survey.AskOne(prompt, &answer); err != nil {
+		return defaultAnswer, err
 	}
 
-	return nil
+	return answer, nil
 }
 
-type cliBouncers struct {}
+type cliBouncers struct {
+	db *database.Client
+	cfg func() *csconfig.Config
+}
 
-func NewCLIBouncers() *cliBouncers {
-	return &cliBouncers{}
+func NewCLIBouncers(getconfig func() *csconfig.Config) *cliBouncers {
+	return &cliBouncers{
+		cfg: getconfig,
+	}
 }
 
-func (cli cliBouncers) NewCommand() *cobra.Command {
+func (cli *cliBouncers) NewCommand() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "bouncers [action]",
 		Short: "Manage bouncers [requires local API]",
@@ -76,94 +57,127 @@ Note: This command requires database direct access, so is intended to be run on
 		Args:              cobra.MinimumNArgs(1),
 		Aliases:           []string{"bouncer"},
 		DisableAutoGenTag: true,
-		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
+		PersistentPreRunE: func(_ *cobra.Command, _ []string) error {
 			var err error
-			if err = require.LAPI(csConfig); err != nil {
+			if err = require.LAPI(cli.cfg()); err != nil {
 				return err
 			}
 
-			dbClient, err = database.NewClient(csConfig.DbConfig)
+			cli.db, err = database.NewClient(cli.cfg().DbConfig)
 			if err != nil {
-				return fmt.Errorf("unable to create new database client: %s", err)
+				return fmt.Errorf("can't connect to the database: %s", err)
 			}
+
 			return nil
 		},
 	}
 
-	cmd.AddCommand(cli.NewListCmd())
-	cmd.AddCommand(cli.NewAddCmd())
-	cmd.AddCommand(cli.NewDeleteCmd())
-	cmd.AddCommand(cli.NewPruneCmd())
+	cmd.AddCommand(cli.newListCmd())
+	cmd.AddCommand(cli.newAddCmd())
+	cmd.AddCommand(cli.newDeleteCmd())
+	cmd.AddCommand(cli.newPruneCmd())
 
 	return cmd
 }
 
-func (cli cliBouncers) NewListCmd() *cobra.Command {
+func (cli *cliBouncers) list() error {
+	out := color.Output
+
+	bouncers, err := cli.db.ListBouncers()
+	if err != nil {
+		return fmt.Errorf("unable to list bouncers: %s", err)
+	}
+
+	switch cli.cfg().Cscli.Output {
+	case "human":
+		getBouncersTable(out, bouncers)
+	case "json":
+		enc := json.NewEncoder(out)
+		enc.SetIndent("", "  ")
+
+		if err := enc.Encode(bouncers); err != nil {
+			return fmt.Errorf("failed to marshal: %w", err)
+		}
+
+		return nil
+	case "raw":
+		csvwriter := csv.NewWriter(out)
+
+		if err := csvwriter.Write([]string{"name", "ip", "revoked", "last_pull", "type", "version", "auth_type"}); err != nil {
+			return fmt.Errorf("failed to write raw header: %w", err)
+		}
+
+		for _, b := range bouncers {
+			valid := "validated"
+			if b.Revoked {
+				valid = "pending"
+			}
+
+			if err := csvwriter.Write([]string{b.Name, b.IPAddress, valid, b.LastPull.Format(time.RFC3339), b.Type, b.Version, b.AuthType}); err != nil {
+				return fmt.Errorf("failed to write raw: %w", err)
+			}
+		}
+
+		csvwriter.Flush()
+	}
+
+	return nil
+}
+
+func (cli *cliBouncers) newListCmd() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:               "list",
 		Short:             "list all bouncers within the database",
 		Example:           `cscli bouncers list`,
 		Args:              cobra.ExactArgs(0),
 		DisableAutoGenTag: true,
-		RunE: func(cmd *cobra.Command, arg []string) error {
-			err := getBouncers(color.Output, dbClient)
-			if err != nil {
-				return fmt.Errorf("unable to list bouncers: %s", err)
-			}
-			return nil
+		RunE: func(_ *cobra.Command, _ []string) error {
+			return cli.list()
 		},
 	}
 
 	return cmd
 }
 
-func (cli cliBouncers) add(cmd *cobra.Command, args []string) error {
-	keyLength := 32
-
-	flags := cmd.Flags()
-
-	key, err := flags.GetString("key")
-	if err != nil {
-		return err
-	}
+func (cli *cliBouncers) add(bouncerName string, key string) error {
+	var err error
 
-	keyName := args[0]
-	var apiKey string
+	keyLength := 32
 
-	if keyName == "" {
-		return fmt.Errorf("please provide a name for the api key")
-	}
-	apiKey = key
 	if key == "" {
-		apiKey, err = middlewares.GenerateAPIKey(keyLength)
-	}
-	if err != nil {
-		return fmt.Errorf("unable to generate api key: %s", err)
+		key, err = middlewares.GenerateAPIKey(keyLength)
+		if err != nil {
+			return fmt.Errorf("unable to generate api key: %s", err)
+		}
 	}
-	_, err = dbClient.CreateBouncer(keyName, "", middlewares.HashSHA512(apiKey), types.ApiKeyAuthType)
+
+	_, err = cli.db.CreateBouncer(bouncerName, "", middlewares.HashSHA512(key), types.ApiKeyAuthType)
 	if err != nil {
 		return fmt.Errorf("unable to create bouncer: %s", err)
 	}
 
-	switch csConfig.Cscli.Output {
+	switch cli.cfg().Cscli.Output {
 	case "human":
-		fmt.Printf("API key for '%s':\n\n", keyName)
-		fmt.Printf("   %s\n\n", apiKey)
+		fmt.Printf("API key for '%s':\n\n", bouncerName)
+		fmt.Printf("   %s\n\n", key)
 		fmt.Print("Please keep this key since you will not be able to retrieve it!\n")
 	case "raw":
-		fmt.Printf("%s", apiKey)
+		fmt.Print(key)
 	case "json":
-		j, err := json.Marshal(apiKey)
+		j, err := json.Marshal(key)
 		if err != nil {
 			return fmt.Errorf("unable to marshal api key")
 		}
-		fmt.Printf("%s", string(j))
+
+		fmt.Print(string(j))
 	}
 
 	return nil
 }
 
-func (cli cliBouncers) NewAddCmd() *cobra.Command {
+func (cli *cliBouncers) newAddCmd() *cobra.Command {
+	var key string
+
 	cmd := &cobra.Command{
 		Use:   "add MyBouncerName",
 		Short: "add a single bouncer to the database",
@@ -171,127 +185,133 @@ func (cli cliBouncers) NewAddCmd() *cobra.Command {
 cscli bouncers add MyBouncerName --key <random-key>`,
 		Args:              cobra.ExactArgs(1),
 		DisableAutoGenTag: true,
-		RunE:              cli.add,
+		RunE: func(_ *cobra.Command, args []string) error {
+			return cli.add(args[0], key)
+		},
 	}
 
 	flags := cmd.Flags()
 	flags.StringP("length", "l", "", "length of the api key")
 	flags.MarkDeprecated("length", "use --key instead")
-	flags.StringP("key", "k", "", "api key for the bouncer")
+	flags.StringVarP(&key, "key", "k", "", "api key for the bouncer")
 
 	return cmd
 }
 
-func (cli cliBouncers) delete(cmd *cobra.Command, args []string) error {
-	for _, bouncerID := range args {
-		err := dbClient.DeleteBouncer(bouncerID)
+func (cli *cliBouncers) deleteValid(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	bouncers, err := cli.db.ListBouncers()
+	if err != nil {
+		cobra.CompError("unable to list bouncers " + err.Error())
+	}
+
+	ret :=[]string{}
+
+	for _, bouncer := range bouncers {
+		if strings.Contains(bouncer.Name, toComplete) && !slices.Contains(args, bouncer.Name) {
+			ret = append(ret, bouncer.Name)
+		}
+	}
+
+	return ret, cobra.ShellCompDirectiveNoFileComp
+}
+
+func (cli *cliBouncers) delete(bouncers []string) error {
+	for _, bouncerID := range bouncers {
+		err := cli.db.DeleteBouncer(bouncerID)
 		if err != nil {
 			return fmt.Errorf("unable to delete bouncer '%s': %s", bouncerID, err)
 		}
+
 		log.Infof("bouncer '%s' deleted successfully", bouncerID)
 	}
 
 	return nil
 }
 
-func (cli cliBouncers) NewDeleteCmd() *cobra.Command {
+func (cli *cliBouncers) newDeleteCmd() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:               "delete MyBouncerName",
 		Short:             "delete bouncer(s) from the database",
 		Args:              cobra.MinimumNArgs(1),
 		Aliases:           []string{"remove"},
 		DisableAutoGenTag: true,
-		ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
-			var err error
-			dbClient, err = getDBClient()
-			if err != nil {
-				cobra.CompError("unable to create new database client: " + err.Error())
-				return nil, cobra.ShellCompDirectiveNoFileComp
-			}
-			bouncers, err := dbClient.ListBouncers()
-			if err != nil {
-				cobra.CompError("unable to list bouncers " + err.Error())
-			}
-			ret := make([]string, 0)
-			for _, bouncer := range bouncers {
-				if strings.Contains(bouncer.Name, toComplete) && !slices.Contains(args, bouncer.Name) {
-					ret = append(ret, bouncer.Name)
-				}
-			}
-			return ret, cobra.ShellCompDirectiveNoFileComp
+		ValidArgsFunction: cli.deleteValid,
+		RunE: func(_ *cobra.Command, args []string) error {
+			return cli.delete(args)
 		},
-		RunE: cli.delete,
 	}
 
 	return cmd
 }
 
-func (cli cliBouncers) NewPruneCmd() *cobra.Command {
-	var parsedDuration time.Duration
+func (cli *cliBouncers) prune(duration time.Duration, force bool) error {
+	if duration < 2*time.Minute {
+		if yes, err := askYesNo(
+				"The duration you provided is less than 2 minutes. " +
+				"This may remove active bouncers. Continue?", false); err != nil {
+			return err
+		} else if !yes {
+			fmt.Println("User aborted prune. No changes were made.")
+			return nil
+		}
+	}
+
+	bouncers, err := cli.db.QueryBouncersLastPulltimeLT(time.Now().UTC().Add(duration))
+	if err != nil {
+		return fmt.Errorf("unable to query bouncers: %w", err)
+	}
+
+	if len(bouncers) == 0 {
+		fmt.Println("No bouncers to prune.")
+		return nil
+	}
+
+	getBouncersTable(color.Output, bouncers)
+
+	if !force {
+		if yes, err := askYesNo(
+				"You are about to PERMANENTLY remove the above bouncers from the database. " +
+				"These will NOT be recoverable. Continue?", false); err != nil {
+			return err
+		} else if !yes {
+			fmt.Println("User aborted prune. No changes were made.")
+			return nil
+		}
+	}
+
+	deleted, err := cli.db.BulkDeleteBouncers(bouncers)
+	if err != nil {
+		return fmt.Errorf("unable to prune bouncers: %s", err)
+	}
+
+	fmt.Fprintf(os.Stderr, "Successfully deleted %d bouncers\n", deleted)
+
+	return nil
+}
+
+func (cli *cliBouncers) newPruneCmd() *cobra.Command {
+	var (
+		duration time.Duration
+		force    bool
+	)
+
+	const defaultDuration = 60 * time.Minute
+
 	cmd := &cobra.Command{
 		Use:               "prune",
 		Short:             "prune multiple bouncers from the database",
 		Args:              cobra.NoArgs,
 		DisableAutoGenTag: true,
-		Example: `cscli bouncers prune -d 60m
-cscli bouncers prune -d 60m --force`,
-		PreRunE: func(cmd *cobra.Command, args []string) error {
-			dur, _ := cmd.Flags().GetString("duration")
-			var err error
-			parsedDuration, err = time.ParseDuration(fmt.Sprintf("-%s", dur))
-			if err != nil {
-				return fmt.Errorf("unable to parse duration '%s': %s", dur, err)
-			}
-			return nil
-		},
-		RunE: func(cmd *cobra.Command, args []string) error {
-			force, _ := cmd.Flags().GetBool("force")
-			if parsedDuration >= 0-2*time.Minute {
-				var answer bool
-				prompt := &survey.Confirm{
-					Message: "The duration you provided is less than or equal 2 minutes this may remove active bouncers continue ?",
-					Default: false,
-				}
-				if err := survey.AskOne(prompt, &answer); err != nil {
-					return fmt.Errorf("unable to ask about prune check: %s", err)
-				}
-				if !answer {
-					fmt.Println("user aborted prune no changes were made")
-					return nil
-				}
-			}
-			bouncers, err := dbClient.QueryBouncersLastPulltimeLT(time.Now().UTC().Add(parsedDuration))
-			if err != nil {
-				return fmt.Errorf("unable to query bouncers: %s", err)
-			}
-			if len(bouncers) == 0 {
-				fmt.Println("no bouncers to prune")
-				return nil
-			}
-			getBouncersTable(color.Output, bouncers)
-			if !force {
-				var answer bool
-				prompt := &survey.Confirm{
-					Message: "You are about to PERMANENTLY remove the above bouncers from the database these will NOT be recoverable, continue ?",
-					Default: false,
-				}
-				if err := survey.AskOne(prompt, &answer); err != nil {
-					return fmt.Errorf("unable to ask about prune check: %s", err)
-				}
-				if !answer {
-					fmt.Println("user aborted prune no changes were made")
-					return nil
-				}
-			}
-			nbDeleted, err := dbClient.BulkDeleteBouncers(bouncers)
-			if err != nil {
-				return fmt.Errorf("unable to prune bouncers: %s", err)
-			}
-			fmt.Printf("successfully delete %d bouncers\n", nbDeleted)
-			return nil
+		Example: `cscli bouncers prune -d 45m
+cscli bouncers prune -d 45m --force`,
+		RunE: func(_ *cobra.Command, _ []string) error {
+			return cli.prune(duration, force)
 		},
 	}
-	cmd.Flags().StringP("duration", "d", "60m", "duration of time since last pull")
-	cmd.Flags().Bool("force", false, "force prune without asking for confirmation")
+
+	flags := cmd.Flags()
+	flags.DurationVarP(&duration, "duration", "d", defaultDuration, "duration of time since last pull")
+	flags.BoolVar(&force, "force", false, "force prune without asking for confirmation")
+
 	return cmd
 }

+ 6 - 1
cmd/crowdsec-cli/main.go

@@ -182,6 +182,11 @@ It is meant to allow you to manage bans, parsers/scenarios/etc, api and generall
 	cmd.Flags().SortFlags = false
 	cmd.PersistentFlags().SortFlags = false
 
+	// we use a getter because the config is not initialized until the Execute() call
+	getconfig := func() *csconfig.Config {
+		return csConfig
+	}
+
 	cmd.AddCommand(NewCLIDoc().NewCommand(cmd))
 	cmd.AddCommand(NewCLIVersion().NewCommand())
 	cmd.AddCommand(NewConfigCmd())
@@ -191,7 +196,7 @@ It is meant to allow you to manage bans, parsers/scenarios/etc, api and generall
 	cmd.AddCommand(NewCLIDecisions().NewCommand())
 	cmd.AddCommand(NewCLIAlerts().NewCommand())
 	cmd.AddCommand(NewCLISimulation().NewCommand())
-	cmd.AddCommand(NewCLIBouncers().NewCommand())
+	cmd.AddCommand(NewCLIBouncers(getconfig).NewCommand())
 	cmd.AddCommand(NewCLIMachines().NewCommand())
 	cmd.AddCommand(NewCLICapi().NewCommand())
 	cmd.AddCommand(NewLapiCmd())

+ 3 - 2
cmd/crowdsec-cli/support.go

@@ -149,10 +149,11 @@ func collectHubItems(hub *cwhub.Hub, itemType string) []byte {
 
 func collectBouncers(dbClient *database.Client) ([]byte, error) {
 	out := bytes.NewBuffer(nil)
-	err := getBouncers(out, dbClient)
+	bouncers, err := dbClient.ListBouncers()
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("unable to list bouncers: %s", err)
 	}
+	getBouncersTable(out, bouncers)
 	return out.Bytes(), nil
 }
 

+ 0 - 12
cmd/crowdsec-cli/utils.go

@@ -8,7 +8,6 @@ import (
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 
-	"github.com/crowdsecurity/crowdsec/pkg/database"
 	"github.com/crowdsecurity/crowdsec/pkg/types"
 )
 
@@ -47,17 +46,6 @@ func manageCliDecisionAlerts(ip *string, ipRange *string, scope *string, value *
 	return nil
 }
 
-func getDBClient() (*database.Client, error) {
-	if err := csConfig.LoadAPIServer(true); err != nil || csConfig.DisableAPI {
-		return nil, err
-	}
-	ret, err := database.NewClient(csConfig.DbConfig)
-	if err != nil {
-		return nil, err
-	}
-	return ret, nil
-}
-
 func removeFromSlice(val string, slice []string) []string {
 	var i int
 	var value string

+ 2 - 2
pkg/csconfig/api.go

@@ -285,7 +285,7 @@ func (c *Config) LoadAPIServer(inCli bool) error {
 		}
 	}
 
-	if c.API.Server.OnlineClient == nil || c.API.Server.OnlineClient.Credentials == nil {
+	if (c.API.Server.OnlineClient == nil || c.API.Server.OnlineClient.Credentials == nil) && !inCli {
 		log.Printf("push and pull to Central API disabled")
 	}
 
@@ -297,7 +297,7 @@ func (c *Config) LoadAPIServer(inCli bool) error {
 		return err
 	}
 
-	if c.API.Server.CapiWhitelistsPath != "" {
+	if c.API.Server.CapiWhitelistsPath != "" && !inCli {
 		log.Infof("loaded capi whitelist from %s: %d IPs, %d CIDRs", c.API.Server.CapiWhitelistsPath, len(c.API.Server.CapiWhitelists.Ips), len(c.API.Server.CapiWhitelists.Cidrs))
 	}
 

+ 2 - 2
pkg/csconfig/config.go

@@ -41,9 +41,9 @@ type Config struct {
 	Hub          *LocalHubCfg        `yaml:"-"`
 }
 
-func NewConfig(configFile string, disableAgent bool, disableAPI bool, quiet bool) (*Config, string, error) {
+func NewConfig(configFile string, disableAgent bool, disableAPI bool, inCli bool) (*Config, string, error) {
 	patcher := yamlpatch.NewPatcher(configFile, ".local")
-	patcher.SetQuiet(quiet)
+	patcher.SetQuiet(inCli)
 	fcontent, err := patcher.MergedPatchContent()
 	if err != nil {
 		return nil, "", err