test(main function): add unit test for ParseCDKMain

pkg/cli/parse.go

@@ -42,7 +42,7 @@ func PassInnerArgs() {
 	os.Args = os.Args[1:]
 }
 
-func ParseCDKMain() {
+func ParseCDKMain() bool {
 
 	if len(os.Args) == 1 {
 		docopt.PrintHelpAndExit(nil, BannerContainer)
@@ -53,15 +53,15 @@ func ParseCDKMain() {
 		// https://github.com/jiguangin/netcat
 		PassInnerArgs()
 		netcat.RunVendorNetcat()
-		return
+		return true
 	}
 
 	// docopt argparse start
 	parseDocopt()
 
 	if Args["auto-escape"].(bool) {
 		plugin.RunSingleTask("auto-escape")
-		return
+		return true
 	}
 
 	// support for cdk eva(Evangelion) and cdk evaluate
@@ -116,7 +116,7 @@ func ParseCDKMain() {
 			evaluate.DumpCgroup()
 
 		}
-		return
+		return true
 	}
 
 	if Args["run"].(bool) {
@@ -128,10 +128,10 @@ func ParseCDKMain() {
 		if plugin.Exploits[name] == nil {
 			fmt.Printf("\nInvalid script name: %s , available scripts:\n", name)
 			plugin.ListAllExploit()
-			return
+			return true
 		}
 		plugin.RunSingleExploit(name)
-		return
+		return true
 	}
 
 	if Args["<tool>"] != nil {
@@ -174,4 +174,6 @@ func ParseCDKMain() {
 			docopt.PrintHelpAndExit(nil, BannerContainer)
 		}
 	}
+
+	return false
 }

pkg/cli/parse_test.go

@@ -0,0 +1,122 @@
+/*
+Copyright 2022 The Authors of https://github.com/CDK-TEAM/CDK .
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package cli
+
+import (
+	"bytes"
+	// "io/ioutil"
+	"log"
+	"os"
+	"testing"
+	"time"
+
+	_ "github.com/cdk-team/CDK/pkg/exploit" // register all exploits
+	_ "github.com/cdk-team/CDK/pkg/task"    // register all task
+)
+
+type testArgsCase struct {
+	name string
+	args []string
+	successStr string
+}
+
+func doParseCDKMainWithTimeout() {
+
+	result := make(chan bool, 1)
+
+	go func ()  {
+		result <- ParseCDKMain()
+	}()
+
+	select {
+	case <-time.After(time.Second * 2):
+		log.Println("check run ok, timeout in 2s, and return.")
+		return
+	case <-result:
+		return 
+	}
+
+}
+
+func TestParseCDKMain(t *testing.T) {
+
+
+	// ./cdk eva 2>&1 | head
+	// ./cdk run test-poc | head
+	// ./cdk ifconfig | head
+
+	tests := []testArgsCase{
+		{
+			name: "./cdk eva",
+			args: []string{"./cdk_cli_path", "eva"},
+			successStr: "current user",
+		},
+		{
+			name: "./cdk run test-poc",
+			args: []string{"./cdk_cli_path", "run", "test-poc"},
+			successStr: "run success",
+		},
+		{
+			name: "./cdk ifconfig",
+			args: []string{"./cdk_cli_path", "ifconfig"},
+			successStr: "GetLocalAddresses",
+		},
+	}
+
+	for _, tt := range tests {
+
+		// fmt.Print and log.Print to buffer, and check output
+		var buf bytes.Buffer
+		log.SetOutput(&buf)
+
+		// hook fmt.X to buffer, hook os.Stdout
+		// oldStdout := os.Stdout
+		// r, w, _ := os.Pipe()
+		// os.Stdout = w
+
+		// hook os.Args
+		args := tt.args
+		os.Args = args
+
+		t.Run(tt.name, func(t *testing.T) {
+			doParseCDKMainWithTimeout()
+			// out, _ := ioutil.ReadAll(r)
+
+			// check success string in buf and out
+			// if !bytes.Contains(buf.Bytes(), []byte(tt.successStr)) && !bytes.Contains(out, []byte(tt.successStr)) {
+			// 	t.Errorf(("parse cdk main failed, name: %s, args: %v, buf: %s, out: %s"), tt.name, tt.args, buf.String()[:1000], string(out)[:1000])
+			// }
+
+			if !bytes.Contains(buf.Bytes(), []byte(tt.successStr)) {
+
+				// get sub string from buf, lenght is 1000
+				str := buf.String()
+				if len(str) > 1000 {
+					str = str[:1000]
+				}
+
+				t.Errorf(("parse cdk main failed, name: %s, args: %v, buf: %s"), tt.name, tt.args, str)
+			}
+
+
+		})
+
+		// return to os.Stdout default
+		// os.Stdout = oldStdout
+		// w.Close()
+	}
+}

pkg/exploit/test_poc.go

@@ -1,4 +1,3 @@
-
 /*
 Copyright 2022 The Authors of https://github.com/CDK-TEAM/CDK .
 
@@ -18,7 +17,8 @@ limitations under the License.
 package exploit
 
 import (
-	"fmt"
+	"log"
+
 	"github.com/cdk-team/CDK/pkg/cli"
 	"github.com/cdk-team/CDK/pkg/plugin"
 )
@@ -32,10 +32,12 @@ func (p TEST) Desc() string {
 func (p TEST) Run() bool {
 	// if your script needs input, handle `var lib.Args["<args>"]` by yourself.
 	// example
-	fmt.Printf("%v", cli.Args["<args>"].([]string))
+	if cli.Args["<args>"] != nil {
+		log.Printf("%v", cli.Args["<args>"].([]string))
+	}
 
 	// functional codes
-	fmt.Printf("\n[test] run success\n")
+	log.Printf("\n[test] run success\n")
 
 	// return [true/false] when [normal exit/fatal]
 	return true

pkg/tool/network/network.go

@@ -24,6 +24,7 @@ import (
 )
 
 func GetLocalAddresses() {
+	log.Printf("[+] run ifconfig, using GetLocalAddresses()")
 	ifaces, err := net.Interfaces()
 	if err != nil {
 		log.Print(fmt.Errorf("localAddresses: %v\n", err.Error()))