probe/discovery/protocol/mdns/mdns.go

package mdns

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"strconv"
	"strings"
	"sync/atomic"
	"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"
	"github.com/grandcat/zeroconf"
)

func Scan(discoverySession session.DiscoverySession) error {

	serviceEntries, err := browse("_services._dns-sd._udp", "local")
	if nil != err {
		log.Print("Cannot find service ", err)
		return nil
	}

	metaIPTypeEnum := omm.ToMetaIPTypeEnum(discoverySession.Zone().MetaIPType)

SERVICE_LOOP:
	for _, serviceEntry := range serviceEntries {
		name := removeDomainName(serviceEntry.Instance, serviceEntry.Domain)
		entries, _err := browse(name, serviceEntry.Domain)
		if nil != _err {
			log.Print("Cannot find entry ", _err)
			continue SERVICE_LOOP
		}

	ENTRY_LOOP:
		for _, entry := range entries {
			// log.Print("serviceEntry ", entry)

			name := entry.Instance   // HP\ LaserJet\ P1505n
			service := entry.Service // _pdl-datastream._tcp
			port := entry.Port       // 9100
			meta := toMeta(entry.Text)
			hostName := removeDomainName(entry.HostName, entry.Domain) // NPIFACA9B
			serviceName, portType, cryptoType := parseService(service)

			var metaPortType *omm.MetaPortType
			switch portType {
			case "tcp":
				metaPortType = omm.ToMetaPortType(omm.MetaPortTypeEnumTCP)
			case "udp":
				metaPortType = omm.ToMetaPortType(omm.MetaPortTypeEnumUDP)
			}

			metaCryptoType := omm.ToMetaCryptoType(omm.MetaCryptoTypeEnumUNKNOWN)
			switch cryptoType {
			case "NONE":
				metaCryptoType = omm.ToMetaCryptoType(omm.MetaCryptoTypeEnumNONE)
			case "TLS":
				metaCryptoType = omm.ToMetaCryptoType(omm.MetaCryptoTypeEnumTLS)
			}

			switch metaIPTypeEnum {
			case omm.MetaIPTypeEnumV4:
				for _, ipv4 := range entry.AddrIPv4 {
					h := discoverySession.AddHost(
						"mDNS",
						&omd.Host{
							MetaIPType:     omm.ToMetaIPType(metaIPTypeEnum),
							Name:           hostName,
							Address:        ipv4.String(),
							Zone:           discoverySession.Zone(),
							DiscoveredDate: omu.NowPtr(),
						},
						meta,
					)

					if 1 > port {
						continue ENTRY_LOOP
					}

					p := discoverySession.AddPort(
						"mDNS",
						&omd.Port{
							MetaPortType:   metaPortType,
							PortNumber:     json.Number(strconv.Itoa(port)),
							Host:           h,
							DiscoveredDate: omu.NowPtr(),
						},
						meta,
					)

					discoverySession.AddService(
						"mDNS",
						&omd.Service{
							MetaCryptoType: metaCryptoType,
							Key:            serviceName,
							Name:           name,
							Port:           p,
							DiscoveredDate: omu.NowPtr(),
						},
						meta,
					)
				}

			case omm.MetaIPTypeEnumV6:
				for _, ipv6 := range entry.AddrIPv6 {
					h := discoverySession.AddHost(
						"mDNS",
						&omd.Host{
							MetaIPType:     omm.ToMetaIPType(metaIPTypeEnum),
							Name:           hostName,
							Address:        ipv6.String(),
							Zone:           discoverySession.Zone(),
							DiscoveredDate: omu.NowPtr(),
						},
						meta,
					)

					if 1 > port {
						continue ENTRY_LOOP
					}

					p := discoverySession.AddPort(
						"mDNS",
						&omd.Port{
							MetaPortType:   metaPortType,
							PortNumber:     json.Number(strconv.Itoa(port)),
							Host:           h,
							DiscoveredDate: omu.NowPtr(),
						},
						meta,
					)

					discoverySession.AddService(
						"mDNS",
						&omd.Service{
							MetaCryptoType: metaCryptoType,
							Key:            serviceName,
							Name:           name,
							Port:           p,
							DiscoveredDate: omu.NowPtr(),
						},
						meta,
					)
				}
			}

		}

		select {
		case <-discoverySession.StopChan():
			return nil
		default:
		}
	}

	return nil
}

func removeDomainName(instance string, domain string) (service string) {
	_instance := strings.TrimRight(instance, ".")
	return strings.TrimRight(_instance, fmt.Sprintf(".%s", domain))
}

func toMeta(text []string) map[string]string {
	if nil == text || 0 == len(text) {
		return nil
	}

	meta := make(map[string]string)
	for _, v := range text {
		ss := strings.SplitN(v, "=", 2)
		if 2 != len(ss) {
			continue
		}
		if nil != ss {
			meta[ss[0]] = ss[1]
		}
	}
	return meta
}

func parseService(service string) (name string, portType string, cryptoType string) {
	ss := strings.SplitN(service, ".", 2)
	if nil == ss {
		return "UNKNOWN", "UNKNOWN", "UNKNOWN"
	}

	_name := strings.TrimLeft(ss[0], "_")
	_portType := strings.TrimLeft(ss[1], "_")

	name = _name
	portType = _portType
	cryptoType = "NONE"
	return
}

func browse(service string, domain string) ([]*zeroconf.ServiceEntry, error) {
	resolver, err := zeroconf.NewResolver(nil)
	if err != nil {
		return nil, fmt.Errorf("Failed to initialize resolver %v", err)
	}

	entryChan := make(chan *zeroconf.ServiceEntry)
	timerStopped := make(chan struct{})

	serviceEntries := make([]*zeroconf.ServiceEntry, 0)

	go func(entryChan <-chan *zeroconf.ServiceEntry) {
		var delay atomic.Value
		delay.Store(false)
		ticker := time.NewTicker(time.Second * 1)

		for {
			select {
			case entry, ok := <-entryChan:
				if !ok {
					return
				}
				delay.Store(true)
				serviceEntries = append(serviceEntries, entry)
			case <-ticker.C:
				if false == delay.Load().(bool) {
					ticker.Stop()
					timerStopped <- struct{}{}
					return
				}
				delay.Store(false)
			}
		}
	}(entryChan)

	ctx, cancel := context.WithTimeout(context.Background(), time.Second*time.Duration(3))
	defer cancel()
	err = resolver.Browse(ctx, service, domain, entryChan)
	if err != nil {
		return nil, fmt.Errorf("Failed to browse %v", err)
	}

	select {
	case <-ctx.Done():
	case <-timerStopped:
	}

	return serviceEntries, nil
}