probe/model/Discovered.go

package model

import (
	"encoding/json"
	"fmt"

	omd "git.loafle.net/overflow/model/discovery"
)

type Discovered interface {
	Zone() *omd.Zone
	AddHost(host *omd.Host) *omd.Host
	AddPort(port *omd.Port) *omd.Port
	AddService(service *omd.Service) *omd.Service
}

func New(zone *omd.Zone) Discovered {
	return &discovered{
		zone:     zone,
		hosts:    make(map[string]*omd.Host),
		ports:    make(map[*omd.Host]map[json.Number]map[string]*omd.Port),
		services: make(map[*omd.Port]map[string]map[string]*omd.Service),
	}
}

type discovered struct {
	zone     *omd.Zone
	hosts    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
}

func (d *discovered) Zone() *omd.Zone {
	return d.zone
}

func (d *discovered) AddHost(host *omd.Host) *omd.Host {
	h := d.findHost(host, true)

	if "" == h.Mac && "" != host.Mac {
		h.Mac = host.Mac
	}

	h.Meta = d.appendMeta(h.Meta, host.Meta)

	return h
}

func (d *discovered) AddPort(port *omd.Port) *omd.Port {
	p := d.findPort(port, true)

	p.Meta = d.appendMeta(p.Meta, port.Meta)

	return p
}

func (d *discovered) AddService(service *omd.Service) *omd.Service {
	s := d.findService(service, true)

	s.Meta = d.appendMeta(s.Meta, service.Meta)

	return s
}

func (d *discovered) findHost(host *omd.Host, add bool) *omd.Host {
	h, ok := d.hosts[host.Address]
	if !ok {
		if !add {
			return nil
		}
		d.hosts[host.Address] = host
		h = host
	}

	return h
}

func (d *discovered) findPort(port *omd.Port, add bool) *omd.Port {
	h := d.findHost(port.Host, false)
	if nil == h {
		return nil
	}

	hostPorts, ok := d.ports[h]
	if !ok {
		if !add {
			return nil
		}

		d.ports[h] = make(map[json.Number]map[string]*omd.Port)
		hostPorts = d.ports[h]
	}

	ports, ok := hostPorts[port.PortNumber]
	if !ok {
		if !add {
			return nil
		}

		hostPorts[port.PortNumber] = make(map[string]*omd.Port)
		ports = hostPorts[port.PortNumber]
	}

	p, ok := ports[port.MetaPortType.Key]
	if !ok {
		if !add {
			return nil
		}

		ports[port.MetaPortType.Key] = port
		p = ports[port.MetaPortType.Key]
	}

	return p
}

func (d *discovered) findService(service *omd.Service, add bool) *omd.Service {
	p := d.findPort(service.Port, false)
	if nil == p {
		return nil
	}

	portServices, ok := d.services[p]
	if !ok {
		if !add {
			return nil
		}

		d.services[p] = make(map[string]map[string]*omd.Service)
		portServices = d.services[p]
	}

	services, ok := portServices[service.Key]
	if !ok {
		if !add {
			return nil
		}

		portServices[service.Key] = make(map[string]*omd.Service)
		services = portServices[service.Key]
	}

	s, ok := services[service.MetaCryptoType.Key]
	if !ok {
		if !add {
			return nil
		}

		services[service.MetaCryptoType.Key] = service
		s = services[service.MetaCryptoType.Key]
	}

	return s
}

func (d *discovered) appendMeta(oriMeta map[string]string, newMeta map[string]string) map[string]string {
	if nil == newMeta {
		return oriMeta
	}
	if nil == oriMeta {
		return newMeta
	}

LOOP:
	for k, v := range oriMeta {
		_v, _ok := oriMeta[k]
		if !_ok {
			oriMeta[k] = v
			continue LOOP
		}
		if v == _v {
			continue LOOP
		}
		oriMeta[k] = fmt.Sprintf("%s|||%s", _v, v)
	}

	return oriMeta
}