refact: refactored module SetRunning method (fixes #49)

modules/api_rest.go

@@ -161,24 +161,21 @@ func (api *RestAPI) Start() error {
 		return err
 	}
 
-	api.SetRunning(true)
+	api.SetRunning(true, func() {
-	go 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 api.SetRunning(false, func() {
-		return session.ErrAlreadyStopped
+		ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
-	}
+		defer cancel()
-	api.SetRunning(false)
+		api.server.Shutdown(ctx)
-	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	})
-	defer cancel()
-	return api.server.Shutdown(ctx)
 }

modules/arp_spoof.go

@@ -173,14 +173,11 @@ func (p *ArpSpoofer) Configure() error {
 }
 
 func (p *ArpSpoofer) Start() error {
-	if p.Running() == true {
+	if err := p.Configure(); err != nil {
-		return session.ErrAlreadyStarted
-	} else if err := p.Configure(); err != nil {
 		return err
 	}
 
-	p.SetRunning(true)
+	return p.SetRunning(true, func() {
-	go 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 p.SetRunning(false, func() {
-		return session.ErrAlreadyStopped
+		log.Info("Waiting for ARP spoofer to stop ...")
-	}
+		<-p.done
-
+		p.unSpoof()
-	log.Info("Waiting for ARP spoofer to stop ...")
+		p.ban = false
-
+	})
-	p.SetRunning(false)
-
-	<-p.done
-
-	p.unSpoof()
-	p.ban = false
-
-	return nil
 }

modules/dhcp6_spoof.go

@@ -369,9 +369,7 @@ func (s *DHCP6Spoofer) Start() error {
 		return err
 	}
 
-	s.SetRunning(true)
+	return s.SetRunning(true, func() {
-
-	go 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 s.SetRunning(false, func() {
-		return session.ErrAlreadyStopped
+		s.Handle.Close()
-	}
+	})
-	s.SetRunning(false)
-	s.Handle.Close()
-	return nil
 }

modules/dns_spoof.go

@@ -266,9 +266,7 @@ func (s *DNSSpoofer) Start() error {
 		return err
 	}
 
-	s.SetRunning(true)
+	return s.SetRunning(true, func() {
-
-	go 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 s.SetRunning(false, func() {
-		return session.ErrAlreadyStopped
+		s.Handle.Close()
-	}
+	})
-	s.SetRunning(false)
-	s.Handle.Close()
-	return nil
 }

modules/events_stream.go

@@ -79,15 +79,11 @@ func (s *EventsStream) Configure() error {
 }
 
 func (s *EventsStream) Start() error {
-	if s.Running() == true {
+	if err := s.Configure(); err != nil {
-		return session.ErrAlreadyStarted
-	} else if err := s.Configure(); err != nil {
 		return err
 	}
 
-	s.SetRunning(true)
+	return s.SetRunning(true, func() {
-
-	go 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 s.SetRunning(false, func() {
-		return session.ErrAlreadyStopped
+		s.quit <- true
-	}
+	})
-	s.SetRunning(false)
-	s.quit <- true
-	return nil
 }

modules/http_proxy.go

@@ -93,17 +93,13 @@ func (p *HttpProxy) Start() error {
 		return err
 	}
 
-	p.SetRunning(true)
+	return p.SetRunning(true, func() {
-	p.proxy.Start()
+		p.proxy.Start()
-
+	})
-	return nil
 }
 
 func (p *HttpProxy) Stop() error {
-	if p.Running() == false {
+	return p.SetRunning(false, func() {
-		return session.ErrAlreadyStopped
+		p.proxy.Stop()
-	}
+	})
-	p.SetRunning(false)
-
-	return p.proxy.Stop()
 }

modules/http_server.go

@@ -103,24 +103,19 @@ func (httpd *HttpServer) Start() error {
 		return err
 	}
 
-	httpd.SetRunning(true)
+	return httpd.SetRunning(true, func() {
-	go 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 httpd.SetRunning(false, func() {
-		return session.ErrAlreadyStopped
+		ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
-	}
+		defer cancel()
-	httpd.SetRunning(false)
+		httpd.server.Shutdown(ctx)
-	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	})
-	defer cancel()
-	return httpd.server.Shutdown(ctx)
 }

modules/https_proxy.go

@@ -131,17 +131,13 @@ func (p *HttpsProxy) Start() error {
 		return err
 	}
 
-	p.SetRunning(true)
+	return p.SetRunning(true, func() {
-	p.proxy.Start()
+		p.proxy.Start()
-
+	})
-	return nil
 }
 
 func (p *HttpsProxy) Stop() error {
-	if p.Running() == false {
+	return p.SetRunning(false, func() {
-		return session.ErrAlreadyStopped
+		p.proxy.Stop()
-	}
+	})
-	p.SetRunning(false)
-
-	return p.proxy.Stop()
 }

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 {
+	if err := mc.Configure(); err != nil {
-		return session.ErrAlreadyStarted
-	} else if err := mc.Configure(); err != nil {
 		return err
 	} else if err := mc.setMac(mc.fakeMac); err != nil {
 		return err
 	}
 
-	mc.SetRunning(true)
+	return mc.SetRunning(true, func() {
-	log.Info("Interface mac address set to %s", core.Bold(mc.fakeMac.String()))
+		log.Info("Interface mac address set to %s", core.Bold(mc.fakeMac.String()))
-
+	})
-	return nil
 }
 
 func (mc *MacChanger) Stop() error {
-	if mc.Running() == false {
+	if err := mc.setMac(mc.originalMac); err != nil {
-		return session.ErrAlreadyStopped
-	} else if err := mc.setMac(mc.originalMac); err != nil {
 		return err
 	}
 
-	mc.SetRunning(false)
+	return mc.SetRunning(false, func() {
-	log.Info("Interface mac address restored to %s", core.Bold(mc.originalMac.String()))
+		log.Info("Interface mac address restored to %s", core.Bold(mc.originalMac.String()))
-	return nil
+	})
 }

modules/net_probe.go

@@ -87,15 +87,11 @@ func (p *Prober) Configure() error {
 }
 
 func (p *Prober) Start() error {
-	if p.Running() == true {
+	if err := p.Configure(); err != nil {
-		return session.ErrAlreadyStarted
-	} else if err := p.Configure(); err != nil {
 		return err
 	}
 
-	p.SetRunning(true)
+	return p.SetRunning(true, func() {
-
-	go 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 p.SetRunning(false, nil)
-		return session.ErrAlreadyStopped
-	}
-	p.SetRunning(false)
-	return nil
 }

modules/net_recon.go

@@ -125,15 +125,11 @@ func (d *Discovery) Configure() error {
 }
 
 func (d *Discovery) Start() error {
-	if d.Running() == true {
+	if err := d.Configure(); err != nil {
-		return session.ErrAlreadyStarted
-	} else if err := d.Configure(); err != nil {
 		return err
 	}
 
-	d.SetRunning(true)
+	return d.SetRunning(true, func() {
-
-	go 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 d.SetRunning(false, func() {
-		return session.ErrAlreadyStopped
+		d.quit <- true
-	}
+	})
-	d.quit <- true
-	d.SetRunning(false)
-	return nil
 }

modules/net_sniff.go

@@ -143,9 +143,7 @@ func (s *Sniffer) Start() error {
 		return err
 	}
 
-	s.SetRunning(true)
+	return s.SetRunning(true, func() {
-
-	go 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 s.SetRunning(false, func() {
-		return session.ErrAlreadyStopped
+		s.Ctx.Close()
-	}
+	})
-	s.SetRunning(false)
-	s.Ctx.Close()
-
-	return nil
 }

modules/ticker.go

@@ -80,8 +80,7 @@ func (t *Ticker) Start() error {
 		return err
 	}
 
-	t.SetRunning(true)
+	return t.SetRunning(true, func() {
-	go 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 t.SetRunning(false, nil)
-		return session.ErrAlreadyStopped
-	}
-	t.SetRunning(false)
-	return nil
 }

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,

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
 }