123 lines
2.5 KiB
Go
123 lines
2.5 KiB
Go
package connection
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net"
|
|
"os/exec"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
omd "git.loafle.net/overflow/model/discovery"
|
|
omm "git.loafle.net/overflow/model/meta"
|
|
omu "git.loafle.net/overflow/model/util"
|
|
"git.loafle.net/overflow_scanner/probe/discovery/session"
|
|
"golang.org/x/sync/semaphore"
|
|
)
|
|
|
|
const (
|
|
defaultUlimit = 1024
|
|
)
|
|
|
|
func Scan(discoverySession session.DiscoverySession, targetHost *omd.Host) error {
|
|
if nil == targetHost || nil == discoverySession.DiscoverPort() || !discoverySession.DiscoverPort().IncludeTCP {
|
|
return nil
|
|
}
|
|
|
|
dp := discoverySession.DiscoverPort()
|
|
|
|
lock := semaphore.NewWeighted(Ulimit())
|
|
var wg sync.WaitGroup
|
|
stopChan := make(chan struct{})
|
|
timeout := 500 * time.Millisecond
|
|
ports := make(map[int]*omd.Port)
|
|
|
|
Loop:
|
|
for portNumber := dp.FirstScanRange; portNumber < dp.LastScanRange; portNumber++ {
|
|
if nil != dp.ExcludePorts {
|
|
for _, exPortNumber := range dp.ExcludePorts {
|
|
if portNumber == exPortNumber {
|
|
continue Loop
|
|
}
|
|
}
|
|
}
|
|
|
|
lock.Acquire(context.TODO(), 1)
|
|
wg.Add(1)
|
|
|
|
go func(port int) {
|
|
defer func() {
|
|
lock.Release(1)
|
|
wg.Done()
|
|
}()
|
|
|
|
tryConnect(discoverySession, ports, targetHost, port, timeout, stopChan)
|
|
}(portNumber)
|
|
|
|
select {
|
|
case <-time.After(time.Microsecond * 100):
|
|
case <-discoverySession.StopChan():
|
|
close(stopChan)
|
|
return nil
|
|
}
|
|
}
|
|
|
|
wg.Wait()
|
|
|
|
return nil
|
|
}
|
|
|
|
func tryConnect(discoverySession session.DiscoverySession, ports map[int]*omd.Port, targetHost *omd.Host, port int, timeout time.Duration, stopChan <-chan struct{}) {
|
|
addr := net.JoinHostPort(targetHost.Address, strconv.Itoa(port))
|
|
conn, err := net.DialTimeout("tcp", addr, timeout)
|
|
dp := discoverySession.DiscoverPort()
|
|
|
|
if err != nil {
|
|
if strings.Contains(err.Error(), "too many open files") {
|
|
select {
|
|
case <-stopChan:
|
|
return
|
|
default:
|
|
}
|
|
time.Sleep(timeout)
|
|
tryConnect(discoverySession, ports, targetHost, port, timeout, stopChan)
|
|
}
|
|
return
|
|
}
|
|
|
|
conn.Close()
|
|
|
|
if _, ok := ports[port]; ok || !dp.Contains(port) {
|
|
return
|
|
}
|
|
|
|
p := &omd.Port{
|
|
MetaPortType: omm.ToMetaPortType(omm.MetaPortTypeEnumTCP),
|
|
PortNumber: json.Number(strconv.Itoa(port)),
|
|
DiscoveredDate: omu.NowPtr(),
|
|
}
|
|
p.Host = targetHost
|
|
|
|
ports[port] = p
|
|
|
|
go discoverySession.AddPort(p)
|
|
}
|
|
|
|
func Ulimit() int64 {
|
|
out, err := exec.Command("ulimit", "-n").Output()
|
|
if nil != err {
|
|
return defaultUlimit
|
|
}
|
|
|
|
s := strings.TrimSpace(string(out))
|
|
|
|
i, err := strconv.ParseInt(s, 10, 64)
|
|
if nil != err {
|
|
return defaultUlimit
|
|
}
|
|
|
|
return i
|
|
}
|