Merge pull request #53 from cdk-team/probe-get-ending

perf(probe): output ending message

pkg/tool/probe/net.go

@@ -1,6 +1,6 @@
+//go:build !no_probe_tool
 // +build !no_probe_tool
 
-
 /*
 Copyright 2022 The Authors of https://github.com/CDK-TEAM/CDK .
 
@@ -22,13 +22,14 @@ package probe
 import (
 	"context"
 	"fmt"
-	"github.com/cdk-team/CDK/conf"
-	"golang.org/x/sync/semaphore"
 	"log"
 	"net"
 	"strings"
 	"sync"
 	"time"
+
+	"github.com/cdk-team/CDK/conf"
+	"golang.org/x/sync/semaphore"
 )
 
 type PortScanner struct {
@@ -54,7 +55,7 @@ func ScanPort(ip string, port int, timeout time.Duration) bool {
 	return true
 }
 
-func (ps *PortScanner) Start(portFromTo []FromTo) {
+func (ps *PortScanner) Start() {
 	wg := sync.WaitGroup{}
 	defer wg.Wait()
 
@@ -68,7 +69,7 @@ func (ps *PortScanner) Start(portFromTo []FromTo) {
 	for ipExt := start; ipExt <= end; ipExt++ {
 		ip := base + "." + fmt.Sprintf("%d", ipExt)
 		// iterate port in task list
-		for _, p := range portFromTo {
+		for _, p := range ps.portRange {
 			// iterate port from A-B
 			for port := p.From; port <= p.To; port++ {
 				// lock down the context
@@ -88,25 +89,34 @@ func (ps *PortScanner) Start(portFromTo []FromTo) {
 
 func TCPScanExploitAPI(ipRange string) {
 	portFromTo, _ := GetTaskPortList()
-	ps := &PortScanner{
-		ipRange:   ipRange,
-		portRange: portFromTo,
-		lock:      semaphore.NewWeighted(conf.TCPScannerConf.MaxParallel),
-		timeout:   conf.TCPScannerConf.Timeout,
-	}
-	log.Printf("scanning %v with pre-defined ports, max parallels:%v, timeout:%v\n", ps.ipRange, conf.TCPScannerConf.MaxParallel, ps.timeout)
-	ps.Start(portFromTo)
+	timeout := time.Duration(conf.TCPScannerConf.Timeout) * time.Millisecond
+
+	TCPPScan(ipRange, portFromTo, conf.TCPScannerConf.MaxParallel, timeout)
 }
 
 func TCPScanToolAPI(ipRange string, portRange string, parallel int64, timeoutMS int) {
 	portFromTo, _ := GetTaskPortListByString(portRange)
+	timeout := time.Duration(timeoutMS) * time.Millisecond
+
+	TCPPScan(ipRange, portFromTo, parallel, timeout)
+}
+
+func TCPPScan(ipRange string, portRange []FromTo, parallel int64, timeout time.Duration) {
 
 	ps := &PortScanner{
 		ipRange:   ipRange,
-		portRange: portFromTo,
+		portRange: portRange,
 		lock:      semaphore.NewWeighted(parallel),
-		timeout:   time.Duration(timeoutMS) * time.Millisecond,
+		timeout:   timeout,
 	}
+
+	startTime := time.Now()
 	log.Printf("scanning %v with user-defined ports, max parallels:%v, timeout:%v\n", ps.ipRange, parallel, ps.timeout)
-	ps.Start(portFromTo)
+	ps.Start()
+
+	endTime := time.Now()
+	useTime := int64(endTime.Sub(startTime).Seconds() * 1000)
+	log.Printf("scanning use time:%vms\n", useTime)
+	log.Printf("ending; @args is ips: %v, max parallels:%v, timeout:%v\n", ps.ipRange, conf.TCPScannerConf.MaxParallel, ps.timeout)
+
 }