脚本编写学习日记：Go（3）

简易端口探测编写练习：增加了-t与-u参数，支持tcp协议与udp协议，默认tcp协议，可写入输出结果到指定文件,支持从文件中读取多个ip或单个ip，可指定范围端口或多个特定端口）

【似乎Udp协议不可用，尝试了好几种方法都有各种各样的问题,，人已经麻了】

package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

var (
	ip       string
	ports    string
	filepath string
	result   string
	tcp      bool
	udp      bool
)

func init() {
	flag.StringVar(&ip, "i", "", "Enter the ip")
	flag.StringVar(&ports, "p", "21,22", "Enter the port")
	flag.StringVar(&filepath, "f", "ip.txt", "Enter the filepath")
	flag.StringVar(&result, "o", "", "Enter the result filename")
	flag.BoolVar(&tcp, "t", false, "Tcp protocol")
	flag.BoolVar(&udp, "u", false, "Udp protocol")
}

func main() {
	flag.Parse()
	if tcp {
		fmt.Println("正在使用tcp协议扫描")
	} else if udp {
		fmt.Println("正在使用udp协议扫描")
	} else {
		fmt.Println("默认使用tcp协议扫描")
	}
	var wg sync.WaitGroup
	starttime := time.Now()
	if strings.TrimSpace(ip) != "" {
		wg.Add(1)
		go scanonly(ip, ports, &wg, tcp, udp)
	} else if strings.TrimSpace(filepath) != "" {
		ips, err := os.Open(filepath)
		if err != nil {
			panic(err)
		}
		defer func(ips *os.File) {
			err = ips.Close()
			if err != nil {

			}
		}(ips)
		ipline := bufio.NewReader(ips)
		for {
			content, _, errs := ipline.ReadLine()
			if errs == io.EOF {
				break
			}
			content1 := string(content)
			wg.Add(1)
			go scanonly(content1, ports, &wg, tcp, udp)
		}
	} else {
		fmt.Println("傻鸟,-h 看一下")
	}
	wg.Wait()
	endtime := time.Since(starttime)
	fmt.Println("扫描时长：", endtime)
}

func scanonly(ip string, ports string, wg *sync.WaitGroup, tcp bool, udp bool) {
	defer wg.Done()
	fmt.Println("当前检验的ip :", ip)
	if strings.Contains(ports, ",") || (!strings.Contains(ports, ",") && !strings.Contains(ports, "-")) {
		portlist := strings.Split(ports, ",")
		for _, port := range portlist {
			wen := make(chan int)
			portint, _ := strconv.Atoi(port)
			wg.Add(1)
			go scan(ip, portint, wen, tcp, udp)
			port_true := <-wen
			if port_true != 0 {
				fmt.Printf("%s:%d is open !!! \n", ip, port_true)
				port_true1 := strconv.Itoa(port_true)
				if strings.TrimSpace(result) != "" {
					go saveResult(result, ip, port_true1, wg, tcp, udp)
				} else {
					wg.Done()
				}
			} else {
				wg.Done()
			}
		}
	} else if strings.Contains(ports, "-") {
		portlist := strings.Split(ports, "-")
		start, _ := strconv.Atoi(portlist[0])
		end, _ := strconv.Atoi(portlist[1])
		portlists := make([]int, 0)
		for i := start; i <= end; i++ {
			portlists = append(portlists, i)
		}
		for _, port := range portlists {
			wen := make(chan int)
			wg.Add(1)
			go scan(ip, port, wen, tcp, udp)
			port_true := <-wen
			if port_true != 0 {
				fmt.Printf("%s:%d is open !!! \n", ip, port_true)
				port_true1 := strconv.Itoa(port_true)
				if strings.TrimSpace(result) != "" {
					go saveResult(result, ip, port_true1, wg, tcp, udp)
				} else {
					wg.Done()
				}
			} else {
				wg.Done()
			}
		}
	}
}

func scan(ip string, port int, wen chan int, tcp bool, udp bool) {
	if tcp {
		conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", ip, port), time.Second*5)
		if err != nil {
			wen <- 0
			return
		}
		defer conn.Close()
	} else if udp {
		addr := &net.UDPAddr{IP: net.ParseIP(ip), Port: port}
		conn, err := net.ListenPacket("udp", addr.String())
		if err != nil {
			wen <- 0
			return
		}
		defer conn.Close()
		conn.SetDeadline(time.Now().Add(time.Second * 5))
		_, err = conn.WriteTo([]byte("hello"), addr)
		if err != nil {
			wen <- 0
			return
		}
		buffer := make([]byte, 1024)
		_, _, err = conn.ReadFrom(buffer)
		if err != nil {
			if err, ok := err.(net.Error); ok && err.Timeout() {
				wen <- 0
				return
			}
		}
		wen <- port
	} else {
		conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", ip, port), time.Second*5)
		if err != nil {
			wen <- 0
			return
		}
		defer conn.Close()
	}
	wen <- port
}

func saveResult(result string, ip string, port string, wg *sync.WaitGroup, tcp bool, udp bool) {
	defer wg.Done()
	file, err := os.OpenFile(result, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		fmt.Println("Error opening file:", err)
		return
	}
	defer file.Close()
	if tcp {
		_, err = file.WriteString("tcp:" + ip + ":" + port + "\n")
		if err != nil {
			fmt.Println("Error writing to file:", err)
		}
	} else if udp {
		_, err = file.WriteString("udp:" + ip + ":" + port + "\n")
		if err != nil {
			fmt.Println("Error writing to file:", err)
		}
	} else {
		_, err = file.WriteString("tcp:" + ip + ":" + port + "\n")
		if err != nil {
			fmt.Println("Error writing to file:", err)
		}
	}
}