/
kv.go
148 lines (116 loc) · 3.72 KB
/
kv.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
package main
import (
"encoding/json"
"flag"
"fmt"
"io/ioutil"
"math/rand"
"time"
"github.com/mitchellh/cli"
"github.com/mkeeler/consul-data-gen/generate/kv"
)
const (
keyTypePetName = "pet-name"
valueTypeRandomB64 = "b64-random"
)
type kvCommand struct {
ui cli.Ui
config kv.Config
keyType stringChoiceValue
valueType stringChoiceValue
numEntries int
randValueMinSize int
randValueMaxSize int
petNameLength int
petNamePrefix string
petNameSeparator string
randSeed int64
flags *flag.FlagSet
help string
}
func newKVCommand(ui cli.Ui) cli.Command {
c := &kvCommand{
config: kv.DefaultConfig(),
ui: ui,
keyType: stringChoiceValue{
choices: []string{keyTypePetName},
value: keyTypePetName,
},
valueType: stringChoiceValue{
choices: []string{valueTypeRandomB64},
value: valueTypeRandomB64,
},
}
flags := flag.NewFlagSet("", flag.ContinueOnError)
flags.Int64Var(&c.randSeed, "seed", 0, "Value to use to seed the pseudo-random number generator with instead of the current time")
flags.IntVar(&c.config.NumEntries, "num-entries", 1024, "Number of KV entries to generate")
flags.IntVar(&c.randValueMinSize, "rand-value-min", 64, "Minimum byte size of random KV values to create")
flags.IntVar(&c.randValueMaxSize, "rand-value-max", 1024, "Maximum byte size of random KV values to create")
flags.IntVar(&c.petNameLength, "pet-name-len", 3, "Length in words of the randomly generated KV key pet names")
flags.StringVar(&c.petNamePrefix, "pet-name-prefix", "", "Prefix for KV key pet names")
flags.StringVar(&c.petNameSeparator, "pet-name-separator", "-", "Separator char for KV key pet names")
c.flags = flags
c.help = genUsage(`Usage: consul-data-gen kv [OPTIONS] [output path]
Generate KV data for consul.
By default the generated output is sent to the console but
an optional output path may be used to cause it to be written
to a file`, c.flags)
return c
}
func (c *kvCommand) Run(args []string) int {
if err := c.flags.Parse(args); err != nil {
c.ui.Error(fmt.Sprintf("Failed to parse command line arguments: %v", err))
return 1
}
args = c.flags.Args()
if c.randSeed == 0 {
c.randSeed = time.Now().UnixNano()
}
rand.Seed(c.randSeed)
switch c.keyType.value {
case keyTypePetName:
if c.petNameLength < 1 {
c.ui.Error(fmt.Sprintf("Invalid pet name length: %v", c.petNameLength))
return 1
}
c.config.KeyGen = kv.PetNameKeyGenerator(c.petNamePrefix, c.petNameLength, c.petNameSeparator)
}
switch c.valueType.value {
case valueTypeRandomB64:
if c.randValueMinSize < 1 {
c.ui.Error(fmt.Sprintf("Invalid random value minimum size: %v. Value must be a positive non-zero integer", c.randValueMinSize))
return 1
}
if c.randValueMaxSize < c.randValueMinSize {
c.ui.Error(fmt.Sprintf("Invalid random value max size: %v. Value must be greater than or equal to the minimum size of %v", c.randValueMaxSize, c.randValueMinSize))
return 1
}
c.config.ValueGen = kv.RandomValueGenerator(c.randValueMinSize, c.randValueMaxSize)
}
data, err := kv.Generate(c.config)
if err != nil {
c.ui.Error(fmt.Sprintf("Failed to generate KV data: %v\n", err))
return 1
}
serialized, err := json.MarshalIndent(data, "", " ")
if err != nil {
c.ui.Error(fmt.Sprintf("Failed to serialize KV data: %v\n", err))
return 1
}
if len(args) > 0 {
if err := ioutil.WriteFile(args[0], serialized, 0644); err != nil {
c.ui.Error(fmt.Sprintf("Failed to write serialized KV data to %q: %v", args[0], err))
return 1
}
c.ui.Info(fmt.Sprintf("KV data written to %s", args[0]))
} else {
c.ui.Info(string(serialized))
}
return 0
}
func (c *kvCommand) Synopsis() string {
return "Generate KV data for Consul"
}
func (c *kvCommand) Help() string {
return c.help
}