1
package pflag
2

3
import (
4
	"fmt"
5
	"io"
6
	"net"
7
	"strings"
8
)
9

10
// -- ipNetSlice Value
11
type ipNetSliceValue struct {
12
	value   *[]net.IPNet
13
	changed bool
14
}
15

16
func newIPNetSliceValue(val []net.IPNet, p *[]net.IPNet) *ipNetSliceValue {
17
	ipnsv := new(ipNetSliceValue)
18
	ipnsv.value = p
19
	*ipnsv.value = val
20
	return ipnsv
21
}
22

23
// Set converts, and assigns, the comma-separated IPNet argument string representation as the []net.IPNet value of this flag.
24
// If Set is called on a flag that already has a []net.IPNet assigned, the newly converted values will be appended.
25
func (s *ipNetSliceValue) Set(val string) error {
26

27
	// remove all quote characters
28
	rmQuote := strings.NewReplacer(`"`, "", `'`, "", "`", "")
29

30
	// read flag arguments with CSV parser
31
	ipNetStrSlice, err := readAsCSV(rmQuote.Replace(val))
32
	if err != nil && err != io.EOF {
33
		return err
34
	}
35

36
	// parse ip values into slice
37
	out := make([]net.IPNet, 0, len(ipNetStrSlice))
38
	for _, ipNetStr := range ipNetStrSlice {
39
		_, n, err := net.ParseCIDR(strings.TrimSpace(ipNetStr))
40
		if err != nil {
41
			return fmt.Errorf("invalid string being converted to CIDR: %s", ipNetStr)
42
		}
43
		out = append(out, *n)
44
	}
45

46
	if !s.changed {
47
		*s.value = out
48
	} else {
49
		*s.value = append(*s.value, out...)
50
	}
51

52
	s.changed = true
53

54
	return nil
55
}
56

57
// Type returns a string that uniquely represents this flag's type.
58
func (s *ipNetSliceValue) Type() string {
59
	return "ipNetSlice"
60
}
61

62
// String defines a "native" format for this net.IPNet slice flag value.
63
func (s *ipNetSliceValue) String() string {
64

65
	ipNetStrSlice := make([]string, len(*s.value))
66
	for i, n := range *s.value {
67
		ipNetStrSlice[i] = n.String()
68
	}
69

70
	out, _ := writeAsCSV(ipNetStrSlice)
71
	return "[" + out + "]"
72
}
73

74
func ipNetSliceConv(val string) (interface{}, error) {
75
	val = strings.Trim(val, "[]")
76
	// Empty string would cause a slice with one (empty) entry
77
	if len(val) == 0 {
78
		return []net.IPNet{}, nil
79
	}
80
	ss := strings.Split(val, ",")
81
	out := make([]net.IPNet, len(ss))
82
	for i, sval := range ss {
83
		_, n, err := net.ParseCIDR(strings.TrimSpace(sval))
84
		if err != nil {
85
			return nil, fmt.Errorf("invalid string being converted to CIDR: %s", sval)
86
		}
87
		out[i] = *n
88
	}
89
	return out, nil
90
}
91

92
// GetIPNetSlice returns the []net.IPNet value of a flag with the given name
93
func (f *FlagSet) GetIPNetSlice(name string) ([]net.IPNet, error) {
94
	val, err := f.getFlagType(name, "ipNetSlice", ipNetSliceConv)
95
	if err != nil {
96
		return []net.IPNet{}, err
97
	}
98
	return val.([]net.IPNet), nil
99
}
100

101
// IPNetSliceVar defines a ipNetSlice flag with specified name, default value, and usage string.
102
// The argument p points to a []net.IPNet variable in which to store the value of the flag.
103
func (f *FlagSet) IPNetSliceVar(p *[]net.IPNet, name string, value []net.IPNet, usage string) {
104
	f.VarP(newIPNetSliceValue(value, p), name, "", usage)
105
}
106

107
// IPNetSliceVarP is like IPNetSliceVar, but accepts a shorthand letter that can be used after a single dash.
108
func (f *FlagSet) IPNetSliceVarP(p *[]net.IPNet, name, shorthand string, value []net.IPNet, usage string) {
109
	f.VarP(newIPNetSliceValue(value, p), name, shorthand, usage)
110
}
111

112
// IPNetSliceVar defines a []net.IPNet flag with specified name, default value, and usage string.
113
// The argument p points to a []net.IPNet variable in which to store the value of the flag.
114
func IPNetSliceVar(p *[]net.IPNet, name string, value []net.IPNet, usage string) {
115
	CommandLine.VarP(newIPNetSliceValue(value, p), name, "", usage)
116
}
117

118
// IPNetSliceVarP is like IPNetSliceVar, but accepts a shorthand letter that can be used after a single dash.
119
func IPNetSliceVarP(p *[]net.IPNet, name, shorthand string, value []net.IPNet, usage string) {
120
	CommandLine.VarP(newIPNetSliceValue(value, p), name, shorthand, usage)
121
}
122

123
// IPNetSlice defines a []net.IPNet flag with specified name, default value, and usage string.
124
// The return value is the address of a []net.IPNet variable that stores the value of that flag.
125
func (f *FlagSet) IPNetSlice(name string, value []net.IPNet, usage string) *[]net.IPNet {
126
	p := []net.IPNet{}
127
	f.IPNetSliceVarP(&p, name, "", value, usage)
128
	return &p
129
}
130

131
// IPNetSliceP is like IPNetSlice, but accepts a shorthand letter that can be used after a single dash.
132
func (f *FlagSet) IPNetSliceP(name, shorthand string, value []net.IPNet, usage string) *[]net.IPNet {
133
	p := []net.IPNet{}
134
	f.IPNetSliceVarP(&p, name, shorthand, value, usage)
135
	return &p
136
}
137

138
// IPNetSlice defines a []net.IPNet flag with specified name, default value, and usage string.
139
// The return value is the address of a []net.IP variable that stores the value of the flag.
140
func IPNetSlice(name string, value []net.IPNet, usage string) *[]net.IPNet {
141
	return CommandLine.IPNetSliceP(name, "", value, usage)
142
}
143

144
// IPNetSliceP is like IPNetSlice, but accepts a shorthand letter that can be used after a single dash.
145
func IPNetSliceP(name, shorthand string, value []net.IPNet, usage string) *[]net.IPNet {
146
	return CommandLine.IPNetSliceP(name, shorthand, value, usage)
147
}
148

149