star/net/scanport.go

package net

import (
	"b612.me/apps/b612/netforward"
	"b612.me/stario"
	"b612.me/starlog"
	"fmt"
	"net"
	"sort"
	"strconv"
	"strings"
	"sync/atomic"
	"time"
)

type ScanPort struct {
	Host    string
	Ports   []int
	Timeout int
	Threads int
	Log     string
	Retry   int
}

func (s *ScanPort) Parse(potStr string) error {
	ports := strings.Split(potStr, ",")
	for _, port := range ports {
		port = strings.TrimSpace(port)
		if strings.Contains(port, "-") {
			// range
			r := strings.Split(port, "-")
			if len(r) != 2 {
				continue
			}
			start, err := strconv.Atoi(r[0])
			if err != nil {
				starlog.Warningf("invalid port: %s\n", r[0])
				continue
			}
			end, err := strconv.Atoi(r[1])
			if err != nil {
				starlog.Warningf("invalid port: %s\n", r[1])
				continue
			}
			for i := start; i <= end; i++ {
				if i < 1 || i > 65535 {
					starlog.Warningf("invalid port: %d\n", i)
					continue
				}
				s.Ports = append(s.Ports, i)
			}
		} else {
			// single port
			tmp, err := strconv.Atoi(port)
			if err != nil {
				starlog.Warningf("invalid port: %s\n", port)
				continue
			}
			if tmp < 1 || tmp > 65535 {
				starlog.Warningf("invalid port: %d\n", tmp)
				continue
			}
			s.Ports = append(s.Ports, tmp)
		}
	}
	return nil
}

func (s *ScanPort) Run() error {
	if s.Threads < 1 {
		s.Threads = 1
	}
	if s.Log != "" {
		starlog.SetLogFile(s.Log, starlog.Std, true)
	}
	sort.Ints(s.Ports)
	starlog.Infof("scan count %d ports for host %v\n", len(s.Ports), s.Host)
	wg := stario.NewWaitGroup(s.Threads)
	localAddr, err := net.ResolveTCPAddr("tcp", ":0")
	if err != nil {
		starlog.Errorln("ResolveTCPAddr error, ", err)
		return err
	}
	count := int32(0)
	allcount := int32(0)
	interrupt := make(chan int)
	go func() {
		for {
			select {
			case <-time.After(time.Second * 2):
				fmt.Printf("scan %d ports, %d open\r", atomic.LoadInt32(&allcount), count)
			case port, opened := <-interrupt:
				if !opened {
					return
				}
				starlog.Infof("port %d is open\n", port)
			}

		}
	}()
	for _, port := range s.Ports {
		wg.Add(1)
		go func(port int) {
			defer wg.Done()
			defer func() {
				atomic.AddInt32(&allcount, 1)
			}()
			for i := 0; i < s.Retry+1; i++ {
				dialer := net.Dialer{
					LocalAddr: localAddr,
					Timeout:   time.Duration(s.Timeout) * time.Millisecond,
					Control:   netforward.ControlSetReUseAddr,
				}
				conn, err := dialer.Dial("tcp", net.JoinHostPort(s.Host, strconv.Itoa(port)))
				if err != nil {
					continue
				}
				conn.(*net.TCPConn).SetLinger(0)
				conn.Close()
				interrupt <- port
				atomic.AddInt32(&count, 1)
				return
			}
		}(port)

	}
	wg.Wait()
	close(interrupt)
	starlog.Infof("scan %d ports, %d open\n", len(s.Ports), count)
	return nil
}