package session import ( "encoding/json" "fmt" "net" "strconv" "sync" omd "git.loafle.net/overflow/model/discovery" ounc "git.loafle.net/overflow/util-go/net/cidr" "git.loafle.net/overflow_scanner/probe/discovery/types" ) type DiscoverySession interface { InitWithRequest(request types.DiscoveryRequest) error DiscoveryRequest() types.DiscoveryRequest Zone() *omd.Zone Host() *omd.Host Port() *omd.Port DiscoveryConfig() *omd.DiscoveryConfig DiscoverHost() *omd.DiscoverHost DiscoverPort() *omd.DiscoverPort DiscoverService() *omd.DiscoverService TargetHosts() []net.IP SetDiscoveryDelegator(chan<- interface{}) AddHost(host *omd.Host) *omd.Host AddPort(port *omd.Port) *omd.Port AddService(service *omd.Service) *omd.Service DiscoveredHost(address string) *omd.Host DiscoveredAllHosts(includeMac bool) map[string]*omd.Host DiscoveredPort(host *omd.Host, portNumber int) map[string]*omd.Port DiscoveredAllPorts() map[*omd.Host]map[json.Number]map[string]*omd.Port DiscoveredService(port *omd.Port, name string) map[string]*omd.Service DiscoveredAllServices() map[*omd.Port]map[string]map[string]*omd.Service StopChan() <-chan struct{} } type ofDiscoverySession struct { discoveryRequest types.DiscoveryRequest zone *omd.Zone host *omd.Host port *omd.Port discoveryConfig *omd.DiscoveryConfig discoverHost *omd.DiscoverHost discoverPort *omd.DiscoverPort discoverService *omd.DiscoverService targetHosts []net.IP discoveryDelegator chan<- interface{} discoveredMtx sync.RWMutex hosts map[string]*omd.Host includeMachosts map[string]*omd.Host ports map[*omd.Host]map[json.Number]map[string]*omd.Port services map[*omd.Port]map[string]map[string]*omd.Service stopChan chan struct{} } func (ds *ofDiscoverySession) init(request types.DiscoveryRequest) { ds.discoveryRequest = request ds.zone = nil ds.host = nil ds.port = nil ds.discoverHost = nil ds.discoverPort = nil ds.discoverService = nil ds.hosts = make(map[string]*omd.Host) ds.includeMachosts = make(map[string]*omd.Host) ds.ports = make(map[*omd.Host]map[json.Number]map[string]*omd.Port) ds.services = make(map[*omd.Port]map[string]map[string]*omd.Service) ds.stopChan = make(chan struct{}) } func (ds *ofDiscoverySession) InitWithRequest(request types.DiscoveryRequest) error { ds.init(request) params := request.Params() switch request.RequestType() { case types.DiscoveryRequestTypeHost: if nil == params || 2 != len(params) { return fmt.Errorf("Parameter is not valid") } zone, ok := params[0].(*omd.Zone) if !ok { return fmt.Errorf("Zone of parameter is not valid") } discoverHost, ok := params[1].(*omd.DiscoverHost) if !ok { return fmt.Errorf("DiscoverHost of parameter is not valid") } if nil == discoverHost.DiscoveryConfig { return fmt.Errorf("DiscoveryConfig of parameter is not valid") } if ts, err := ds.getTargetHosts(zone, discoverHost); nil != err { return err } else { ds.targetHosts = ts } ds.discoveryConfig = discoverHost.DiscoveryConfig ds.setZone(zone) ds.setDiscoverHost(discoverHost) case types.DiscoveryRequestTypePort: if nil == params || 2 != len(params) { return fmt.Errorf("Parameter is not valid") } host, ok := params[0].(*omd.Host) if !ok { return fmt.Errorf("Host of parameter is not valid") } discoverPort, ok := params[1].(*omd.DiscoverPort) if !ok { return fmt.Errorf("DiscoverPort of parameter is not valid") } if nil == discoverPort.DiscoveryConfig { return fmt.Errorf("DiscoveryConfig of parameter is not valid") } ds.discoveryConfig = discoverPort.DiscoveryConfig ds.setHost(host) ds.setDiscoverPort(discoverPort) case types.DiscoveryRequestTypeService: if nil == params || 2 != len(params) { return fmt.Errorf("Parameter is not valid") } port, ok := params[0].(*omd.Port) if !ok { return fmt.Errorf("Port of parameter is not valid") } discoverService, ok := params[1].(*omd.DiscoverService) if !ok { return fmt.Errorf("DiscoverService of parameter is not valid") } if nil == discoverService.DiscoveryConfig { return fmt.Errorf("DiscoveryConfig of parameter is not valid") } ds.discoveryConfig = discoverService.DiscoveryConfig ds.setPort(port) ds.setDiscoverService(discoverService) } return nil } func (ds *ofDiscoverySession) SetDiscoveryDelegator(discoveryDelegator chan<- interface{}) { ds.discoveryDelegator = discoveryDelegator } func (ds *ofDiscoverySession) DiscoveryRequest() types.DiscoveryRequest { return ds.discoveryRequest } func (ds *ofDiscoverySession) Zone() *omd.Zone { return ds.zone } func (ds *ofDiscoverySession) Host() *omd.Host { return ds.host } func (ds *ofDiscoverySession) Port() *omd.Port { return ds.port } func (ds *ofDiscoverySession) DiscoveryConfig() *omd.DiscoveryConfig { return ds.discoveryConfig } func (ds *ofDiscoverySession) DiscoverHost() *omd.DiscoverHost { return ds.discoverHost } func (ds *ofDiscoverySession) DiscoverPort() *omd.DiscoverPort { return ds.discoverPort } func (ds *ofDiscoverySession) DiscoverService() *omd.DiscoverService { return ds.discoverService } func (ds *ofDiscoverySession) TargetHosts() []net.IP { return ds.targetHosts } func (ds *ofDiscoverySession) AddHost(host *omd.Host) *omd.Host { ds.discoveredMtx.Lock() defer ds.discoveredMtx.Unlock() h, modified := ds.findHost(host) if "" == h.Mac && "" != host.Mac { h.Mac = host.Mac modified = true } meta, metaModified := ds.appendMeta(h.Meta, host.Meta) if metaModified { h.Meta = meta modified = metaModified } if modified && nil != ds.discoveryDelegator { ds.discoveryDelegator <- h } return h } func (ds *ofDiscoverySession) AddPort(port *omd.Port) *omd.Port { ds.discoveredMtx.Lock() defer ds.discoveredMtx.Unlock() p, modified := ds.findPort(port) meta, metaModified := ds.appendMeta(p.Meta, port.Meta) if metaModified { p.Meta = meta modified = metaModified } if modified && nil != ds.discoveryDelegator { ds.discoveryDelegator <- p } return p } func (ds *ofDiscoverySession) AddService(service *omd.Service) *omd.Service { ds.discoveredMtx.Lock() defer ds.discoveredMtx.Unlock() s, modified := ds.findService(service) meta, metaModified := ds.appendMeta(s.Meta, service.Meta) if metaModified { s.Meta = meta modified = metaModified } if modified && nil != ds.discoveryDelegator { ds.discoveryDelegator <- s } return s } func (ds *ofDiscoverySession) DiscoveredHost(address string) *omd.Host { ds.discoveredMtx.RLock() defer ds.discoveredMtx.RUnlock() h, ok := ds.hosts[address] if !ok { return nil } return h } func (ds *ofDiscoverySession) DiscoveredAllHosts(includeMac bool) map[string]*omd.Host { if includeMac { return ds.includeMachosts } return ds.hosts } func (ds *ofDiscoverySession) DiscoveredPort(host *omd.Host, portNumber int) map[string]*omd.Port { ds.discoveredMtx.RLock() defer ds.discoveredMtx.RUnlock() h, _ := ds.findHost(host) if nil == h { return nil } hostPorts, ok := ds.ports[h] if !ok { return nil } ports, ok := hostPorts[json.Number(strconv.Itoa(portNumber))] if !ok { return nil } return ports } func (ds *ofDiscoverySession) DiscoveredAllPorts() map[*omd.Host]map[json.Number]map[string]*omd.Port { return ds.ports } func (ds *ofDiscoverySession) DiscoveredService(port *omd.Port, name string) map[string]*omd.Service { ds.discoveredMtx.RLock() defer ds.discoveredMtx.RUnlock() p, _ := ds.findPort(port) if nil == p { return nil } portServices, ok := ds.services[p] if !ok { return nil } services, ok := portServices[name] if !ok { return nil } return services } func (ds *ofDiscoverySession) DiscoveredAllServices() map[*omd.Port]map[string]map[string]*omd.Service { return ds.services } func (ds *ofDiscoverySession) findHost(host *omd.Host) (h *omd.Host, modified bool) { modified = false var ok bool h, ok = ds.hosts[host.Address] if !ok { ds.hosts[host.Address] = host if "" != host.Mac { ds.includeMachosts[host.Address] = host } h = host modified = true } return } func (ds *ofDiscoverySession) StopChan() <-chan struct{} { return ds.stopChan } func (ds *ofDiscoverySession) findPort(port *omd.Port) (p *omd.Port, modified bool) { modified = false var ok bool h, _ := ds.findHost(port.Host) if nil == h { return } hostPorts, ok := ds.ports[h] if !ok { ds.ports[h] = make(map[json.Number]map[string]*omd.Port) hostPorts = ds.ports[h] } ports, ok := hostPorts[port.PortNumber] if !ok { hostPorts[port.PortNumber] = make(map[string]*omd.Port) ports = hostPorts[port.PortNumber] } p, ok = ports[port.MetaPortType.Key] if !ok { ports[port.MetaPortType.Key] = port p = ports[port.MetaPortType.Key] modified = true } return } func (ds *ofDiscoverySession) findService(service *omd.Service) (s *omd.Service, modified bool) { modified = false var ok bool p, _ := ds.findPort(service.Port) if nil == p { return } portServices, ok := ds.services[p] if !ok { ds.services[p] = make(map[string]map[string]*omd.Service) portServices = ds.services[p] } services, ok := portServices[service.Key] if !ok { portServices[service.Key] = make(map[string]*omd.Service) services = portServices[service.Key] } s, ok = services[service.MetaCryptoType.Key] if !ok { services[service.MetaCryptoType.Key] = service s = services[service.MetaCryptoType.Key] modified = true } return } func (ds *ofDiscoverySession) appendMeta(oriMeta map[string]string, newMeta map[string]string) (resultMap map[string]string, modified bool) { modified = false if nil == newMeta { resultMap = oriMeta return } if nil == oriMeta { resultMap = newMeta modified = true return } LOOP: for k, v := range newMeta { _v, _ok := oriMeta[k] if !_ok { oriMeta[k] = v modified = true continue LOOP } if v == _v { continue LOOP } oriMeta[k] = fmt.Sprintf("%s|||%s", _v, v) modified = true } resultMap = oriMeta return } func (ds *ofDiscoverySession) setZone(zone *omd.Zone) { if nil == zone { return } ds.zone = zone } func (ds *ofDiscoverySession) setHost(host *omd.Host) { if nil == host { return } ds.setZone(host.Zone) ds.host = host } func (ds *ofDiscoverySession) setPort(port *omd.Port) { if nil == port { return } ds.setHost(port.Host) ds.port = port } func (ds *ofDiscoverySession) setDiscoverHost(discoverHost *omd.DiscoverHost) { if nil == discoverHost { return } ds.discoverHost = discoverHost ds.setDiscoverPort(discoverHost.DiscoverPort) } func (ds *ofDiscoverySession) setDiscoverPort(discoverPort *omd.DiscoverPort) { if nil == discoverPort { return } ds.discoverPort = discoverPort ds.setDiscoverService(discoverPort.DiscoverService) } func (ds *ofDiscoverySession) setDiscoverService(discoverService *omd.DiscoverService) { if nil == discoverService { return } ds.discoverService = discoverService } func (ds *ofDiscoverySession) getTargetHosts(zone *omd.Zone, discoverHost *omd.DiscoverHost) ([]net.IP, error) { cr, err := ounc.NewCIDRRanger(zone.Network) if nil != err { return nil, err } var firstIP net.IP if "" != discoverHost.FirstScanRange { firstIP = net.ParseIP(discoverHost.FirstScanRange) if nil == firstIP { return nil, fmt.Errorf("IP(%v) of FirstScanRange host is not valid", firstIP) } } var lastIP net.IP if "" != discoverHost.LastScanRange { lastIP = net.ParseIP(discoverHost.LastScanRange) if nil == lastIP { return nil, fmt.Errorf("IP(%v) of LastScanRange host is not valid", lastIP) } } includeIPs := make([]net.IP, 0) for _, iHost := range discoverHost.IncludeHosts { iIP := net.ParseIP(iHost) if nil == iIP { return nil, fmt.Errorf("IP(%v) of include host is not valid", iHost) } includeIPs = append(includeIPs, iIP) } excludeIPs := make([]net.IP, 0) for _, eHost := range discoverHost.ExcludeHosts { eIP := net.ParseIP(eHost) if nil == eIP { return nil, fmt.Errorf("IP(%v) of exclude host is not valid", eHost) } excludeIPs = append(excludeIPs, eIP) } ranges, err := cr.Ranges(firstIP, lastIP, includeIPs, excludeIPs) if nil != err { return nil, err } return ranges, nil } var discoverySessionPool sync.Pool func RetainDiscoverySession() *ofDiscoverySession { v := discoverySessionPool.Get() var ds *ofDiscoverySession if v == nil { ds = &ofDiscoverySession{} } else { ds = v.(*ofDiscoverySession) } return ds } func ReleaseDiscoverySession(ds *ofDiscoverySession) { close(ds.stopChan) ds.discoveryRequest = nil ds.zone = nil ds.host = nil ds.port = nil ds.discoverHost = nil ds.discoverPort = nil ds.discoverService = nil ds.includeMachosts = nil ds.hosts = nil ds.ports = nil ds.services = nil discoverySessionPool.Put(ds) }