正对#16优化

cmd/lmap/main.go

@@ -29,11 +29,30 @@ import (
 func main() {
 	isVerbose := false
 	flag.BoolVar(&isVerbose, "v", false, "be verbose")
+
+	subnets := multiArg{}
+	flag.Var(&subnets, "subnet", "network to scan (can be specified multiple times)")
+
+	excludes := multiArg{}
+	flag.Var(&excludes, "exclude", "IP or subnet to exclude (can be specified multiple times)")
+
 	flag.Parse()
-	args := flag.Args()
-	if len(args) < 1 {
-		_, _ = fmt.Fprintf(os.Stderr, "使用方法：%s [-v] <网络号>/<CIDR>\n", os.Args[0])
+
+	if len(subnets) < 1 {
+		_, _ = fmt.Fprintf(os.Stderr, "使用方法：%s [-v] -subnet <网络号>/<CIDR> [-subnet ...] [-exclude <IP>/<CIDR>] [-exclude ...]\n", os.Args[0])
 		os.Exit(-1)
 	}
-	lmap.CheckIP(args[0], isVerbose)
+
+	lmap.CheckIP(subnets, excludes, isVerbose)
+}
+
+type multiArg []string
+
+func (m *multiArg) String() string {
+	return fmt.Sprintf("%v", *m)
+}
+
+func (m *multiArg) Set(value string) error {
+	*m = append(*m, value)
+	return nil
 }

pkg/lmap/check_ip.go

@@ -27,33 +27,39 @@ import (
 
 const OUTPUT_IP_PER_LINE = 3
 
-func CheckIP(subnet string, isVerbose bool) {
+func CheckIP(subnets, excludes []string, isVerbose bool) {
 	checkerGroup := &sync.WaitGroup{}
 	t := time.Now()
-	hosts, _ := GetAllIPsFromCIDR(subnet)
-	for index := range hosts {
-		//time.Sleep(500)
+
+	var allHosts []HostInfo
+	for _, subnet := range subnets {
+		hosts, _ := GetAllIPsFromCIDR(subnet, excludes)
+		allHosts = append(allHosts, hosts...)
+	}
+
+	for index := range allHosts {
 		checkerGroup.Add(1)
 
 		go func(index int) {
 			defer checkerGroup.Done()
-			hosts[index].isUsed = Ping(hosts[index].host)
+			allHosts[index].isUsed = Ping(allHosts[index].host)
 			if isVerbose {
-				if hosts[index].isUsed {
-					fmt.Println("已使用IP：", hosts[index].host.String())
+				if allHosts[index].isUsed {
+					fmt.Println("已使用IP：", allHosts[index].host.String())
 				} else {
-					fmt.Println("未使用IP：", hosts[index].host.String())
+					fmt.Println("未使用IP：", allHosts[index].host.String())
 				}
 			}
 		}(index)
 	}
+
 	checkerGroup.Wait()
 	elapsed := time.Since(t)
 	_, _ = fmt.Fprintln(os.Stderr, "IP扫描完成,耗时", elapsed)
 	fmt.Println("已使用IP：")
-	printIPList(hosts, true)
+	printIPList(allHosts, true)
 	fmt.Println("未使用IP：")
-	printIPList(hosts, false)
+	printIPList(allHosts, false)
 }
 
 func printIPList(hosts []HostInfo, boolFilter bool) {

pkg/lmap/parse_subnet.go

@@ -18,27 +18,55 @@
 
 package lmap
 
-import "net"
+import (
+	"net"
+	"strings"
+)
 
-func GetAllIPsFromCIDR(cidr string) ([]HostInfo, error) {
+func GetAllIPsFromCIDR(cidr string, excludes []string) ([]HostInfo, error) {
 	ip, ipNet, err := net.ParseCIDR(cidr)
 	if err != nil {
 		return nil, err
 	}
 
 	var ips []HostInfo
 	for ip := ip.Mask(ipNet.Mask); ipNet.Contains(ip); inc(ip) {
-		ips = append(ips, HostInfo{
-			host:   dupIP(ip),
-			isUsed: false,
-		})
+		ipStr := ip.String()
+		if !isExcluded(ipStr, excludes) {
+			ips = append(ips, HostInfo{
+				host:   dupIP(ip),
+				isUsed: false,
+			})
+		}
 	}
 	if len(ips) <= 2 {
 		return ips, nil
 	}
 	return ips[0 : len(ips)-1], nil
 }
 
+func isExcluded(ip string, excludes []string) bool {
+	for _, exclude := range excludes {
+		if strings.Contains(exclude, "/") {
+			_, excludeNet, err := net.ParseCIDR(exclude)
+			if err != nil {
+				//fmt.Printf("解析排除规则 %s 失败: %v\n", exclude, err)
+				continue
+			}
+			if excludeNet.Contains(net.ParseIP(ip)) {
+				//fmt.Printf("IP %s 被排除规则 %s 排除\n", ip, exclude)
+				return true
+			}
+		} else {
+			if ip == exclude {
+				//fmt.Printf("IP %s 被排除规则 %s 排除\n", ip, exclude)
+				return true
+			}
+		}
+	}
+	return false
+}
+
 func inc(ip net.IP) {
 	for j := len(ip) - 1; j >= 0; j-- {
 		ip[j]++
@@ -49,7 +77,6 @@ func inc(ip net.IP) {
 }
 
 func dupIP(ip net.IP) net.IP {
-	// To save space, try and only use 4 bytes
 	if x := ip.To4(); x != nil {
 		ip = x
 	}