container_discovery/internal/discoverer/discoverer_test.go

package discoverer_test

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

	logging "git.loafle.net/commons/logging-go"
	occc "git.loafle.net/overflow/commons-go/core/constants"
	occu "git.loafle.net/overflow/commons-go/core/util"
	ocmd "git.loafle.net/overflow/commons-go/model/discovery"
	"git.loafle.net/overflow/container_discovery/internal/discoverer"
)

var (
	dz = &ocmd.DiscoverZone{
		ExcludePatterns: []string{},
	}
	z = &ocmd.Zone{
		Network: "192.168.1.0/24",
		IPV4:    "192.168.1.101",
		Iface:   "enp3s0",
		Mac:     "44:8a:5b:f1:f1:f3",
	}
	dh = &ocmd.DiscoverHost{
		FirstScanRangeV4: "192.168.1.1",
		LastScanRangeV4:  "192.168.1.254",
	}
	h = &ocmd.Host{
		Zone: z,
		IPV4: "192.168.1.10",
		Mac:  "d0:50:99:97:5d:99",
	}
	dp = &ocmd.DiscoverPort{
		FirstScanRange: 1,
		LastScanRange:  50000,
		ExcludePorts: []int{
			631,
		},
		IncludeTCP: true,
		IncludeUDP: true,
	}
	p = &ocmd.Port{
		Host:           h,
		PortType:       occc.PortTypeTCP,
		PortNumber:     json.Number(strconv.Itoa(80)),
		DiscoveredDate: occu.NowPtr(),
	}
	ds = &ocmd.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 = &ocmd.DiscoverHost{
		FirstScanRangeV4: "192.168.1.0",
		LastScanRangeV4:  "192.168.1.255",
		DiscoverPort: &ocmd.DiscoverPort{
			FirstScanRange: 1,
			LastScanRange:  50000,
			ExcludePorts: []int{
				631,
			},
			IncludeTCP:      true,
			IncludeUDP:      false,
			DiscoverService: ds,
		},
	}
)

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

// func TestMain(m *testing.M) {

// 	go func() {
// 		os.Exit(m.Run())
// 	}()
// }

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)
			}
			data.Release()
		}
	}
}