"vault operator usage" CLI for client count reporting

command/commands.go

@@ -407,6 +407,11 @@ func initCommands(ui, serverCmdUi cli.Ui, runOpts *RunOptions) {
 				BaseCommand: getBaseCommand(),
 			}, nil
 		},
+		"operator usage": func() (cli.Command, error) {
+			return &OperatorUsageCommand{
+				BaseCommand: getBaseCommand(),
+			}, nil
+		},
 		"operator unseal": func() (cli.Command, error) {
 			return &OperatorUnsealCommand{
 				BaseCommand: getBaseCommand(),

command/operator_usage.go

@@ -0,0 +1,286 @@
+package command
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"sort"
+	"strings"
+	"time"
+
+	"github.com/mitchellh/cli"
+	"github.com/posener/complete"
+	"github.com/ryanuber/columnize"
+)
+
+var _ cli.Command = (*OperatorUsageCommand)(nil)
+var _ cli.CommandAutocomplete = (*OperatorUsageCommand)(nil)
+
+type OperatorUsageCommand struct {
+	*BaseCommand
+	flagStartTime time.Time
+	flagEndTime   time.Time
+}
+
+func (c *OperatorUsageCommand) Synopsis() string {
+	return "Lists historical client counts"
+}
+
+func (c *OperatorUsageCommand) Help() string {
+	helpText := `
+Usage: vault operator usage
+
+  List the client counts for the default reporting period.
+
+	  $ vault operator usage
+
+  List the client counts for a specific time period.
+
+          $ vault operator usage -start-time=2020-10 -end-time=2020-11
+
+` + c.Flags().Help()
+
+	return strings.TrimSpace(helpText)
+}
+
+func (c *OperatorUsageCommand) Flags() *FlagSets {
+	set := c.flagSet(FlagSetHTTP | FlagSetOutputFormat)
+
+	f := set.NewFlagSet("Command Options")
+
+	f.TimeVar(&TimeVar{
+		Name:       "start-time",
+		Usage:      "Start of report period (defaults to default_reporting_period before end time.)",
+		Target:     &c.flagStartTime,
+		Completion: complete.PredictNothing,
+		Default:    time.Time{},
+		Formats:    TimeVar_TimeOrDay | TimeVar_Month,
+	})
+	f.TimeVar(&TimeVar{
+		Name:       "end-time",
+		Usage:      "End of report period (defaults to end of last month.)",
+		Target:     &c.flagEndTime,
+		Completion: complete.PredictNothing,
+		Default:    time.Time{},
+		Formats:    TimeVar_TimeOrDay | TimeVar_Month,
+	})
+
+	return set
+}
+
+func (c *OperatorUsageCommand) AutocompleteArgs() complete.Predictor {
+	return complete.PredictAnything
+}
+
+func (c *OperatorUsageCommand) AutocompleteFlags() complete.Flags {
+	return c.Flags().Completions()
+}
+
+func (c *OperatorUsageCommand) Run(args []string) int {
+	f := c.Flags()
+
+	if err := f.Parse(args); err != nil {
+		c.UI.Error(err.Error())
+		return 1
+	}
+
+	data := make(map[string][]string)
+	if !c.flagStartTime.IsZero() {
+		data["start_time"] = []string{c.flagStartTime.Format(time.RFC3339)}
+	}
+	if !c.flagEndTime.IsZero() {
+		data["end_time"] = []string{c.flagEndTime.Format(time.RFC3339)}
+	}
+
+	client, err := c.Client()
+	if err != nil {
+		c.UI.Error(err.Error())
+		return 2
+	}
+
+	resp, err := client.Logical().ReadWithData("sys/internal/counters/activity", data)
+	if err != nil {
+		c.UI.Error(fmt.Sprintf("Error retrieving client counts: %v", err))
+		return 2
+	}
+
+	if resp == nil || resp.Data == nil {
+		c.UI.Warn("No data is available for the given time range.")
+		// No further output
+		// TODO: report if any data at all is available?
+		return 0
+	}
+
+	switch Format(c.UI) {
+	case "table":
+	default:
+		// Handle JSON, YAML, etc.
+		return OutputData(c.UI, resp)
+	}
+
+	// Show this before the headers
+	c.outputTimestamps(resp.Data)
+
+	out := []string{
+		"Namespace path | Distinct entities | Non-Entity tokens | Active clients",
+	}
+
+	out = append(out, c.namespacesOutput(resp.Data)...)
+
+	out = c.addTotalToOutput(out, resp.Data)
+
+	colConfig := columnize.DefaultConfig()
+	colConfig.Empty = " " // Do not show n/a on intentional blank lines
+	colConfig.Glue = "   "
+	c.UI.Output(tableOutput(out, colConfig))
+	return 0
+}
+
+func (c *OperatorUsageCommand) outputTimestamps(data map[string]interface{}) {
+	c.UI.Output(fmt.Sprintf("Period start: %v\nPeriod end: %v\n",
+		data["start_time"].(string),
+		data["end_time"].(string)))
+}
+
+type UsageCommandNamespace struct {
+	formattedLine string
+	sortOrder     string
+
+	// Sort order:
+	// -- root first
+	// -- namespaces in lexicographic order
+	// -- deleted namespace "xxxxx" last
+}
+
+type UsageResponse struct {
+	namespacePath string
+	entityCount   int64
+	tokenCount    int64
+	clientCount   int64
+}
+
+func jsonNumberOK(m map[string]interface{}, key string) (int64, bool) {
+	val, ok := m[key].(json.Number)
+	if !ok {
+		return 0, false
+	}
+	intVal, err := val.Int64()
+	if err != nil {
+		return 0, false
+	}
+	return intVal, true
+}
+
+// TODO: provide a function in the API module for doing this conversion?
+func (c *OperatorUsageCommand) parseNamespaceCount(rawVal interface{}) (UsageResponse, error) {
+	var ret UsageResponse
+
+	val, ok := rawVal.(map[string]interface{})
+	if !ok {
+		return ret, errors.New("value is not a map")
+	}
+
+	ret.namespacePath, ok = val["namespace_path"].(string)
+	if !ok {
+		return ret, errors.New("bad namespace path")
+	}
+
+	counts, ok := val["counts"].(map[string]interface{})
+	if !ok {
+		return ret, errors.New("missing counts")
+	}
+
+	ret.entityCount, ok = jsonNumberOK(counts, "distinct_entities")
+	if !ok {
+		return ret, errors.New("missing distinct_entities")
+	}
+
+	ret.tokenCount, ok = jsonNumberOK(counts, "non_entity_tokens")
+	if !ok {
+		return ret, errors.New("missing non_entity_tokens")
+	}
+
+	ret.clientCount, ok = jsonNumberOK(counts, "clients")
+	if !ok {
+		return ret, errors.New("missing clients")
+	}
+
+	return ret, nil
+
+}
+
+func (c *OperatorUsageCommand) namespacesOutput(data map[string]interface{}) []string {
+	byNs, ok := data["by_namespace"].([]interface{})
+	if !ok {
+		c.UI.Error("missing namespace breakdown in response")
+		return nil
+	}
+
+	nsOut := make([]UsageCommandNamespace, 0, len(byNs))
+
+	for _, rawVal := range byNs {
+		val, err := c.parseNamespaceCount(rawVal)
+		if err != nil {
+			c.UI.Error(fmt.Sprintf("malformed namespace in response: %v", err))
+			continue
+		}
+
+		sortOrder := "1" + val.namespacePath
+		if val.namespacePath == "" {
+			val.namespacePath = "[root]"
+			sortOrder = "0"
+		} else if strings.HasPrefix(val.namespacePath, "deleted namespace") {
+			sortOrder = "2" + val.namespacePath
+		}
+
+		formattedLine := fmt.Sprintf("%s | %d | %d | %d",
+			val.namespacePath, val.entityCount, val.tokenCount, val.clientCount)
+		nsOut = append(nsOut, UsageCommandNamespace{
+			formattedLine: formattedLine,
+			sortOrder:     sortOrder,
+		})
+	}
+
+	sort.Slice(nsOut, func(i, j int) bool {
+		return nsOut[i].sortOrder < nsOut[j].sortOrder
+	})
+
+	out := make([]string, len(nsOut))
+	for i := range nsOut {
+		out[i] = nsOut[i].formattedLine
+	}
+
+	return out
+}
+
+func (c *OperatorUsageCommand) addTotalToOutput(out []string, data map[string]interface{}) []string {
+	// blank line separating it from namespaces
+	out = append(out, "  |  |  |  ")
+
+	total, ok := data["total"].(map[string]interface{})
+	if !ok {
+		c.UI.Error("missing total in response")
+		return out
+	}
+
+	entityCount, ok := jsonNumberOK(total, "distinct_entities")
+	if !ok {
+		c.UI.Error("missing distinct_entities in total")
+		return out
+	}
+
+	tokenCount, ok := jsonNumberOK(total, "non_entity_tokens")
+	if !ok {
+		c.UI.Error("missing non_entity_tokens in total")
+		return out
+	}
+	clientCount, ok := jsonNumberOK(total, "clients")
+	if !ok {
+		c.UI.Error("missing clients in total")
+		return out
+	}
+
+	out = append(out, fmt.Sprintf("Total | %d | %d | %d",
+		entityCount, tokenCount, clientCount))
+	return out
+}