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
parent d2f13a3293
commit 209725d623
5 changed files with 43 additions and 2 deletions

3
.gitignore vendored
View file

@ -16,4 +16,5 @@ stage/
.idea .idea
/cover.out /cover.out
/can-data /can-data
/test*.yml /test*.yml
/zerochaos.js

View file

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

View file

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

View file

@ -95,6 +95,27 @@ func jsOnEventFunc(call otto.FunctionCall) otto.Value {
return js.NullValue 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 { func jsSaveJSONFunc(call otto.FunctionCall) otto.Value {
argv := call.ArgumentList argv := call.ArgumentList
argc := len(argv) argc := len(argv)

View file

@ -328,6 +328,7 @@ func (s *Session) Start() error {
plugin.Defines["fileExists"] = jsFileExistsFunc plugin.Defines["fileExists"] = jsFileExistsFunc
plugin.Defines["loadJSON"] = jsLoadJSONFunc plugin.Defines["loadJSON"] = jsLoadJSONFunc
plugin.Defines["saveJSON"] = jsSaveJSONFunc plugin.Defines["saveJSON"] = jsSaveJSONFunc
plugin.Defines["saveToFile"] = jsSaveToFileFunc
plugin.Defines["onEvent"] = jsOnEventFunc plugin.Defines["onEvent"] = jsOnEventFunc
plugin.Defines["session"] = s plugin.Defines["session"] = s