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

modules/zerogod/zerogod.go

@@ -0,0 +1,157 @@
+package zerogod
+
+import (
+	"github.com/bettercap/bettercap/v2/session"
+	"github.com/bettercap/bettercap/v2/tls"
+)
+
+type ZeroGod struct {
+	session.SessionModule
+	browser    *Browser
+	advertiser *Advertiser
+}
+
+func NewZeroGod(s *session.Session) *ZeroGod {
+	mod := &ZeroGod{
+		SessionModule: session.NewSessionModule("zerogod", s),
+		browser:       nil,
+		advertiser:    nil,
+	}
+
+	mod.SessionModule.Requires("net.recon")
+
+	mod.AddHandler(session.NewModuleHandler("zerogod.discovery on", "",
+		"Start DNS-SD / mDNS discovery.",
+		func(args []string) error {
+			return mod.Start()
+		}))
+
+	mod.AddHandler(session.NewModuleHandler("zerogod.discovery off", "",
+		"Stop DNS-SD / mDNS discovery.",
+		func(args []string) error {
+			return mod.Stop()
+		}))
+
+	mod.AddHandler(session.NewModuleHandler("zerogod.show", "",
+		"Show discovered services.",
+		func(args []string) error {
+			return mod.show("", false)
+		}))
+
+	mod.AddHandler(session.NewModuleHandler("zerogod.show-full", "",
+		"Show discovered services and their DNS records.",
+		func(args []string) error {
+			return mod.show("", true)
+		}))
+
+	// TODO: add autocomplete
+	mod.AddHandler(session.NewModuleHandler("zerogod.show ADDRESS", "zerogod.show (.+)",
+		"Show discovered services given an ip address.",
+		func(args []string) error {
+			return mod.show(args[0], false)
+		}))
+
+	mod.AddHandler(session.NewModuleHandler("zerogod.show-full ADDRESS", "zerogod.show-full (.+)",
+		"Show discovered services and DNS records given an ip address.",
+		func(args []string) error {
+			return mod.show(args[0], true)
+		}))
+
+	mod.AddHandler(session.NewModuleHandler("zerogod.save ADDRESS FILENAME", "zerogod.save (.+) (.+)",
+		"Save the mDNS information of a given ADDRESS in the FILENAME yaml file.",
+		func(args []string) error {
+			return mod.save(args[0], args[1])
+		}))
+
+	mod.AddHandler(session.NewModuleHandler("zerogod.advertise FILENAME", "zerogod.advertise (.+)",
+		"Start advertising the mDNS services from the FILENAME yaml file.",
+		func(args []string) error {
+			if args[0] == "off" {
+				return mod.stopAdvertiser()
+			}
+			return mod.startAdvertiser(args[0])
+		}))
+
+	mod.AddHandler(session.NewModuleHandler("zerogod.advertise off", "",
+		"Start a previously started advertiser.",
+		func(args []string) error {
+			return mod.stopAdvertiser()
+		}))
+
+	mod.AddParam(session.NewStringParameter("zerogod.advertise.certificate",
+		"~/.bettercap-zerogod.cert.pem",
+		"",
+		"TLS certificate file (will be auto generated if filled but not existing) to use for advertised TCP services."))
+
+	mod.AddParam(session.NewStringParameter("zerogod.advertise.key",
+		"~/.bettercap-zerogod.key.pem",
+		"",
+		"TLS key file (will be auto generated if filled but not existing) to use for advertised TCP services."))
+
+	tls.CertConfigToModule("zerogod.advertise", &mod.SessionModule, tls.DefaultLegitConfig)
+
+	mod.AddParam(session.NewStringParameter("zerogod.ipp.save_path",
+		"~/.bettercap/zerogod/documents/",
+		"",
+		"If an IPP acceptor is started, this setting defines where to save documents received for printing."))
+
+	return mod
+}
+
+func (mod *ZeroGod) Name() string {
+	return "zerogod"
+}
+
+func (mod *ZeroGod) Description() string {
+	return "A DNS-SD / mDNS / Bonjour / Zeroconf module for discovery and spoofing."
+}
+
+func (mod *ZeroGod) Author() string {
+	return "Simone Margaritelli <evilsocket@gmail.com>"
+}
+
+func (mod *ZeroGod) Configure() (err error) {
+	if mod.Running() {
+		return session.ErrAlreadyStarted(mod.Name())
+	}
+
+	if mod.browser != nil {
+		mod.browser.Stop(false)
+	}
+
+	mod.browser = NewBrowser()
+
+	return
+}
+
+func (mod *ZeroGod) Start() (err error) {
+	if err = mod.Configure(); err != nil {
+		return err
+	}
+
+	// start the root discovery
+	if err = mod.startResolver(DNSSD_DISCOVERY_SERVICE); err != nil {
+		return err
+	}
+
+	return mod.SetRunning(true, func() {
+		mod.Info("service discovery started")
+		mod.browser.Wait()
+		mod.Info("service discovery stopped")
+	})
+}
+
+func (mod *ZeroGod) Stop() error {
+	return mod.SetRunning(false, func() {
+		mod.stopAdvertiser()
+		if mod.browser != nil {
+			mod.Debug("stopping discovery")
+
+			mod.browser.Stop(true)
+
+			mod.Debug("stopped")
+
+			mod.browser = nil
+		}
+	})
+}

