probe/discoverer/discoverer_test.go

package discoverer_test

import (
	"encoding/json"
	"log"
	"strconv"
	"testing"

	logging "git.loafle.net/commons/logging-go"
	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/discoverer"
)

var (
	dz = &omd.DiscoverZone{
		ExcludePatterns: []string{},
	}
	z = &omd.Zone{
		Network:    "192.168.1.0/24",
		MetaIPType: omm.ToMetaIPType(omm.MetaIPTypeEnumV4),
		Address:    "192.168.1.201",
		Iface:      " \\Device\\NPF_{1924FA2B-6927-4BA5-AF43-876C3F8853CE}",
		Mac:        "30:9C:23:15:A3:09",
	}
	dh = &omd.DiscoverHost{
		MetaIPType:     omm.ToMetaIPType(omm.MetaIPTypeEnumV4),
		FirstScanRange: "192.168.1.1",
		LastScanRange:  "192.168.1.254",
	}
	h = &omd.Host{
		Zone:       z,
		MetaIPType: omm.ToMetaIPType(omm.MetaIPTypeEnumV4),
		Address:    "192.168.1.229",
		Mac:        "50:E5:49:46:93:28",
	}
	dp = &omd.DiscoverPort{
		FirstScanRange: 1,
		LastScanRange:  50000,
		ExcludePorts: []int{
			631,
		},
		IncludeTCP: true,
		IncludeUDP: true,
	}
	p = &omd.Port{
		Host:           h,
		MetaPortType:   omm.ToMetaPortType(omm.MetaPortTypeEnumTCP),
		PortNumber:     json.Number(strconv.Itoa(21)),
		DiscoveredDate: omu.NowPtr(),
	}
	ds = &omd.DiscoverService{
		IncludeServices: []string{
			"ACTIVEDIRECTORY",
			"CASSANDRA",
			"DHCP",
			"DNS",
			"FTP",
			"HTTP",
			"IMAP",
			"LDAP",
			"MONGODB",
			"MSSQL",
			"MYSQL",
			"NETBIOS",
			"ORACLE",
			"POP",
			"POSTGRESQL",
			"REDIS",
			"RMI",
			"SMB",
			"SMTP",
			"SNMP",
			"SSH",
			"TELNET",
			"WMI",
			"UNKNOWN",
			"SSH",
			"WMI",
			"SNMP V3",
		},
	}

	dha = &omd.DiscoverHost{
		MetaIPType:     omm.ToMetaIPType(omm.MetaIPTypeEnumV4),
		FirstScanRange: "192.168.1.0",
		LastScanRange:  "192.168.1.255",
		DiscoverPort: &omd.DiscoverPort{
			FirstScanRange: 1,
			LastScanRange:  50000,
			ExcludePorts: []int{
				631,
			},
			IncludeTCP:      true,
			IncludeUDP:      false,
			DiscoverService: ds,
		},
	}
)

func init() {
	// runtime.LockOSThread()
	logging.InitializeLogger("")
}

func TestDiscoverZone(t *testing.T) {
	_discoverer := discoverer.GetDiscoverer()

	handleDiscovery(_discoverer, func(dataChan chan *discoverer.DiscoveryData) {
		_discoverer.DiscoverZone(dataChan, dz)
	})
}

func TestDiscoverHost(t *testing.T) {
	_discoverer := discoverer.GetDiscoverer()

	handleDiscovery(_discoverer, func(dataChan chan *discoverer.DiscoveryData) {
		_discoverer.DiscoverHost(dataChan, z, dh)
	})
}

func TestDiscoverPort(t *testing.T) {
	_discoverer := discoverer.GetDiscoverer()

	handleDiscovery(_discoverer, func(dataChan chan *discoverer.DiscoveryData) {
		_discoverer.DiscoverPort(dataChan, h, dp)
	})
}

func TestDiscoverService(t *testing.T) {
	_discoverer := discoverer.GetDiscoverer()

	handleDiscovery(_discoverer, func(dataChan chan *discoverer.DiscoveryData) {
		_discoverer.DiscoverSerice(dataChan, p, ds)
	})
}

func TestDiscoverHostAll(t *testing.T) {
	_discoverer := discoverer.GetDiscoverer()

	handleDiscovery(_discoverer, func(dataChan chan *discoverer.DiscoveryData) {
		_discoverer.DiscoverHost(dataChan, z, dha)
	})
}

func handleDiscovery(_discoverer discoverer.Discoverer, discoveryFunc func(dataChan chan *discoverer.DiscoveryData)) error {
	var dataChan chan *discoverer.DiscoveryData
	retainChan := make(chan struct{})

	go func() {
		dataChan = _discoverer.Retain()
		close(retainChan)
	}()

	select {
	case <-retainChan:
	}
	defer func() {
		_discoverer.Release(dataChan)
	}()

	go discoveryFunc(dataChan)

	for {
		select {
		case data, ok := <-dataChan:
			if !ok {
				return nil
			}
			switch data.Type {
			case discoverer.DiscoveryDataTypeStart:
				log.Printf("DiscoveryService.DiscoveryStart")
			case discoverer.DiscoveryDataTypeStop:
				log.Printf("DiscoveryService.DiscoveryStop")
				data.Release()
				return nil
			case discoverer.DiscoveryDataTypeError:
				log.Printf("DiscoveryService.DiscoveryDataTypeError %v", data.Error)
			case discoverer.DiscoveryDataTypeZone:
				log.Printf("DiscoveryService.DiscoveryDataTypeZone %v", data.Result)
			case discoverer.DiscoveryDataTypeHost:
				log.Printf("DiscoveryService.DiscoveryDataTypeHost %v", data.Result)
			case discoverer.DiscoveryDataTypePort:
				log.Printf("DiscoveryService.DiscoveryDataTypePort %v", data.Result)
			case discoverer.DiscoveryDataTypeService:
				// log.Printf("DiscoveryService.DiscoveryDataTypeService %v", data.Result)
				log.Println(data.Result)
			}
			data.Release()
		}
	}
}