go lang编程练习 TCP 扫瞄器

端口扫描 测试结构体

package TestPort

import (
	"errors"
	"fmt"
	"net"
	"sort"
	"strconv"
	"strings"
	"sync"
)

/*
 功能需求： 尽可能在设备承受负载情况下 尽快的TCP扫描 指定ip 下的端口是否打开

0. 根据 go 下边的net 包进行去 Dial 方法就去测试
1. 使用 goroutine 并发模式去遍历。
2. 使用channel 保证线程之间的操作安全
3. 使用flag 包进行接收参数信息，不符合参数类型 提示
*/

//返回来的端口信息
type DataPortInfo struct {
	IsOpen  bool
	PortInt int
}

//TestPort
type TestPort struct {
	CurIp        string
	PortChan     chan int
	ResultChan   chan DataPortInfo
	OpenPort     []int
	ClosePort    []int
	Wg           sync.WaitGroup
	mutex        sync.Mutex
	threadNumber int
	StaPorNu     int
	EndPorNu     int
}

func init() {

}

//自己的方法判断ip地址是否合法
func IsVaildIp(strIp string) bool {
	ipStrArr := strings.Split(strIp, ".")
	if len(ipStrArr) == 4 {
		var (
			tempInt int
			err     error
		)
		for idx := 0; idx < len(ipStrArr); idx++ {
			tempInt, err = strconv.Atoi(ipStrArr[idx])
			if err != nil {
				return false
			} else if tempInt < 0 || tempInt > 255 {
				return false
			}
		}
		return true
	}
	return false
}

//创建一个 TestPort 结构体
func CreateTestPort(stratPort, endPort, thNu int, ip string) (ts TestPort, err error) {
	ts = TestPort{}
	if endPort < stratPort {
		err = errors.New("开始端口号大于结束端口号！！")
		return
	}
	if thNu <= 0 {
		err = errors.New("线程数不能小于等于零！")
		return
	}
	if IPv4 := net.ParseIP(ip); IPv4 == nil {
		err = errors.New("ip地址不合法！")
		return
	}
	ts.PortChan = make(chan int, 100)
	ts.ResultChan = make(chan DataPortInfo, endPort-stratPort+1)
	ts.StaPorNu = stratPort
	ts.EndPorNu = endPort
	ts.CurIp = ip
	ts.Wg = sync.WaitGroup{}
	ts.threadNumber = thNu
	err = nil
	return
}

//开始测试端口
func (its *TestPort) Test() bool {
	if len(its.CurIp) == 0 || its.ResultChan == nil || its.threadNumber <= 0 {
		return false
	}

	// its.mutex.Lock()
	// defer its.mutex.Unlock()

	//0. 结果集重新初始化
	its.OpenPort = make([]int, 0)
	its.ClosePort = make([]int, 0)
	its.ResultChan = make(chan DataPortInfo, its.EndPorNu-its.StaPorNu+1)
	// 1. 先创建 等待工作线程
	for i := 0; i < its.threadNumber; i++ {
		its.Wg.Add(1)
		go myGoroutine(&its.PortChan, its.CurIp, &its.ResultChan, &its.Wg)
	}

	//2. 遍历接口线程 往端口管道中放入要测试的端口
	var pushDataWg sync.WaitGroup = sync.WaitGroup{}
	pushDataWg.Add(1)
	go func(ch *chan int, wg *sync.WaitGroup, startPor, endPor int) {
		for i := startPor; i < endPor+1; i++ {
			*ch <- i
		}
		wg.Done()
	}(&its.PortChan, &pushDataWg, its.StaPorNu, its.EndPorNu)
	//3. 遍历结果集从结果集 中拿出结果 放到各自的切片集合中 （开个协程做最好）
	var resultWg sync.WaitGroup = sync.WaitGroup{}
	resultWg.Add(1)
	go func(resChan *chan DataPortInfo, openSlice, closeSlice *[]int, wg *sync.WaitGroup) {
		for {
			v, ok := <-*resChan
			if ok {
				if v.IsOpen {
					*openSlice = append(*openSlice, v.PortInt)
				} else {
					*closeSlice = append(*closeSlice, v.PortInt)
				}
			} else {
				break
			}
		}
		wg.Done()
	}(&its.ResultChan, &its.OpenPort, &its.ClosePort, &resultWg)

	pushDataWg.Wait()    /* 感觉是不是有点堵屁眼的感觉， 虽然各个工作都开了各自协程去工作，但是 时间有先后顺序 这样能加快速度吗？？*/
	//4. 关闭缓冲放入端口管道
	close(its.PortChan)
	its.Wg.Wait()
	//5. 关闭收集结果通道
	close(its.ResultChan)
	//6. 等待结果协程完成
	resultWg.Wait()
	//7. 对结果进行切片排序
	sort.Ints(its.OpenPort)
	sort.Ints(its.ClosePort)
	return true
}

func myGoroutine(ch *chan int, ipStr string, resultChan *chan DataPortInfo, wg *sync.WaitGroup) {
	for {
		port, ok := <-*ch
		if ok {
			con, err := net.Dial("tcp", fmt.Sprintf("%s:%d", ipStr, port))
			portInfo := DataPortInfo{}
			portInfo.PortInt = port
			if err != nil {
				//说明此IP地址的端口没有被TCP server占用
				portInfo.IsOpen = false
			} else {
				portInfo.IsOpen = true
				con.Close()
			}
			//入结果管道
			*resultChan <- portInfo
		} else {
			break
		}
	}
	wg.Done()
}

主程序代码：

package main

import (
	"One/TestPort"
	"flag"
	"fmt"
	"sort"
	"strings"
	"sync"
	"time"
)

func main() {
	Test()
}

func Test() {
	var (
		inputIP  string
		sport    int
		eport    int
		threadNu int
	)

	flag.StringVar(&inputIP, "ip", "", "输入待测试的IP地址")
	flag.IntVar(&sport, "sport", 1, "测试开始端口号")
	flag.IntVar(&eport, "eport", 2, "测试后结束端口号")
	flag.IntVar(&threadNu, "threadNu", 10, "并发线程")
	flag.Parse()
	if len(inputIP) == 0 {
		fmt.Println("应该输入参数 -ip （172.16.3.194）  -SPort（开始端口号） -EPort（结束端口号） -ThreadNu（并发线程个数）  ")
		return
	}
	myTestPort, err := TestPort.CreateTestPort(sport, eport, threadNu, inputIP)
	if err != nil {
		fmt.Println("创建结果失败！" + err.Error())
		return
	}
	startTime := time.Now()
	bl := myTestPort.Test()
	timeSeconds := time.Since(startTime) / 1e9
	if !bl {
		fmt.Println("测试端口失败！")
		return
	}
	fmt.Println(fmt.Sprintf("扫描完成，耗时：%d s", timeSeconds))
	var sb strings.Builder
	sb.WriteString(inputIP)
	sb.WriteString(" 地址  TCP 被打开了端口：")
	sb.WriteString(fmt.Sprintln(myTestPort.OpenPort))
	sb.WriteString(" TCP 关闭端口：")
	sb.WriteString(fmt.Sprintln(myTestPort.ClosePort))
	fmt.Println(sb.String())
}

嗯感觉不太好