misc: small fix or general refactoring i did not bother commenting

modules/zerogod/zerogod_browser.go

@@ -0,0 +1,119 @@
+package zerogod
+
+import (
+	"context"
+	"sort"
+
+	"github.com/bettercap/bettercap/v2/zeroconf"
+	"github.com/evilsocket/islazy/tui"
+)
+
+const DNSSD_DISCOVERY_SERVICE = "_services._dns-sd._udp"
+
+type AddressServices struct {
+	Address  string
+	Services []*zeroconf.ServiceEntry
+}
+
+type Browser struct {
+	resolvers    map[string]*zeroconf.Resolver
+	servicesByIP map[string]map[string]*zeroconf.ServiceEntry
+	context      context.Context
+	cancel       context.CancelFunc
+}
+
+func NewBrowser() *Browser {
+	servicesByIP := make(map[string]map[string]*zeroconf.ServiceEntry)
+	resolvers := make(map[string]*zeroconf.Resolver)
+	context, cancel := context.WithCancel(context.Background())
+	return &Browser{
+		resolvers:    resolvers,
+		servicesByIP: servicesByIP,
+		context:      context,
+		cancel:       cancel,
+	}
+}
+
+func (b *Browser) Wait() {
+	<-b.context.Done()
+}
+
+func (b *Browser) Stop(wait bool) {
+	b.cancel()
+	if wait {
+		b.Wait()
+	}
+}
+
+func (b *Browser) HasResolverFor(service string) bool {
+	_, found := b.resolvers[service]
+	return found
+}
+
+func (b *Browser) AddServiceFor(ip string, svc *zeroconf.ServiceEntry) {
+	if ipServices, found := b.servicesByIP[ip]; found {
+		ipServices[svc.ServiceInstanceName()] = svc
+	} else {
+		b.servicesByIP[ip] = map[string]*zeroconf.ServiceEntry{
+			svc.ServiceInstanceName(): svc,
+		}
+	}
+}
+
+func (b *Browser) GetServicesFor(ip string) map[string]*zeroconf.ServiceEntry {
+	if ipServices, found := b.servicesByIP[ip]; found {
+		return ipServices
+	}
+	return nil
+}
+
+func (b *Browser) StartBrowsing(service string, domain string, mod *ZeroGod) (chan *zeroconf.ServiceEntry, error) {
+	resolver, err := zeroconf.NewResolver(nil)
+	if err != nil {
+		return nil, err
+	}
+
+	b.resolvers[service] = resolver
+	ch := make(chan *zeroconf.ServiceEntry)
+
+	// start browsing
+	go func() {
+		if err := resolver.Browse(b.context, service, domain, ch); err != nil {
+			mod.Error("%v", err)
+		}
+		mod.Debug("resolver for service %s stopped", tui.Yellow(service))
+	}()
+
+	return ch, nil
+}
+
+func (b *Browser) ServicesByAddress(filter string) []AddressServices {
+	// convert to list for sorting
+	entries := make([]AddressServices, 0)
+
+	for ip, services := range b.servicesByIP {
+		if filter == "" || ip == filter {
+			// collect and sort services by name
+			svcList := make([]*zeroconf.ServiceEntry, 0)
+			for _, svc := range services {
+				svcList = append(svcList, svc)
+			}
+
+			sort.Slice(svcList, func(i, j int) bool {
+				return svcList[i].ServiceInstanceName() < svcList[j].ServiceInstanceName()
+			})
+
+			entries = append(entries, AddressServices{
+				Address:  ip,
+				Services: svcList,
+			})
+		}
+	}
+
+	// sort entries by ip
+	sort.Slice(entries, func(i, j int) bool {
+		return entries[i].Address < entries[j].Address
+	})
+
+	return entries
+}