142 lines
2.1 KiB
Go
142 lines
2.1 KiB
Go
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
|
|
}
|