From dd2fd0f30f307b821cd62c16fa2f2e4e673cc2c8 Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Mon, 8 Jan 2018 09:15:36 +0100
Subject: [PATCH] refact: minor refactoring of net.recon module

---
 session/modules/net_recon.go | 90 ++++++++++++++++++++----------------
 1 file changed, 49 insertions(+), 41 deletions(-)

diff --git a/session/modules/net_recon.go b/session/modules/net_recon.go
index efbba46a..1f21e1d9 100644
--- a/session/modules/net_recon.go
+++ b/session/modules/net_recon.go
@@ -65,6 +65,54 @@ func (d Discovery) OnSessionEnded(s *session.Session) {
 	}
 }
 
+func (d *Discovery) checkShared(new net.ArpTable) {
+	n_gw_shared := 0
+	for ip, mac := range new {
+		if ip != d.Session.Gateway.IpAddress && mac == d.Session.Gateway.HwAddress {
+			n_gw_shared++
+		}
+	}
+
+	if n_gw_shared > 0 {
+		a := ""
+		b := ""
+		if n_gw_shared == 1 {
+			a = ""
+			b = "s"
+		} else {
+			a = "s"
+			b = ""
+		}
+
+		d.Session.Events.Log(session.WARNING, "Found %d endpoint%s which share%s the same MAC of the gateway (%s), there're might be some IP isolation going on, skipping.", n_gw_shared, a, b, d.Session.Gateway.HwAddress)
+	}
+}
+
+func (d *Discovery) runDiff() {
+	var new net.ArpTable = make(net.ArpTable)
+	var rem net.ArpTable = make(net.ArpTable)
+
+	if d.before != nil {
+		new = net.ArpDiff(d.current, d.before)
+		rem = net.ArpDiff(d.before, d.current)
+	} else {
+		new = d.current
+	}
+
+	if len(new) > 0 || len(rem) > 0 {
+		d.checkShared(new)
+
+		// refresh target pool
+		for ip, mac := range rem {
+			d.Session.Targets.Remove(ip, mac)
+		}
+
+		for ip, mac := range new {
+			d.Session.Targets.AddIfNotExist(ip, mac)
+		}
+	}
+}
+
 func (d *Discovery) Start() error {
 	if d.Running() == false {
 		d.SetRunning(true)
@@ -80,47 +128,7 @@ func (d *Discovery) Start() error {
 						continue
 					}
 
-					var new net.ArpTable = make(net.ArpTable)
-					var rem net.ArpTable = make(net.ArpTable)
-
-					if d.before != nil {
-						new = net.ArpDiff(d.current, d.before)
-						rem = net.ArpDiff(d.before, d.current)
-					} else {
-						new = d.current
-					}
-
-					if len(new) > 0 || len(rem) > 0 {
-						n_gw_shared := 0
-						for ip, mac := range new {
-							if ip != d.Session.Gateway.IpAddress && mac == d.Session.Gateway.HwAddress {
-								n_gw_shared++
-							}
-						}
-
-						if n_gw_shared > 0 {
-							a := ""
-							b := ""
-							if n_gw_shared == 1 {
-								a = ""
-								b = "s"
-							} else {
-								a = "s"
-								b = ""
-							}
-
-							d.Session.Events.Log(session.WARNING, "WARNING: Found %d endpoint%s which share%s the same MAC of the gateway (%s), there're might be some IP isolation going on, skipping.", n_gw_shared, a, b, d.Session.Gateway.HwAddress)
-						}
-
-						// refresh target pool
-						for ip, mac := range rem {
-							d.Session.Targets.Remove(ip, mac)
-						}
-
-						for ip, mac := range new {
-							d.Session.Targets.AddIfNotExist(ip, mac)
-						}
-					}
+					d.runDiff()
 
 					d.before = d.current