Add 'RandIpAddress'.

This commit is contained in:
Kit Ewbank 2019-08-02 13:35:28 -04:00 committed by Kit Ewbank
parent 3cda2bd7a5
commit 949a6ae695
2 changed files with 94 additions and 0 deletions

View File

@ -10,10 +10,13 @@ import (
"fmt"
"math/big"
"math/rand"
"net"
"strings"
"time"
"golang.org/x/crypto/ssh"
"github.com/apparentlymart/go-cidr/cidr"
)
func init() {
@ -105,6 +108,39 @@ func RandTLSCert(orgName string) (string, string, error) {
return certPEM, privateKeyPEM, nil
}
// RandIpAddress returns a random IP address in the specified CIDR block.
// The prefix length must be less than 31.
func RandIpAddress(s string) (string, error) {
_, network, err := net.ParseCIDR(s)
if err != nil {
return "", err
}
firstIp, lastIp := cidr.AddressRange(network)
first := &big.Int{}
first.SetBytes([]byte(firstIp))
last := &big.Int{}
last.SetBytes([]byte(lastIp))
r := &big.Int{}
r.Sub(last, first)
if len := r.BitLen(); len > 31 {
return "", fmt.Errorf("CIDR range is too large: %d", len)
}
max := int(r.Int64())
if max == 0 {
// panic: invalid argument to Int31n
return firstIp.String(), nil
}
host, err := cidr.Host(network, RandIntRange(0, max))
if err != nil {
return "", err
}
return host.String(), nil
}
func genPrivateKey() (*rsa.PrivateKey, string, error) {
privateKey, err := rsa.GenerateKey(crand.Reader, 1024)
if err != nil {

View File

@ -0,0 +1,58 @@
package acctest
import (
"regexp"
"testing"
)
func TestRandIpAddress(t *testing.T) {
testCases := []struct {
s string
expected *regexp.Regexp
expectedErr string
}{
{
s: "1.1.1.1/32",
expected: regexp.MustCompile(`^1\.1\.1\.1$`),
},
{
s: "10.0.0.0/8",
expected: regexp.MustCompile(`^10\.\d{1,3}\.\d{1,3}\.\d{1,3}$`),
},
{
s: "0.0.0.0/0",
expectedErr: "CIDR range is too large: 32",
},
{
s: "449d:e5f1:14b1:ddf3:8525:7e9e:4a0d:4a82/128",
expected: regexp.MustCompile(`^449d:e5f1:14b1:ddf3:8525:7e9e:4a0d:4a82$`),
},
{
s: "2001:db8::/112",
expected: regexp.MustCompile(`^2001:db8::[[:xdigit:]]{1,4}$`),
},
{
s: "2001:db8::/64",
expectedErr: "CIDR range is too large: 64",
},
{
s: "abcdefg",
expectedErr: "invalid CIDR address: abcdefg",
},
}
for i, tc := range testCases {
v, err := RandIpAddress(tc.s)
if err != nil {
msg := err.Error()
if tc.expectedErr == "" {
t.Fatalf("expected test case %d to succeed but got error %q, ", i, msg)
}
if msg != tc.expectedErr {
t.Fatalf("expected test case %d to fail with %q but got %q", i, tc.expectedErr, msg)
}
} else if !tc.expected.MatchString(v) {
t.Fatalf("expected test case %d to return %q but got %q", i, tc.expected, v)
}
}
}