diff --git a/pkg/cli/parse.go b/pkg/cli/parse.go index f800ffe..be79a8e 100644 --- a/pkg/cli/parse.go +++ b/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,7 +53,7 @@ func ParseCDKMain() { // https://github.com/jiguangin/netcat PassInnerArgs() netcat.RunVendorNetcat() - return + return true } // docopt argparse start @@ -61,7 +61,7 @@ func ParseCDKMain() { 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[""] != nil { @@ -174,4 +174,6 @@ func ParseCDKMain() { docopt.PrintHelpAndExit(nil, BannerContainer) } } + + return false } diff --git a/pkg/cli/parse_test.go b/pkg/cli/parse_test.go new file mode 100644 index 0000000..ee74a5f --- /dev/null +++ b/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() + } +} diff --git a/pkg/exploit/test_poc.go b/pkg/exploit/test_poc.go index 3b11232..80a592f 100644 --- a/pkg/exploit/test_poc.go +++ b/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[""]` by yourself. // example - fmt.Printf("%v", cli.Args[""].([]string)) + if cli.Args[""] != nil { + log.Printf("%v", cli.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 diff --git a/pkg/tool/network/network.go b/pkg/tool/network/network.go index 099a8c0..83b26d3 100644 --- a/pkg/tool/network/network.go +++ b/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()))