scan.go

   1 package main
   2
   3 import (
   4         "context"
   5         "log"
   6         "os"
   7         "os/signal"
   8         "sync"
   9         "sync/atomic"
  10         "time"
  11 )
  12
  13 type ScanRecord struct {
  14         IP  string
  15         RTT time.Duration
  16 }
  17
  18 type ScanRecords struct {
  19         recordMutex sync.RWMutex
  20         records     []*ScanRecord
  21         scanCounter int32
  22 }
  23
  24 func (srs *ScanRecords) AddRecord(rec *ScanRecord) {
  25         srs.recordMutex.Lock()
  26         srs.records = append(srs.records, rec)
  27         srs.recordMutex.Unlock()
  28         log.Printf("Found a record: IP=%s, RTT=%s\n", rec.IP, rec.RTT.String())
  29 }
  30
  31 func (srs *ScanRecords) IncScanCounter() {
  32         scanCount := atomic.AddInt32(&srs.scanCounter, 1)
  33         if scanCount%1000 == 0 {
  34                 log.Printf("Scanned %d IPs, Found %d records\n", scanCount, srs.RecordSize())
  35         }
  36 }
  37
  38 func (srs *ScanRecords) RecordSize() int {
  39         srs.recordMutex.RLock()
  40         defer srs.recordMutex.RUnlock()
  41         return len(srs.records)
  42 }
  43
  44 func (srs *ScanRecords) ScanCount() int32 {
  45         return atomic.LoadInt32(&srs.scanCounter)
  46 }
  47
  48 var testIPFunc func(ip string, config *ScanConfig, record *ScanRecord) bool
  49
  50 func testip(ip string, config *ScanConfig) *ScanRecord {
  51         record := new(ScanRecord)
  52         for i := 0; i < config.ScanCountPerIP; i++ {
  53                 if !testIPFunc(ip, config, record) {
  54                         return nil
  55                 }
  56         }
  57         record.IP = ip
  58         record.RTT = record.RTT / time.Duration(config.ScanCountPerIP)
  59         return record
  60 }
  61
  62 func testip_worker(ctx context.Context, ch chan string, gcfg *GScanConfig, cfg *ScanConfig, srs *ScanRecords, wg *sync.WaitGroup) {
  63         defer wg.Done()
  64
  65         timer := time.NewTimer(cfg.ScanMaxRTT + 100*time.Millisecond)
  66         defer timer.Stop()
  67
  68         ctx, cancal := context.WithCancel(ctx)
  69         defer cancal()
  70
  71         for ip := range ch {
  72                 srs.IncScanCounter()
  73
  74                 if gcfg.VerifyPing {
  75                         start := time.Now()
  76                         if err := Ping(ip, gcfg.ScanMaxPingRTT); err != nil {
  77                                 continue
  78                         }
  79                         if time.Since(start) < gcfg.ScanMinPingRTT {
  80                                 continue
  81                         }
  82                 }
  83
  84                 done := make(chan struct{}, 1)
  85                 go func() {
  86                         r := testip(ip, cfg)
  87                         if r != nil {
  88                                 if srs.RecordSize() >= cfg.RecordLimit {
  89                                         close(done)
  90                                         return
  91                                 }
  92                                 srs.AddRecord(r)
  93                         }
  94                         done <- struct{}{}
  95                 }()
  96
  97                 timer.Reset(cfg.ScanMaxRTT + 100*time.Millisecond)
  98                 select {
  99                 case <-ctx.Done():
 100                         return
 101                 case <-timer.C:
 102                         log.Println(ip, "timeout")
 103                 case <-done:
 104                 }
 105         }
 106 }
 107
 108 func StartScan(srs *ScanRecords, gcfg *GScanConfig, cfg *ScanConfig, ipqueue chan string) {
 109         var wg sync.WaitGroup
 110         wg.Add(gcfg.ScanWorker)
 111
 112         interrupt := make(chan os.Signal, 1)
 113         signal.Notify(interrupt, os.Interrupt)
 114
 115         ctx, cancel := context.WithCancel(context.Background())
 116         defer cancel()
 117
 118         go func() {
 119                 <-interrupt
 120                 cancel()
 121         }()
 122
 123         ch := make(chan string, 100)
 124         for i := 0; i < gcfg.ScanWorker; i++ {
 125                 go testip_worker(ctx, ch, gcfg, cfg, srs, &wg)
 126         }
 127
 128         for ip := range ipqueue {
 129                 select {
 130                 case ch <- ip:
 131                 case <-ctx.Done():
 132                         return
 133                 }
 134                 if srs.RecordSize() >= cfg.RecordLimit {
 135                         break
 136                 }
 137         }
 138
 139         close(ch)
 140         wg.Wait()
 141 }