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

This commit is contained in:
evilsocket 2018-02-11 21:04:00 +01:00
parent c2cd198d37
commit 2d26f77acb

View file

@ -10,21 +10,11 @@ import (
type Discovery struct { type Discovery struct {
session.SessionModule session.SessionModule
refresh int
before net.ArpTable
current net.ArpTable
quit chan bool
} }
func NewDiscovery(s *session.Session) *Discovery { func NewDiscovery(s *session.Session) *Discovery {
d := &Discovery{ d := &Discovery{
SessionModule: session.NewSessionModule("net.recon", s), SessionModule: session.NewSessionModule("net.recon", s),
refresh: 1,
before: nil,
current: nil,
quit: make(chan bool),
} }
d.AddHandler(session.NewModuleHandler("net.recon on", "", d.AddHandler(session.NewModuleHandler("net.recon on", "",
@ -40,25 +30,25 @@ func NewDiscovery(s *session.Session) *Discovery {
})) }))
d.AddHandler(session.NewModuleHandler("net.show", "", d.AddHandler(session.NewModuleHandler("net.show", "",
"Show current hosts list (default sorting by ip).", "Show cache hosts list (default sorting by ip).",
func(args []string) error { func(args []string) error {
return d.Show("address") return d.Show("address")
})) }))
d.AddHandler(session.NewModuleHandler("net.show by seen", "", d.AddHandler(session.NewModuleHandler("net.show by seen", "",
"Show current hosts list (sort by last seen).", "Show cache hosts list (sort by last seen).",
func(args []string) error { func(args []string) error {
return d.Show("seen") return d.Show("seen")
})) }))
d.AddHandler(session.NewModuleHandler("net.show by sent", "", d.AddHandler(session.NewModuleHandler("net.show by sent", "",
"Show current hosts list (sort by sent packets).", "Show cache hosts list (sort by sent packets).",
func(args []string) error { func(args []string) error {
return d.Show("sent") return d.Show("sent")
})) }))
d.AddHandler(session.NewModuleHandler("net.show by rcvd", "", d.AddHandler(session.NewModuleHandler("net.show by rcvd", "",
"Show current hosts list (sort by received packets).", "Show cache hosts list (sort by received packets).",
func(args []string) error { func(args []string) error {
return d.Show("rcvd") return d.Show("rcvd")
})) }))
@ -78,11 +68,11 @@ func (d Discovery) Author() string {
return "Simone Margaritelli <evilsocket@protonmail.com>" return "Simone Margaritelli <evilsocket@protonmail.com>"
} }
func (d *Discovery) runDiff() { func (d *Discovery) runDiff(cache net.ArpTable) {
// check for endpoints who disappeared // check for endpoints who disappeared
var rem net.ArpTable = make(net.ArpTable) var rem net.ArpTable = make(net.ArpTable)
for mac, t := range d.Session.Targets.Targets { for mac, t := range d.Session.Targets.Targets {
if _, found := d.current[mac]; found == false { if _, found := cache[mac]; found == false {
rem[mac] = t.IpAddress rem[mac] = t.IpAddress
} }
} }
@ -92,7 +82,7 @@ func (d *Discovery) runDiff() {
} }
// now check for new friends ^_^ // now check for new friends ^_^
for ip, mac := range d.current { for ip, mac := range cache {
d.Session.Targets.AddIfNew(ip, mac) d.Session.Targets.AddIfNew(ip, mac)
} }
} }
@ -107,29 +97,20 @@ func (d *Discovery) Start() error {
} }
return d.SetRunning(true, func() { return d.SetRunning(true, func() {
for { every := time.Duration(1) * time.Second
select { iface := d.Session.Interface.Name()
case <-time.After(time.Duration(d.refresh) * time.Second):
var err error
if d.current, err = net.ArpUpdate(d.Session.Interface.Name()); err != nil { for d.Running() {
log.Error("%s", err) if table, err := net.ArpUpdate(iface); err != nil {
continue log.Error("%s", err)
} } else {
d.runDiff(table)
d.runDiff()
d.before = d.current
case <-d.quit:
return
} }
time.Sleep(every)
} }
}) })
} }
func (d *Discovery) Stop() error { func (d *Discovery) Stop() error {
return d.SetRunning(false, func() { return d.SetRunning(false, nil)
d.quit <- true
})
} }