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

This commit is contained in:
Simone Margaritelli 2024-09-22 17:40:23 +02:00
commit 209725d623
5 changed files with 43 additions and 2 deletions

View file

@ -3,6 +3,7 @@ package zerogod
import (
"context"
"sort"
"sync"
"github.com/bettercap/bettercap/v2/modules/zerogod/zeroconf"
"github.com/evilsocket/islazy/tui"
@ -16,6 +17,8 @@ type AddressServices struct {
}
type Browser struct {
sync.RWMutex
resolvers map[string]*zeroconf.Resolver
servicesByIP map[string]map[string]*zeroconf.ServiceEntry
context context.Context
@ -46,11 +49,16 @@ func (b *Browser) Stop(wait bool) {
}
func (b *Browser) HasResolverFor(service string) bool {
b.RLock()
defer b.RUnlock()
_, found := b.resolvers[service]
return found
}
func (b *Browser) AddServiceFor(ip string, svc *zeroconf.ServiceEntry) {
b.Lock()
defer b.Unlock()
if ipServices, found := b.servicesByIP[ip]; found {
ipServices[svc.ServiceInstanceName()] = svc
} else {
@ -61,6 +69,9 @@ func (b *Browser) AddServiceFor(ip string, svc *zeroconf.ServiceEntry) {
}
func (b *Browser) GetServicesFor(ip string) map[string]*zeroconf.ServiceEntry {
b.RLock()
defer b.RUnlock()
if ipServices, found := b.servicesByIP[ip]; found {
return ipServices
}
@ -68,11 +79,15 @@ func (b *Browser) GetServicesFor(ip string) map[string]*zeroconf.ServiceEntry {
}
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.Lock()
defer b.Unlock()
b.resolvers[service] = resolver
ch := make(chan *zeroconf.ServiceEntry)
@ -88,6 +103,9 @@ func (b *Browser) StartBrowsing(service string, domain string, mod *ZeroGod) (ch
}
func (b *Browser) ServicesByAddress(filter string) []AddressServices {
b.RLock()
defer b.RUnlock()
// convert to list for sorting
entries := make([]AddressServices, 0)

View file

@ -254,7 +254,7 @@ func (mod *ZeroGod) onPacket(pkt gopacket.Packet) {
services := make([]string, 0)
for _, q := range dns.Questions {
services = append(services, tui.Yellow(string(q.Name)))
services = append(services, string(q.Name))
}
instances := make([]string, 0)