实现Ping和TCP端口扫描功能

2026-07-17 19:00:47 +08:00
parent 4a2086c6f1
commit c858e8e311
6 changed files with 516 additions and 23 deletions
@@ -0,0 +1,141 @@
+package scan
+
+import (
+	"strconv"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+type Result struct {
+	Text string
+}
+
+type Progress struct {
+	Done  int
+	Total int
+}
+
+type Dispatcher struct {
+	IPs      []string
+	Ports    []int
+	Timeout  time.Duration
+	Workers  int
+	PingOnly bool
+
+	OnResult   func(Result)
+	OnProgress func(Progress)
+	OnFinish   func()
+
+	stop    chan struct{}
+	done    sync.WaitGroup
+	counter int64
+}
+
+func (d *Dispatcher) Start() {
+	d.stop = make(chan struct{})
+	tasks := d.buildTasks()
+	total := len(tasks)
+	if total == 0 {
+		if d.OnFinish != nil {
+			d.OnFinish()
+		}
+		return
+	}
+
+	workers := d.Workers
+	if workers <= 0 {
+		workers = 1
+	}
+
+	taskCh := make(chan task)
+	for i := 0; i < workers; i++ {
+		d.done.Add(1)
+		go d.worker(taskCh, total)
+	}
+
+feed:
+	for _, t := range tasks {
+		select {
+		case <-d.stop:
+			break feed
+		case taskCh <- t:
+		}
+	}
+	close(taskCh)
+	d.done.Wait()
+	if d.OnFinish != nil {
+		d.OnFinish()
+	}
+}
+
+func (d *Dispatcher) worker(taskCh <-chan task, total int) {
+	defer d.done.Done()
+
+	var pinger *Pinger
+	if d.PingOnly {
+		pinger, _ = NewPinger()
+		defer pinger.Close()
+	}
+
+	for task := range taskCh {
+		select {
+		case <-d.stop:
+			return
+		default:
+		}
+
+		ok := false
+		if task.ping {
+			if pinger != nil {
+				ok = pinger.Ping(task.ip, d.Timeout)
+			}
+		} else {
+			ok = TCPProbe(task.ip, task.port, d.Timeout)
+		}
+
+		if ok && d.OnResult != nil {
+			d.OnResult(Result{Text: task.label})
+		}
+
+		done := int(atomic.AddInt64(&d.counter, 1))
+		if d.OnProgress != nil {
+			d.OnProgress(Progress{Done: done, Total: total})
+		}
+	}
+}
+
+func (d *Dispatcher) Stop() {
+	select {
+	case <-d.stop:
+	default:
+		close(d.stop)
+	}
+}
+
+type task struct {
+	ip    string
+	port  int
+	label string
+	ping  bool
+}
+
+func (d *Dispatcher) buildTasks() []task {
+	var tasks []task
+	if d.PingOnly {
+		for _, ip := range d.IPs {
+			tasks = append(tasks, task{ip: ip, label: ip, ping: true})
+		}
+	} else {
+		for _, ip := range d.IPs {
+			for _, port := range d.Ports {
+				tasks = append(tasks, task{
+					ip:    ip,
+					port:  port,
+					label: ip + ":" + strconv.Itoa(port),
+				})
+			}
+		}
+	}
+	return tasks
+}