modules/zerogod/zerogod_advertise.go

@@ -10,23 +10,16 @@ import (
 	"time"
 
 	tls_utils "github.com/bettercap/bettercap/v2/tls"
-	"github.com/bettercap/bettercap/v2/zeroconf"
 	"github.com/evilsocket/islazy/fs"
 	yaml "gopkg.in/yaml.v3"
 )
 
 type Advertiser struct {
-	Filename string
+	Filename  string
-
 	Services  []*ServiceData
 	Acceptors []*Acceptor
 }
 
-type setupResult struct {
-	err    error
-	server *zeroconf.Server
-}
-
 func (mod *ZeroGod) loadTLSConfig() (*tls.Config, error) {
 	var certFile string
 	var keyFile string

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
+}

modules/zerogod/zerogod_discovery.go

@@ -1,141 +1,14 @@
 package zerogod
 
 import (
-	"context"
 	"strings"
 
 	"github.com/bettercap/bettercap/v2/network"
 	"github.com/bettercap/bettercap/v2/session"
-	"github.com/bettercap/bettercap/v2/tls"
 	"github.com/bettercap/bettercap/v2/zeroconf"
 	"github.com/evilsocket/islazy/tui"
 )
 
-type ZeroGod struct {
-	session.SessionModule
-
-	advertiser  *Advertiser
-	rootContext context.Context
-	rootCancel  context.CancelFunc
-	resolvers   map[string]*zeroconf.Resolver
-	mapping     map[string]map[string]*zeroconf.ServiceEntry
-}
-
-func NewZeroGod(s *session.Session) *ZeroGod {
-	mod := &ZeroGod{
-		SessionModule: session.NewSessionModule("zerogod", s),
-		mapping:       make(map[string]map[string]*zeroconf.ServiceEntry),
-		resolvers:     make(map[string]*zeroconf.Resolver),
-	}
-
-	mod.SessionModule.Requires("net.recon")
-
-	mod.AddHandler(session.NewModuleHandler("zerogod.discovery on", "",
-		"Start DNS-SD / mDNS discovery.",
-		func(args []string) error {
-			return mod.Start()
-		}))
-
-	mod.AddHandler(session.NewModuleHandler("zerogod.discovery off", "",
-		"Stop DNS-SD / mDNS discovery.",
-		func(args []string) error {
-			return mod.Stop()
-		}))
-
-	mod.AddHandler(session.NewModuleHandler("zerogod.show", "",
-		"Show discovered services.",
-		func(args []string) error {
-			return mod.show("", false)
-		}))
-
-	mod.AddHandler(session.NewModuleHandler("zerogod.show-full", "",
-		"Show discovered services and their DNS records.",
-		func(args []string) error {
-			return mod.show("", true)
-		}))
-
-	// TODO: add autocomplete
-	mod.AddHandler(session.NewModuleHandler("zerogod.show ADDRESS", "zerogod.show (.+)",
-		"Show discovered services given an ip address.",
-		func(args []string) error {
-			return mod.show(args[0], false)
-		}))
-
-	mod.AddHandler(session.NewModuleHandler("zerogod.show-full ADDRESS", "zerogod.show-full (.+)",
-		"Show discovered services and DNS records given an ip address.",
-		func(args []string) error {
-			return mod.show(args[0], true)
-		}))
-
-	mod.AddHandler(session.NewModuleHandler("zerogod.save ADDRESS FILENAME", "zerogod.save (.+) (.+)",
-		"Save the mDNS information of a given ADDRESS in the FILENAME yaml file.",
-		func(args []string) error {
-			return mod.save(args[0], args[1])
-		}))
-
-	mod.AddHandler(session.NewModuleHandler("zerogod.advertise FILENAME", "zerogod.advertise (.+)",
-		"Start advertising the mDNS services from the FILENAME yaml file.",
-		func(args []string) error {
-			if args[0] == "off" {
-				return mod.stopAdvertiser()
-			}
-			return mod.startAdvertiser(args[0])
-		}))
-
-	mod.AddHandler(session.NewModuleHandler("zerogod.advertise off", "",
-		"Start a previously started advertiser.",
-		func(args []string) error {
-			return mod.stopAdvertiser()
-		}))
-
-	mod.AddParam(session.NewStringParameter("zerogod.advertise.certificate",
-		"~/.bettercap-zerogod.cert.pem",
-		"",
-		"TLS certificate file (will be auto generated if filled but not existing) to use for advertised TCP services."))
-
-	mod.AddParam(session.NewStringParameter("zerogod.advertise.key",
-		"~/.bettercap-zerogod.key.pem",
-		"",
-		"TLS key file (will be auto generated if filled but not existing) to use for advertised TCP services."))
-
-	tls.CertConfigToModule("zerogod.advertise", &mod.SessionModule, tls.DefaultLegitConfig)
-
-	mod.AddParam(session.NewStringParameter("zerogod.ipp.save_path",
-		"~/.bettercap/zerogod/documents/",
-		"",
-		"If an IPP acceptor is started, this setting defines where to save documents received for printing."))
-
-	return mod
-}
-
-func (mod *ZeroGod) Name() string {
-	return "zerogod"
-}
-
-func (mod *ZeroGod) Description() string {
-	return "A DNS-SD / mDNS / Bonjour / Zeroconf module for discovery and spoofing."
-}
-
-func (mod *ZeroGod) Author() string {
-	return "Simone Margaritelli <evilsocket@gmail.com>"
-}
-
-func (mod *ZeroGod) Configure() (err error) {
-	if mod.Running() {
-		return session.ErrAlreadyStarted(mod.Name())
-	}
-
-	if mod.rootContext != nil {
-		mod.rootCancel()
-	}
-
-	mod.mapping = make(map[string]map[string]*zeroconf.ServiceEntry)
-	mod.resolvers = make(map[string]*zeroconf.Resolver)
-	mod.rootContext, mod.rootCancel = context.WithCancel(context.Background())
-
-	return
-}
-
 type ServiceDiscoveryEvent struct {
 	Service  zeroconf.ServiceEntry `json:"service"`
 	Endpoint *network.Endpoint     `json:"endpoint"`
@@ -144,9 +17,9 @@ type ServiceDiscoveryEvent struct {
 func (mod *ZeroGod) onServiceDiscovered(svc *zeroconf.ServiceEntry) {
 	mod.Debug("%++v", *svc)
 
-	if svc.Service == "_services._dns-sd._udp" && len(svc.AddrIPv4) == 0 && len(svc.AddrIPv6) == 0 {
+	if svc.Service == DNSSD_DISCOVERY_SERVICE && len(svc.AddrIPv4) == 0 && len(svc.AddrIPv6) == 0 {
 		svcName := strings.Replace(svc.Instance, ".local", "", 1)
-		if _, found := mod.resolvers[svcName]; !found {
+		if !mod.browser.HasResolverFor(svcName) {
 			mod.Debug("discovered service %s", tui.Green(svcName))
 			if err := mod.startResolver(svcName); err != nil {
 				mod.Error("%v", err)
@@ -172,17 +45,10 @@ func (mod *ZeroGod) onServiceDiscovered(svc *zeroconf.ServiceEntry) {
 	for _, ip := range addresses {
 		address := ip.String()
 		if event.Endpoint = mod.Session.Lan.GetByIp(address); event.Endpoint != nil {
+			// update internal mapping
+			mod.browser.AddServiceFor(address, svc)
 			// update endpoint metadata
 			mod.updateEndpointMeta(address, event.Endpoint, svc)
-
-			// update internal module mapping
-			if ipServices, found := mod.mapping[address]; found {
-				ipServices[svc.ServiceInstanceName()] = svc
-			} else {
-				mod.mapping[address] = map[string]*zeroconf.ServiceEntry{
-					svc.ServiceInstanceName(): svc,
-				}
-			}
 			break
 		}
 	}
@@ -199,66 +65,16 @@ func (mod *ZeroGod) onServiceDiscovered(svc *zeroconf.ServiceEntry) {
 func (mod *ZeroGod) startResolver(service string) error {
 	mod.Debug("starting resolver for service %s", tui.Yellow(service))
 
-	resolver, err := zeroconf.NewResolver(nil)
+	if ch, err := mod.browser.StartBrowsing(service, "local.", mod); err != nil {
-	if err != nil {
 		return err
+	} else {
+		// start listening
+		go func() {
+			for entry := range ch {
+				mod.onServiceDiscovered(entry)
+			}
+		}()
 	}
 
-	// start listening
-	channel := make(chan *zeroconf.ServiceEntry)
-	go func() {
-		for entry := range channel {
-			mod.onServiceDiscovered(entry)
-		}
-	}()
-
-	// start browsing
-	go func() {
-		err = resolver.Browse(mod.rootContext, service, "local.", channel)
-		if err != nil {
-			mod.Error("%v", err)
-		}
-		mod.Debug("resolver for service %s stopped", tui.Yellow(service))
-	}()
-
-	mod.resolvers[service] = resolver
-
 	return nil
 }
-
-func (mod *ZeroGod) Start() (err error) {
-	if err = mod.Configure(); err != nil {
-		return err
-	}
-
-	// start the root discovery
-	if err = mod.startResolver("_services._dns-sd._udp"); err != nil {
-		return err
-	}
-
-	return mod.SetRunning(true, func() {
-		mod.Info("service discovery started")
-
-		<-mod.rootContext.Done()
-
-		mod.Info("service discovery stopped")
-	})
-}
-
-func (mod *ZeroGod) Stop() error {
-	return mod.SetRunning(false, func() {
-		mod.stopAdvertiser()
-
-		if mod.rootCancel != nil {
-			mod.Debug("stopping discovery")
-
-			mod.rootCancel()
-			<-mod.rootContext.Done()
-
-			mod.Debug("stopped")
-
-			mod.rootContext = nil
-			mod.rootCancel = nil
-		}
-	})
-}

modules/zerogod/zerogod_generic_handler.go

@@ -55,7 +55,7 @@ func viewString(b []byte) string {
 func handleGenericTCP(ctx *HandlerContext) {
 	defer ctx.client.Close()
 
-	ctx.mod.Debug("accepted generic tcp connection for service %s (port %d): %v", tui.Green(ctx.service), ctx.srvPort, ctx.client.RemoteAddr())
+	ctx.mod.Info("accepted generic tcp connection for service %s (port %d): %v", tui.Green(ctx.service), ctx.srvPort, ctx.client.RemoteAddr())
 
 	buf := make([]byte, 1024)
 	for {

modules/zerogod/zerogod_save.go

@@ -1,6 +1,7 @@
 package zerogod
 
 import (
+	"errors"
 	"fmt"
 	"io/ioutil"
 
@@ -32,6 +33,10 @@ func svcEntriesToData(services map[string]*zeroconf.ServiceEntry) []ServiceData
 }
 
 func (mod *ZeroGod) save(address, filename string) error {
+	if mod.browser == nil {
+		return errors.New("use 'zerogod.discovery on' to start the discovery first")
+	}
+
 	if address == "" {
 		return fmt.Errorf("address cannot be empty")
 	}
@@ -39,7 +44,7 @@ func (mod *ZeroGod) save(address, filename string) error {
 		return fmt.Errorf("filename cannot be empty")
 	}
 
-	if ipServices, found := mod.mapping[address]; found {
+	if ipServices := mod.browser.GetServicesFor(address); ipServices != nil {
 		services := svcEntriesToData(ipServices)
 		data, err := yaml.Marshal(services)
 		if err != nil {

modules/zerogod/zerogod_service.go

@@ -54,7 +54,7 @@ func (svc *ServiceData) Register(mod *ZeroGod, localHostName string) (err error)
 		if addr, err := net.LookupAddr(svc.Responder); err == nil && len(addr) > 0 {
 			responderHostName = addr[0]
 		} else {
-			mod.Debug("could not get responder %s reverse dns entry: %v", svc.Responder, err)
+			mod.Debug("could not get responder %s hostname (%v)", svc.Responder, err)
 		}
 
 		// if we don't have a host, create a .nip.io representation

modules/zerogod/zerogod_show.go

@@ -1,44 +1,33 @@
 package zerogod
 
 import (
+	"errors"
 	"fmt"
-	"sort"
+	"strings"
 
-	"github.com/bettercap/bettercap/v2/zeroconf"
 	"github.com/evilsocket/islazy/str"
 	"github.com/evilsocket/islazy/tui"
 )
 
-type entry struct {
-	ip       string
-	services map[string]*zeroconf.ServiceEntry
-}
-
 func (mod *ZeroGod) show(filter string, withData bool) error {
-	fmt.Fprintf(mod.Session.Events.Stdout, "\n")
+	if mod.browser == nil {
-
+		return errors.New("use 'zerogod.discovery on' to start the discovery first")
-	// convert to list for sorting
-	entries := make([]entry, 0)
-	for ip, services := range mod.mapping {
-		if filter == "" || ip == filter {
-			entries = append(entries, entry{ip, services})
-		}
 	}
 
-	sort.Slice(entries, func(i, j int) bool {
+	fmt.Fprintf(mod.Session.Events.Stdout, "\n")
-		return entries[i].ip < entries[j].ip
+
-	})
+	entries := mod.browser.ServicesByAddress(filter)
 
 	for _, entry := range entries {
-		if endpoint := mod.Session.Lan.GetByIp(entry.ip); endpoint != nil {
+		if endpoint := mod.Session.Lan.GetByIp(entry.Address); endpoint != nil {
 			fmt.Fprintf(mod.Session.Events.Stdout, "* %s (%s)\n", tui.Bold(endpoint.IpAddress), tui.Dim(endpoint.Vendor))
 		} else {
-			fmt.Fprintf(mod.Session.Events.Stdout, "* %s\n", tui.Bold(entry.ip))
+			fmt.Fprintf(mod.Session.Events.Stdout, "* %s\n", tui.Bold(entry.Address))
 		}
 
-		for name, svc := range entry.services {
+		for _, svc := range entry.Services {
 			fmt.Fprintf(mod.Session.Events.Stdout, "  %s (%s) [%v / %v]:%s\n",
-				tui.Green(name),
+				tui.Green(svc.ServiceInstanceName()),
 				tui.Dim(svc.HostName),
 				svc.AddrIPv4,
 				svc.AddrIPv6,
@@ -48,11 +37,22 @@ func (mod *ZeroGod) show(filter string, withData bool) error {
 			numFields := len(svc.Text)
 			if withData {
 				if numFields > 0 {
+					columns := []string{"key", "value"}
+					rows := make([][]string, 0)
+
 					for _, field := range svc.Text {
 						if field = str.Trim(field); len(field) > 0 {
-							fmt.Fprintf(mod.Session.Events.Stdout, "    %s\n", field)
+							keyval := strings.SplitN(field, "=", 2)
+							rows = append(rows, []string{
+								keyval[0],
+								keyval[1],
+							})
 						}
 					}
+
+					tui.Table(mod.Session.Events.Stdout, columns, rows)
+					fmt.Fprintf(mod.Session.Events.Stdout, "\n")
+
 				} else {
 					fmt.Fprintf(mod.Session.Events.Stdout, "    %s\n", tui.Dim("no data"))
 				}