network dns support multiple domain

Signed-off-by: sanxun0325 <bbz17640380550@163.com>

Author: binbin0325

exec/bin/changedns/changedns.go

@@ -20,13 +20,15 @@ import (
 	"context"
 	"flag"
 	"fmt"
-
 	"github.com/chaosblade-io/chaosblade-spec-go/channel"
 	"github.com/chaosblade-io/chaosblade-spec-go/spec"
+	"strings"
 
 	"github.com/chaosblade-io/chaosblade-exec-os/exec/bin"
 )
 
+const sep = ","
+
 var dnsDomain, dnsIp string
 var changeDnsStart, changeDnsStop bool
 
@@ -57,13 +59,18 @@ var cl = channel.NewLocalChannel()
 // startChangeDns by the domain and ip
 func startChangeDns(domain, ip string) {
 	ctx := context.Background()
-	dnsPair := createDnsPair(domain, ip)
-	response := cl.Run(ctx, "grep", fmt.Sprintf(`-q "%s" %s`, dnsPair, hosts))
-	if response.Success {
-		bin.PrintErrAndExit(fmt.Sprintf("%s has been exist", dnsPair))
-		return
+	domains := strings.Split(domain, sep)
+	var dnsPairs string
+	for _, v := range domains {
+		dnsPair := createDnsPair(v, ip)
+		response := cl.Run(ctx, "grep", fmt.Sprintf(`-q "%s" %s`, dnsPair, hosts))
+		if response.Success {
+			bin.PrintErrAndExit(fmt.Sprintf("%s has been exist", dnsPair))
+			return
+		}
+		dnsPairs = dnsPairs + "\n" + dnsPair
 	}
-	response = cl.Run(ctx, "echo", fmt.Sprintf(`"%s" >> %s`, dnsPair, hosts))
+	response := cl.Run(ctx, "echo", fmt.Sprintf(`"%s" >> %s`, dnsPairs, hosts))
 	if !response.Success {
 		bin.PrintErrAndExit(response.Err)
 		return
@@ -74,22 +81,24 @@ func startChangeDns(domain, ip string) {
 // recoverDns
 func recoverDns(domain, ip string) {
 	ctx := context.Background()
-	dnsPair := createDnsPair(domain, ip)
-	response := cl.Run(ctx, "grep", fmt.Sprintf(`-q "%s" %s`, dnsPair, hosts))
-	if !response.Success {
-		bin.PrintOutputAndExit("nothing to do")
-		return
-	}
+	domains := strings.Split(domain, sep)
+	for _, v := range domains {
+		dnsPair := createDnsPair(v, ip)
+		response := cl.Run(ctx, "grep", fmt.Sprintf(`-q "%s" %s`, dnsPair, hosts))
+		if !response.Success {
+			bin.PrintOutputAndExit("nothing to do")
+			return
+		}
 
-	if !cl.IsCommandAvailable("cat") {
-		bin.PrintErrAndExit(spec.CommandCatNotFound.Msg)
-	}
-
-	response = cl.Run(ctx, "cat", fmt.Sprintf(`%s | grep -v "%s" > %s && cat %s > %s`,
-		hosts, dnsPair, tmpHosts, tmpHosts, hosts))
-	if !response.Success {
-		bin.PrintErrAndExit(response.Err)
-		return
+		if !cl.IsCommandAvailable("cat") {
+			bin.PrintErrAndExit(spec.CommandCatNotFound.Msg)
+		}
+		response = cl.Run(ctx, "cat", fmt.Sprintf(`%s | grep -v "%s" > %s && cat %s > %s`,
+			hosts, dnsPair, tmpHosts, tmpHosts, hosts))
+		if !response.Success {
+			bin.PrintErrAndExit(response.Err)
+			return
+		}
 	}
 	cl.Run(ctx, "rm", fmt.Sprintf(`-rf %s`, tmpHosts))
 }

exec/bin/changedns/changedns_test.go

@@ -113,3 +113,83 @@ func Test_recoverDns_failed(t *testing.T) {
 		t.Errorf("unexpected commands: %+v, expected commands: %+v", actualCommands, expectedCommands)
 	}
 }
+
+func Test_startChangeMultipleDns(t *testing.T) {
+	type args struct {
+		domain string
+		ip     string
+	}
+
+	as := &args{
+		domain: "abc.com,bbc.com",
+		ip:     "208.80.152.2",
+	}
+
+	var exitCode int
+	bin.ExitFunc = func(code int) {
+		exitCode = code
+	}
+	cl = channel.NewMockLocalChannel()
+	mockChannel := cl.(*channel.MockLocalChannel)
+	actualCommands := make([]string, 0)
+	mockChannel.RunFunc = func(ctx context.Context, script, args string) *spec.Response {
+		actualCommands = append(actualCommands, fmt.Sprintf("%s %s", script, args))
+		if script == "echo" {
+			return spec.ReturnSuccess("")
+		}
+		return spec.ReturnFail(spec.CodeType{}, "")
+	}
+	expectedCommands := []string{`grep -q "208.80.152.2 abc.com #chaosblade" /etc/hosts`,
+		`grep -q "208.80.152.2 bbc.com #chaosblade" /etc/hosts`,
+		`echo "
+208.80.152.2 abc.com #chaosblade
+208.80.152.2 bbc.com #chaosblade" >> /etc/hosts`}
+
+	startChangeDns(as.domain, as.ip)
+	if exitCode != 0 {
+		t.Errorf("unexpected result: %d, expected result: %d", exitCode, 0)
+	}
+	if !reflect.DeepEqual(expectedCommands, actualCommands) {
+		t.Errorf("unexpected commands: %+v, expected commands: %+v", actualCommands, expectedCommands)
+	}
+}
+
+func Test_recoverChangeMultipleDns(t *testing.T) {
+	type args struct {
+		domain string
+		ip     string
+	}
+
+	as := &args{
+		domain: "abc.com,bbc.com",
+		ip:     "208.80.152.2",
+	}
+
+	var exitCode int
+	bin.ExitFunc = func(code int) {
+		exitCode = code
+	}
+	cl = channel.NewMockLocalChannel()
+	mockChannel := cl.(*channel.MockLocalChannel)
+	actualCommands := make([]string, 0)
+	mockChannel.RunFunc = func(ctx context.Context, script, args string) *spec.Response {
+		actualCommands = append(actualCommands, fmt.Sprintf("%s %s", script, args))
+		return spec.ReturnSuccess("")
+	}
+	mockChannel.IsCommandAvailableFunc = func(commandName string) bool {
+		return commandName == "cat"
+	}
+	expectedCommands := []string{`grep -q "208.80.152.2 abc.com #chaosblade" /etc/hosts`,
+		`cat /etc/hosts | grep -v "208.80.152.2 abc.com #chaosblade" > /tmp/chaos-hosts.tmp && cat /tmp/chaos-hosts.tmp > /etc/hosts`,
+		`grep -q "208.80.152.2 bbc.com #chaosblade" /etc/hosts`,
+		`cat /etc/hosts | grep -v "208.80.152.2 bbc.com #chaosblade" > /tmp/chaos-hosts.tmp && cat /tmp/chaos-hosts.tmp > /etc/hosts`,
+		`rm -rf /tmp/chaos-hosts.tmp`}
+
+	recoverDns(as.domain, as.ip)
+	if exitCode != 0 {
+		t.Errorf("unexpected result: %d, expected result: %d", exitCode, 0)
+	}
+	if !reflect.DeepEqual(expectedCommands, actualCommands) {
+		t.Errorf("unexpected commands: %+v, expected commands: %+v", actualCommands, expectedCommands)
+	}
+}