Merge pull request #4 from misaka4e21/master

Implement receiving commandline arguments.

Author: sswolf

lmap.go

@@ -1,6 +1,9 @@
 package main
 
 import (
+	"os"
+	"flag"
+	"sort"
 	"bytes"
 	"encoding/binary"
 	"fmt"
@@ -12,8 +15,6 @@ import (
 
 var icmp ICMP
 
-var wg sync.WaitGroup
-
 type ICMP struct {
 	Type        uint8
 	Code        uint8
@@ -23,41 +24,64 @@ type ICMP struct {
 }
 
 func main() {
-	wg.Add(1)
-	go func() {
-		CheckIP()
-		wg.Done()
-	}()
-	wg.Wait()
-	CheckIP()
+	isVerbose := false
+	flag.BoolVar(&isVerbose, "v", false, "be verbose")
+	flag.Parse()
+	args := flag.Args()
+	if len(args) < 1 {
+		fmt.Printf("使用方法：%s [-v] <网络号>/<CIDR>\n", os.Args[0])
+		os.Exit(-1)
+	}
+	CheckIP(args[0], isVerbose)
 }
 
-func CheckIP() {
+func CheckIP(subnet string, isVerbose bool) {
+	var usedIP []string
+	var unusedIP []string
 	checkerGroup := &sync.WaitGroup{}
 	t := time.Now()
-	hosts, _ := hosts("10.150.1.1/24")
-	usedIP := make([]string, 0, len(hosts))
-	unusedIP := make([]string, 0, len(hosts))
+	hosts, _ := getAllHostsFromCIDR(subnet)
 	for _, ip := range hosts {
+		time.Sleep(500)
 		checkerGroup.Add(1)
 		go func(data string) {
-			pong := ping(data)
-			if pong {
+			defer checkerGroup.Done()
+			isUsed := ping(data)
+			if isUsed {
 				usedIP = append(usedIP, data)
+				if isVerbose {
+					fmt.Println("已使用IP：", usedIP)
+				}
 			} else {
 				unusedIP = append(unusedIP, data)
+				if isVerbose {
+					fmt.Println("未使用IP：", unusedIP)
+				}
 			}
-			checkerGroup.Done()
 		}(ip)
 	}
 	checkerGroup.Wait()
-	fmt.Println("已使用IP：", usedIP)
-	fmt.Println("未使用IP：", unusedIP)
 	elapsed := time.Since(t)
 	fmt.Println("IP扫描完成,耗时", elapsed)
+	fmt.Println("已使用IP：", sortIPList(usedIP))
+	fmt.Println("未使用IP：", sortIPList(unusedIP))
+}
+
+func sortIPList(ipStrings []string) (result []string) {
+	var ips []net.IP
+	for _, ipString := range ipStrings {
+		ips = append(ips, net.ParseIP(ipString))
+	}
+	sort.Slice(ips, func(i, j int) bool {
+		return bytes.Compare(ips[i], ips[j]) < 0
+	})
+	for _, ip := range ips {
+		result = append(result, ip.String())
+	}
+	return
 }
 
-func hosts(cidr string) ([]string, error) {
+func getAllHostsFromCIDR(cidr string) ([]string, error) {
 	ip, ipnet, err := net.ParseCIDR(cidr)
 	if err != nil {
 		return nil, err
@@ -97,8 +121,7 @@ func ping(ip string) bool {
 	buffer.Reset()
 	binary.Write(&buffer, binary.BigEndian, icmp)
 
-	Time, _ := time.ParseDuration("1s")
-	conn, err := net.DialTimeout("ip4:icmp", ip, Time)
+	conn, err := net.DialTimeout("ip4:icmp", ip, 1 * time.Second)
 	if err != nil {
 		return false
 	}