From 9c6eb70eb36d6364ed167aeb0849f93d4cd41451 Mon Sep 17 00:00:00 2001 From: evilsocket Date: Sun, 11 Feb 2018 01:43:48 +0100 Subject: [PATCH] refact: refactored module SetRunning method (fixes #49) --- modules/api_rest.go | 17 +++++++---------- modules/arp_spoof.go | 31 +++++++++---------------------- modules/dhcp6_spoof.go | 17 +++++------------ modules/dns_spoof.go | 17 +++++------------ modules/events_stream.go | 21 ++++++--------------- modules/http_proxy.go | 16 ++++++---------- modules/http_server.go | 19 +++++++------------ modules/https_proxy.go | 16 ++++++---------- modules/mac_changer.go | 21 ++++++++------------- modules/net_probe.go | 18 ++++-------------- modules/net_recon.go | 21 ++++++--------------- modules/net_sniff.go | 18 +++++------------- modules/ticker.go | 13 +++---------- modules/wol.go | 11 ++++++----- session/module.go | 21 ++++++++++++++++++++- 15 files changed, 103 insertions(+), 174 deletions(-) diff --git a/modules/api_rest.go b/modules/api_rest.go index 6640809f..90904256 100644 --- a/modules/api_rest.go +++ b/modules/api_rest.go @@ -161,24 +161,21 @@ func (api *RestAPI) Start() error { return err } - api.SetRunning(true) - go func() { + api.SetRunning(true, func() { log.Info("API server starting on https://%s", api.server.Addr) err := api.server.ListenAndServeTLS(api.certFile, api.keyFile) if err != nil && err != http.ErrServerClosed { panic(err) } - }() + }) return nil } func (api *RestAPI) Stop() error { - if api.Running() == false { - return session.ErrAlreadyStopped - } - api.SetRunning(false) - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() - return api.server.Shutdown(ctx) + return api.SetRunning(false, func() { + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + api.server.Shutdown(ctx) + }) } diff --git a/modules/arp_spoof.go b/modules/arp_spoof.go index cba3ae35..9f53a4fb 100644 --- a/modules/arp_spoof.go +++ b/modules/arp_spoof.go @@ -173,14 +173,11 @@ func (p *ArpSpoofer) Configure() error { } func (p *ArpSpoofer) Start() error { - if p.Running() == true { - return session.ErrAlreadyStarted - } else if err := p.Configure(); err != nil { + if err := p.Configure(); err != nil { return err } - p.SetRunning(true) - go func() { + return p.SetRunning(true, func() { from := p.Session.Gateway.IP from_hw := p.Session.Interface.HW @@ -192,24 +189,14 @@ func (p *ArpSpoofer) Start() error { } p.done <- true - }() - - return nil + }) } func (p *ArpSpoofer) Stop() error { - if p.Running() == false { - return session.ErrAlreadyStopped - } - - log.Info("Waiting for ARP spoofer to stop ...") - - p.SetRunning(false) - - <-p.done - - p.unSpoof() - p.ban = false - - return nil + return p.SetRunning(false, func() { + log.Info("Waiting for ARP spoofer to stop ...") + <-p.done + p.unSpoof() + p.ban = false + }) } diff --git a/modules/dhcp6_spoof.go b/modules/dhcp6_spoof.go index aa13d706..b1fbb358 100644 --- a/modules/dhcp6_spoof.go +++ b/modules/dhcp6_spoof.go @@ -369,9 +369,7 @@ func (s *DHCP6Spoofer) Start() error { return err } - s.SetRunning(true) - - go func() { + return s.SetRunning(true, func() { defer s.Handle.Close() src := gopacket.NewPacketSource(s.Handle, s.Handle.LinkType()) @@ -382,16 +380,11 @@ func (s *DHCP6Spoofer) Start() error { s.onPacket(packet) } - }() - - return nil + }) } func (s *DHCP6Spoofer) Stop() error { - if s.Running() == false { - return session.ErrAlreadyStopped - } - s.SetRunning(false) - s.Handle.Close() - return nil + return s.SetRunning(false, func() { + s.Handle.Close() + }) } diff --git a/modules/dns_spoof.go b/modules/dns_spoof.go index 93a23206..15d2a880 100644 --- a/modules/dns_spoof.go +++ b/modules/dns_spoof.go @@ -266,9 +266,7 @@ func (s *DNSSpoofer) Start() error { return err } - s.SetRunning(true) - - go func() { + return s.SetRunning(true, func() { defer s.Handle.Close() src := gopacket.NewPacketSource(s.Handle, s.Handle.LinkType()) @@ -279,16 +277,11 @@ func (s *DNSSpoofer) Start() error { s.onPacket(packet) } - }() - - return nil + }) } func (s *DNSSpoofer) Stop() error { - if s.Running() == false { - return session.ErrAlreadyStopped - } - s.SetRunning(false) - s.Handle.Close() - return nil + return s.SetRunning(false, func() { + s.Handle.Close() + }) } diff --git a/modules/events_stream.go b/modules/events_stream.go index 8510f319..dc4ee561 100644 --- a/modules/events_stream.go +++ b/modules/events_stream.go @@ -79,15 +79,11 @@ func (s *EventsStream) Configure() error { } func (s *EventsStream) Start() error { - if s.Running() == true { - return session.ErrAlreadyStarted - } else if err := s.Configure(); err != nil { + if err := s.Configure(); err != nil { return err } - s.SetRunning(true) - - go func() { + return s.SetRunning(true, func() { for { var e session.Event select { @@ -99,9 +95,7 @@ func (s *EventsStream) Start() error { return } } - }() - - return nil + }) } func (s *EventsStream) Show(limit int) error { @@ -123,10 +117,7 @@ func (s *EventsStream) Show(limit int) error { } func (s *EventsStream) Stop() error { - if s.Running() == false { - return session.ErrAlreadyStopped - } - s.SetRunning(false) - s.quit <- true - return nil + return s.SetRunning(false, func() { + s.quit <- true + }) } diff --git a/modules/http_proxy.go b/modules/http_proxy.go index 9f1075a7..1861b304 100644 --- a/modules/http_proxy.go +++ b/modules/http_proxy.go @@ -93,17 +93,13 @@ func (p *HttpProxy) Start() error { return err } - p.SetRunning(true) - p.proxy.Start() - - return nil + return p.SetRunning(true, func() { + p.proxy.Start() + }) } func (p *HttpProxy) Stop() error { - if p.Running() == false { - return session.ErrAlreadyStopped - } - p.SetRunning(false) - - return p.proxy.Stop() + return p.SetRunning(false, func() { + p.proxy.Stop() + }) } diff --git a/modules/http_server.go b/modules/http_server.go index 7e55e2bc..15483f99 100644 --- a/modules/http_server.go +++ b/modules/http_server.go @@ -103,24 +103,19 @@ func (httpd *HttpServer) Start() error { return err } - httpd.SetRunning(true) - go func() { + return httpd.SetRunning(true, func() { log.Info("httpd server starting on http://%s", httpd.server.Addr) err := httpd.server.ListenAndServe() if err != nil && err != http.ErrServerClosed { panic(err) } - }() - - return nil + }) } func (httpd *HttpServer) Stop() error { - if httpd.Running() == false { - return session.ErrAlreadyStopped - } - httpd.SetRunning(false) - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() - return httpd.server.Shutdown(ctx) + return httpd.SetRunning(false, func() { + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + httpd.server.Shutdown(ctx) + }) } diff --git a/modules/https_proxy.go b/modules/https_proxy.go index 1bd04406..897779d4 100644 --- a/modules/https_proxy.go +++ b/modules/https_proxy.go @@ -131,17 +131,13 @@ func (p *HttpsProxy) Start() error { return err } - p.SetRunning(true) - p.proxy.Start() - - return nil + return p.SetRunning(true, func() { + p.proxy.Start() + }) } func (p *HttpsProxy) Stop() error { - if p.Running() == false { - return session.ErrAlreadyStopped - } - p.SetRunning(false) - - return p.proxy.Stop() + return p.SetRunning(false, func() { + p.proxy.Stop() + }) } diff --git a/modules/mac_changer.go b/modules/mac_changer.go index 92c9b775..6df4536a 100644 --- a/modules/mac_changer.go +++ b/modules/mac_changer.go @@ -103,28 +103,23 @@ func (mc *MacChanger) setMac(mac net.HardwareAddr) error { } func (mc *MacChanger) Start() error { - if mc.Running() == true { - return session.ErrAlreadyStarted - } else if err := mc.Configure(); err != nil { + if err := mc.Configure(); err != nil { return err } else if err := mc.setMac(mc.fakeMac); err != nil { return err } - mc.SetRunning(true) - log.Info("Interface mac address set to %s", core.Bold(mc.fakeMac.String())) - - return nil + return mc.SetRunning(true, func() { + log.Info("Interface mac address set to %s", core.Bold(mc.fakeMac.String())) + }) } func (mc *MacChanger) Stop() error { - if mc.Running() == false { - return session.ErrAlreadyStopped - } else if err := mc.setMac(mc.originalMac); err != nil { + if err := mc.setMac(mc.originalMac); err != nil { return err } - mc.SetRunning(false) - log.Info("Interface mac address restored to %s", core.Bold(mc.originalMac.String())) - return nil + return mc.SetRunning(false, func() { + log.Info("Interface mac address restored to %s", core.Bold(mc.originalMac.String())) + }) } diff --git a/modules/net_probe.go b/modules/net_probe.go index 4d774986..022f2840 100644 --- a/modules/net_probe.go +++ b/modules/net_probe.go @@ -87,15 +87,11 @@ func (p *Prober) Configure() error { } func (p *Prober) Start() error { - if p.Running() == true { - return session.ErrAlreadyStarted - } else if err := p.Configure(); err != nil { + if err := p.Configure(); err != nil { return err } - p.SetRunning(true) - - go func() { + return p.SetRunning(true, func() { list, err := iprange.Parse(p.Session.Interface.CIDR()) if err != nil { log.Fatal("%s", err) @@ -121,15 +117,9 @@ func (p *Prober) Start() error { time.Sleep(5 * time.Second) } - }() - - return nil + }) } func (p *Prober) Stop() error { - if p.Running() == false { - return session.ErrAlreadyStopped - } - p.SetRunning(false) - return nil + return p.SetRunning(false, nil) } diff --git a/modules/net_recon.go b/modules/net_recon.go index 8bee565a..8681f864 100644 --- a/modules/net_recon.go +++ b/modules/net_recon.go @@ -125,15 +125,11 @@ func (d *Discovery) Configure() error { } func (d *Discovery) Start() error { - if d.Running() == true { - return session.ErrAlreadyStarted - } else if err := d.Configure(); err != nil { + if err := d.Configure(); err != nil { return err } - d.SetRunning(true) - - go func() { + return d.SetRunning(true, func() { for { select { case <-time.After(time.Duration(d.refresh) * time.Second): @@ -152,16 +148,11 @@ func (d *Discovery) Start() error { return } } - }() - - return nil + }) } func (d *Discovery) Stop() error { - if d.Running() == false { - return session.ErrAlreadyStopped - } - d.quit <- true - d.SetRunning(false) - return nil + return d.SetRunning(false, func() { + d.quit <- true + }) } diff --git a/modules/net_sniff.go b/modules/net_sniff.go index 36fa6dfd..2363492d 100644 --- a/modules/net_sniff.go +++ b/modules/net_sniff.go @@ -143,9 +143,7 @@ func (s *Sniffer) Start() error { return err } - s.SetRunning(true) - - go func() { + return s.SetRunning(true, func() { s.Stats = NewSnifferStats() src := gopacket.NewPacketSource(s.Ctx.Handle, s.Ctx.Handle.LinkType()) @@ -180,17 +178,11 @@ func (s *Sniffer) Start() error { } } } - }() - - return nil + }) } func (s *Sniffer) Stop() error { - if s.Running() == false { - return session.ErrAlreadyStopped - } - s.SetRunning(false) - s.Ctx.Close() - - return nil + return s.SetRunning(false, func() { + s.Ctx.Close() + }) } diff --git a/modules/ticker.go b/modules/ticker.go index 5114806e..d323cf26 100644 --- a/modules/ticker.go +++ b/modules/ticker.go @@ -80,8 +80,7 @@ func (t *Ticker) Start() error { return err } - t.SetRunning(true) - go func() { + return t.SetRunning(true, func() { log.Info("Ticker running with period %.fs.", t.Period.Seconds()) tick := time.Tick(t.Period) for _ = range tick { @@ -95,15 +94,9 @@ func (t *Ticker) Start() error { } } } - }() - - return nil + }) } func (t *Ticker) Stop() error { - if t.Running() == false { - return session.ErrAlreadyStopped - } - t.SetRunning(false) - return nil + return t.SetRunning(false, nil) } diff --git a/modules/wol.go b/modules/wol.go index 5de7bdb8..b4dbaca8 100644 --- a/modules/wol.go +++ b/modules/wol.go @@ -101,11 +101,11 @@ func buildPayload(mac string) []byte { } func (w *WOL) wolETH(mac string) error { + w.SetRunning(true, nil) + defer w.SetRunning(false, nil) + payload := buildPayload(mac) log.Info("Sending %d bytes of ethernet WOL packet to %s", len(payload), core.Bold(mac)) - w.SetRunning(true) - defer w.SetRunning(false) - eth := layers.Ethernet{ SrcMAC: w.Session.Interface.HW, DstMAC: layers.EthernetBroadcast, @@ -126,10 +126,11 @@ func (w *WOL) wolETH(mac string) error { } func (w *WOL) wolUDP(mac string) error { + w.SetRunning(true, nil) + defer w.SetRunning(false, nil) + payload := buildPayload(mac) log.Info("Sending %d bytes of UDP WOL packet to %s", len(payload), core.Bold(mac)) - w.SetRunning(true) - defer w.SetRunning(false) eth := layers.Ethernet{ SrcMAC: w.Session.Interface.HW, diff --git a/session/module.go b/session/module.go index 9a2d343d..2bc0ad50 100644 --- a/session/module.go +++ b/session/module.go @@ -122,9 +122,16 @@ func (m *SessionModule) Running() bool { return m.Started } -func (m *SessionModule) SetRunning(running bool) { +func (m *SessionModule) SetRunning(running bool, cb func()) error { m.StatusLock.Lock() defer m.StatusLock.Unlock() + + if running && m.Started == true { + return ErrAlreadyStarted + } else if running == false && m.Started == false { + return ErrAlreadyStopped + } + m.Started = running if *m.Session.Options.Debug == true { @@ -134,4 +141,16 @@ func (m *SessionModule) SetRunning(running bool) { m.Session.Events.Add("mod.stopped", m.Name) } } + + if cb != nil { + if running == true { + // this is the worker, start async + go cb() + } else { + // stop callback, this is sync + cb() + } + } + + return nil }