From 209725d6230ceacfecad4837725203ddb803f192 Mon Sep 17 00:00:00 2001
From: Simone Margaritelli <evilsocket@gmail.com>
Date: Sun, 22 Sep 2024 17:40:23 +0200
Subject: [PATCH] misc: small fix or general refactoring i did not bother
 commenting

---
 .gitignore                           |  3 ++-
 modules/zerogod/zerogod_browser.go   | 18 ++++++++++++++++++
 modules/zerogod/zerogod_discovery.go |  2 +-
 session/script_builtin_runtime.go    | 21 +++++++++++++++++++++
 session/session.go                   |  1 +
 5 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/.gitignore b/.gitignore
index 29ef182b..8b655142 100644
--- a/.gitignore
+++ b/.gitignore
@@ -16,4 +16,5 @@ stage/
 .idea
 /cover.out
 /can-data
-/test*.yml
\ No newline at end of file
+/test*.yml
+/zerochaos.js
\ No newline at end of file
diff --git a/modules/zerogod/zerogod_browser.go b/modules/zerogod/zerogod_browser.go
index 98c87bca..e7e72707 100644
--- a/modules/zerogod/zerogod_browser.go
+++ b/modules/zerogod/zerogod_browser.go
@@ -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)
 
diff --git a/modules/zerogod/zerogod_discovery.go b/modules/zerogod/zerogod_discovery.go
index b85345e3..4e805271 100644
--- a/modules/zerogod/zerogod_discovery.go
+++ b/modules/zerogod/zerogod_discovery.go
@@ -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)
diff --git a/session/script_builtin_runtime.go b/session/script_builtin_runtime.go
index 5394315c..ffcd14b4 100644
--- a/session/script_builtin_runtime.go
+++ b/session/script_builtin_runtime.go
@@ -95,6 +95,27 @@ func jsOnEventFunc(call otto.FunctionCall) otto.Value {
 	return js.NullValue
 }
 
+func jsSaveToFileFunc(call otto.FunctionCall) otto.Value {
+	argv := call.ArgumentList
+	argc := len(argv)
+	if argc != 2 {
+		return js.ReportError("saveToFile accepts two string arguments")
+	} else if argv[0].IsString() == false {
+		return js.ReportError("saveToFile accepts two string arguments")
+	} else if argv[1].IsString() == false {
+		return js.ReportError("saveToFile accepts two string arguments")
+	}
+
+	fileName := argv[0].String()
+	data := argv[1].String()
+
+	if err := ioutil.WriteFile(fileName, []byte(data), os.ModePerm); err != nil {
+		return js.ReportError("error writing to '%s': %v", fileName, err)
+	}
+
+	return js.NullValue
+}
+
 func jsSaveJSONFunc(call otto.FunctionCall) otto.Value {
 	argv := call.ArgumentList
 	argc := len(argv)
diff --git a/session/session.go b/session/session.go
index 55f0c76e..4d0d0ef6 100644
--- a/session/session.go
+++ b/session/session.go
@@ -328,6 +328,7 @@ func (s *Session) Start() error {
 	plugin.Defines["fileExists"] = jsFileExistsFunc
 	plugin.Defines["loadJSON"] = jsLoadJSONFunc
 	plugin.Defines["saveJSON"] = jsSaveJSONFunc
+	plugin.Defines["saveToFile"] = jsSaveToFileFunc
 	plugin.Defines["onEvent"] = jsOnEventFunc
 	plugin.Defines["session"] = s