forked from sandia-minimega/discovery
/
os.go
134 lines (103 loc) · 2.5 KB
/
os.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
// Copyright 2018 National Technology & Engineering Solutions of Sandia, LLC
// (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the U.S.
// Government retains certain rights in this software.
package main
import (
"flag"
"fmt"
"math/rand"
"strconv"
"strings"
"pkg/commands"
log "pkg/minilog"
)
type OS struct {
Name string
// Prob is percentage (0 <= Prob <= 1)
Prob float64
}
type OSes struct {
// list of OS
oses []OS
// computed ranges for RNG:
// if ranges[i-1] < rand.Float64() < ranges[i] => pick OS i
ranges []float64
}
type CommandOS struct {
commands.Base // embed
}
func init() {
base := commands.Base{
Flags: flag.NewFlagSet("os", flag.ExitOnError),
Usage: "os <OS,WEIGHT>...",
Short: "annotate with random oses",
Long: `
Annotate nodes with OSes based on the distribution specified on the command
line. For example, to get 50% Windows and 50% Linux:
os Linux:0.5 Windows:0.5
`,
}
commands.Append(&CommandOS{base})
}
func (o OSes) Rand(r *rand.Rand) string {
v := r.Float64()
for i, v2 := range o.ranges {
if v < v2 {
return o.oses[i].Name
}
}
log.Error("no OSes?")
return "???"
}
func (c *CommandOS) Run() error {
oses, err := parseOSes(c.Flags.Args())
if err != nil {
return err
}
endpoints, err := dc.GetEndpoints("", "")
if err != nil {
return err
}
for _, v := range endpoints {
if _, ok := v.D["os"]; ok && !*f_overwrite {
continue
}
v.D["os"] = oses.Rand(rng)
log.Debug("updating node %v with os %v", v, v.D["os"])
if err := UpdateEndpoint(v); err != nil {
return err
}
}
return nil
}
// parseOSes parses OS,WEIGHT args and returns a map with normalized values
func parseOSes(args []string) (*OSes, error) {
oses := []OS{}
for _, f := range args {
parts := strings.SplitN(f, ",", 2)
if len(parts) != 2 {
return nil, fmt.Errorf("expected OS,WEIGHT not `%v`", f)
}
v, err := strconv.ParseFloat(parts[1], 64)
if err != nil {
return nil, fmt.Errorf("expected float not `%v`", parts[1])
}
name := parts[0]
log.Info("adding OS %v with weight %v", name, v)
oses = append(oses, OS{Name: name, Prob: v})
}
// normalize weights to create probabilities
sum := 0.0
for _, v := range oses {
sum += v.Prob
}
prev := 0.0
ranges := []float64{}
for i := range oses {
oses[i].Prob = oses[i].Prob / sum
log.Debug("converted OS %v weight to probability %v%%", oses[i].Name, oses[i].Prob*100)
prev = prev + oses[i].Prob
ranges = append(ranges, prev)
}
return &OSes{oses: oses, ranges: ranges}, nil
}