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

package discovery

import (
	"context"
	"fmt"
	"io"
	"net"
	"os/exec"
	"strings"
	"sync"
	"time"

	"auto-check/pkg/model"
	"golang.org/x/text/encoding/simplifiedchinese"
	"golang.org/x/text/transform"
)

// Scanner 网段扫描器（并发 ping 探测，支持 CIDR 和 IP 范围格式）
type Scanner struct {
	CIDR        string
	Timeout     time.Duration // 单 IP 探测超时
	Concurrency int           // 并发数，默认 30
}

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

// Scan 并发 ping 扫描网段，返回存活设备列表（仅含 IP，MAC 留空待 SSH 阶段回填）
func (s *Scanner) Scan() []model.Device {
	ips, ok := s.candidateIPs()
	if !ok {
		return nil
	}
	fmt.Printf("[扫描] 目标 %s，共 %d 个IP，%d 并发 ping 探测...\n", s.CIDR, len(ips), s.Concurrency)

	jobs := make(chan string, len(ips))
	results := make(chan model.Device, len(ips))

	var wg sync.WaitGroup
	for i := 0; i < s.Concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for ip := range jobs {
				if s.ping(ip) {
					results <- model.Device{IP: ip, MAC: ""}
				}
			}
		}()
	}

	for _, ip := range ips {
		jobs <- ip
	}
	close(jobs)

	wg.Wait()
	close(results)

	var devices []model.Device
	for dev := range results {
		devices = append(devices, dev)
	}

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

// candidateIPs 解析网段并生成候选 IP 列表。
// 支持两种格式：
//   - CIDR: "10.0.3.0/24"（排除网络地址和广播地址）
//   - IP 范围: "10.0.3.1-10.0.3.100"（包含起始和结束 IP）
func (s *Scanner) candidateIPs() ([]string, bool) {
	// 检查是否为 IP 范围格式
	if strings.Contains(s.CIDR, "-") {
		return s.parseIPRange()
	}
	// 否则按 CIDR 解析
	return s.parseCIDR()
}

// parseCIDR 解析 CIDR 格式并生成候选 IP 列表
func (s *Scanner) parseCIDR() ([]string, bool) {
	_, ipnet, err := net.ParseCIDR(s.CIDR)
	if err != nil {
		fmt.Printf("[扫描] 网段解析失败: %v\n", err)
		return nil, false
	}
	ip := ipnet.IP.Mask(ipnet.Mask)
	var ips []string
	for ; ipnet.Contains(ip); inc(ip) {
		ips = append(ips, ip.String())
	}
	if len(ips) > 2 {
		ips = ips[1 : len(ips)-1]
	}
	return ips, true
}

// parseIPRange 解析 IP 范围格式（如 "10.0.3.1-10.0.3.100"）
func (s *Scanner) parseIPRange() ([]string, bool) {
	parts := strings.SplitN(s.CIDR, "-", 2)
	if len(parts) != 2 {
		fmt.Printf("[扫描] IP 范围格式无效: %s\n", s.CIDR)
		return nil, false
	}
	startIP := net.ParseIP(strings.TrimSpace(parts[0]))
	endIP := net.ParseIP(strings.TrimSpace(parts[1]))
	if startIP == nil || endIP == nil {
		fmt.Printf("[扫描] IP 地址解析失败: %s\n", s.CIDR)
		return nil, false
	}
	// 转换为 4 字节 IPv4
	start := startIP.To4()
	end := endIP.To4()
	if start == nil || end == nil {
		fmt.Printf("[扫描] 仅支持 IPv4 范围: %s\n", s.CIDR)
		return nil, false
	}
	// 确保起始 <= 结束
	if ipToUint32(start) > ipToUint32(end) {
		fmt.Printf("[扫描] 起始 IP 大于结束 IP: %s\n", s.CIDR)
		return nil, false
	}
	var ips []string
	for ip := make(net.IP, len(start)); ; {
		copy(ip, start)
		ips = append(ips, ip.String())
		if start.Equal(end) {
			break
		}
		inc(start)
	}
	return ips, true
}

// ipToUint32 将 IPv4 地址转换为 uint32 用于比较
func ipToUint32(ip net.IP) uint32 {
	ip = ip.To4()
	return uint32(ip[0])<<24 | uint32(ip[1])<<16 | uint32(ip[2])<<8 | uint32(ip[3])
}

// ping 探测单个 IP 是否存活（ICMP，经系统 ping 命令）。
// 仅凭 exit code 不可靠（Windows 在某些代答/虚拟网卡场景会返回 0），
// 因此解析输出，要求至少收到 1 个回包。
func (s *Scanner) ping(ip string) bool {
	// ping 命令的等待超时（毫秒）：固定 1s 足够局域网探测
	const pingWait = 1000
	// context 超时设为 ping 等待的 2 倍，留足 ping 自然退出时间，
	// 避免 context 先于 ping -w 杀进程导致拿不到输出。
	ctx, cancel := context.WithTimeout(context.Background(), 2*pingWait*time.Millisecond)
	defer cancel()
	cmd := exec.CommandContext(ctx, "ping", "-n", "1", "-w", fmt.Sprintf("%d", pingWait), ip)
	out, err := cmd.Output()
	if err != nil {
		return false
	}
	// Windows 下 ping 输出为 GBK 编码，解码为 UTF-8 后再解析
	text := decodeGBK(out)
	return parsePingReply(text)
}

// parsePingReply 判断 ping 输出是否表示存活。
// 兼容中英文 Windows 输出：英文 "Received = 1" / 中文 "已接收 = 1"。
// 匹配接收数关键字后取其后数字，>=1 即视为存活。
func parsePingReply(out string) bool {
	for _, line := range strings.Split(out, "\n") {
		l := strings.ToLower(strings.TrimSpace(line))
		// 提取关键字后的接收数，定位到 '=' 之后
		kw := "received ="
		idx := strings.Index(l, kw)
		if idx < 0 {
			kw = "已接收 ="
			idx = strings.Index(l, kw)
		}
		if idx < 0 {
			continue
		}
		eq := strings.Index(l[idx:], "=")
		if eq < 0 {
			continue
		}
		n := strings.TrimSpace(l[idx+eq+1:])
		cnt := 0
		for _, c := range n {
			if c < '0' || c > '9' {
				break
			}
			cnt = cnt*10 + int(c-'0')
		}
		return cnt >= 1
	}
	return false
}

// inc IP 递增（原地修改，调用方负责复制）
func inc(ip net.IP) {
	for j := len(ip) - 1; j >= 0; j-- {
		ip[j]++
		if ip[j] > 0 {
			break
		}
	}
}

// decodeGBK 将 GBK 编码字节转为 UTF-8 字符串（Windows ping 输出为 GBK）。
// 若解码失败则原样返回。
func decodeGBK(b []byte) string {
	r := transform.NewReader(strings.NewReader(string(b)), simplifiedchinese.GBK.NewDecoder())
	out, err := io.ReadAll(r)
	if err != nil {
		return string(b)
	}
	return string(out)
}