membank/auto-check/pkg/discovery/discovery.go

package discovery

import (
	"fmt"
	"net"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"sync"
	"time"

	"auto-check/pkg/model"
)

// Devices 设备表（IP → 设备），每轮 Discover() 后更新，包级公开
var Devices = make(map[string]model.Device)

// Scanner 网段扫描器
type Scanner struct {
	CIDR        string
	Timeout     time.Duration
	Concurrency int
}

// NewScanner 创建扫描器
func NewScanner(cidr string, timeout time.Duration, concurrency int) *Scanner {
	return &Scanner{
		CIDR:        cidr,
		Timeout:     timeout,
		Concurrency: concurrency,
	}
}

// Scan 扫描网段，返回存活设备列表
func (s *Scanner) Scan() []model.Device {
	ips, ok := s.candidateIPs()
	if !ok {
		return nil
	}

	fmt.Printf("[扫描] 网段 %s，共 %d 个IP，开始探测...\n", s.CIDR, len(ips))

	results := s.probeAll(ips)

	fmt.Printf("[扫描] 完成，发现 %d 台存活设备\n", len(results))
	return results
}

// candidateIPs 解析网段并生成候选 IP 列表（排除网络地址和广播地址）
// 网段解析失败时返回 false
func (s *Scanner) candidateIPs() ([]net.IP, bool) {
	ip, ipnet, err := net.ParseCIDR(s.CIDR)
	if err != nil {
		fmt.Printf("[扫描] 网段解析失败: %v\n", err)
		return nil, false
	}

	var ips []net.IP
	for ip := ip.Mask(ipnet.Mask); ipnet.Contains(ip); inc(ip) {
		dst := make(net.IP, len(ip))
		copy(dst, ip)
		ips = append(ips, dst)
	}
	// 排除网络地址和广播地址
	if len(ips) > 2 {
		ips = ips[1 : len(ips)-1]
	}
	return ips, true
}

// probeAll 并发探测所有候选 IP，返回存活设备列表
func (s *Scanner) probeAll(ips []net.IP) []model.Device {
	var mu sync.Mutex
	var wg sync.WaitGroup
	var results []model.Device

	sem := make(chan struct{}, s.Concurrency)
	for _, ip := range ips {
		wg.Add(1)
		sem <- struct{}{}
		go func(target net.IP) {
			defer wg.Done()
			defer func() { <-sem }()

			dev, ok := s.probeDevice(target)
			if !ok {
				return
			}
			mu.Lock()
			results = append(results, dev)
			mu.Unlock()
			desc := "无MAC"
			if dev.MAC != "" {
				desc = dev.MAC
			}
			fmt.Printf("  [+] %s 存活 (MAC: %s)\n", target, desc)
		}(ip)
	}
	wg.Wait()
	return results
}

// probeDevice 探测单个 IP，存活则返回设备信息
func (s *Scanner) probeDevice(ip net.IP) (model.Device, bool) {
	if !s.probe(ip) {
		return model.Device{}, false
	}
	// 获取 MAC 地址（ping 探测已触发 ARP 解析）
	dev := model.Device{
		IP:  ip.String(),
		MAC: getMAC(ip.String()),
	}
	return dev, true
}

// probe 用 ping 探测单个 IP 是否存活，返回 bool
func (s *Scanner) probe(ip net.IP) bool {
	waitMs := s.Timeout.Milliseconds()
	if waitMs < 1000 {
		waitMs = 1000 // ping 等待至少 1 秒
	}
	// 各平台 ping 参数不同：Linux -W 秒，macOS/Windows -W/-w 毫秒
	var args []string
	switch runtime.GOOS {
	case "windows":
		args = []string{"-n", "1", "-w", fmt.Sprintf("%d", waitMs), ip.String()}
	case "darwin":
		args = []string{"-c", "1", "-W", fmt.Sprintf("%d", waitMs), ip.String()}
	default:
		args = []string{"-c", "1", "-W", fmt.Sprintf("%d", waitMs/1000), ip.String()}
	}
	return exec.Command("ping", args...).Run() == nil
}

// getMAC 从 ARP 表获取 IP 对应的 MAC 地址（跨平台）
// 返回空字符串表示无法获取
func getMAC(ip string) string {
	switch runtime.GOOS {
	case "linux":
		return macFromProcNetArp(ip)
	case "windows":
		return macFromWindowsArp(ip)
	case "darwin":
		return macFromUnixArp(ip)
	default:
		return macFromUnixArp(ip)
	}
}

// macFromProcNetArp 解析 Linux /proc/net/arp
func macFromProcNetArp(ip string) string {
	data, err := os.ReadFile("/proc/net/arp")
	if err != nil {
		return ""
	}
	for _, line := range strings.Split(string(data), "\n")[1:] {
		fields := strings.Fields(line)
		if len(fields) >= 4 && fields[0] == ip {
			mac := fields[3]
			if mac != "" && mac != "00:00:00:00:00:00" {
				return mac
			}
		}
	}
	return ""
}

// macFromWindowsArp 解析 Windows `arp -a` 输出
func macFromWindowsArp(ip string) string {
	out, err := exec.Command("arp", "-a").Output()
	if err != nil {
		return ""
	}
	for _, line := range strings.Split(string(out), "\n") {
		fields := strings.Fields(line)
		// Windows 格式: IP  MAC  Type
		if len(fields) >= 2 && fields[0] == ip {
			mac := strings.ReplaceAll(fields[1], "-", ":")
			if mac != "" && mac != "00:00:00:00:00:00" {
				return strings.ToLower(mac)
			}
		}
	}
	return ""
}

// macFromUnixArp 解析 macOS/Linux `arp -n` 输出
func macFromUnixArp(ip string) string {
	out, err := exec.Command("arp", "-n", ip).Output()
	if err != nil {
		return ""
	}
	for _, line := range strings.Split(string(out), "\n") {
		fields := strings.Fields(line)
		// macOS 格式: ? (IP) at MAC on en0 ifscope [ethernet]
		for i, f := range fields {
			if f == "("+ip+")" && i+2 < len(fields) {
				mac := fields[i+2]
				if mac != "" && mac != "ff:ff:ff:ff:ff:ff" {
					return strings.ToLower(mac)
				}
			}
		}
	}
	return ""
}

// Discover 扫描并更新包级 Devices map（过滤无 MAC 设备）
func (s *Scanner) Discover() {
	devices := s.Scan()
	Devices = make(map[string]model.Device, len(devices))
	for _, dev := range devices {
		if dev.MAC == "" {
			fmt.Printf("  [!] %s 无 MAC（ARP 未解析），下轮重试\n", dev.IP)
			continue
		}
		Devices[dev.IP] = dev
	}
}

// inc IP 递增
func inc(ip net.IP) {
	for j := len(ip) - 1; j >= 0; j-- {
		ip[j]++
		if ip[j] > 0 {
			break
		}
	}
}