From 7356bf229562bdaec24f4c566f10483e24b80086 Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Sat, 17 Feb 2018 00:43:31 +0100
Subject: [PATCH] misc: small fix or general refactoring i did not bother
 commenting

---
 modules/wifi_recon.go | 38 ++++++++++++--------------------------
 packets/dot11.go      | 25 +++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 26 deletions(-)

diff --git a/modules/wifi_recon.go b/modules/wifi_recon.go
index 3d50c432..f8eb5b5d 100644
--- a/modules/wifi_recon.go
+++ b/modules/wifi_recon.go
@@ -321,7 +321,7 @@ func (w *WiFiRecon) startDeauth() error {
 	return errors.New("No base station or client set.")
 }
 
-func (w *WiFiRecon) discoverAccessPoints(radiotapLayer gopacket.Layer, dot11 *layers.Dot11, packet gopacket.Packet) {
+func (w *WiFiRecon) discoverAccessPoints(radiotap *layers.RadioTap, dot11 *layers.Dot11, packet gopacket.Packet) {
 	dot11infoLayer := packet.Layer(layers.LayerTypeDot11InformationElement)
 	if dot11infoLayer == nil {
 		return
@@ -335,7 +335,6 @@ func (w *WiFiRecon) discoverAccessPoints(radiotapLayer gopacket.Layer, dot11 *la
 	dst := dot11.Address1
 	// packet sent to broadcast mac with a SSID set?
 	if bytes.Compare(dst, network.BroadcastMac) == 0 && len(dot11info.Info) > 0 {
-		radiotap, _ := radiotapLayer.(*layers.RadioTap)
 		ssid := string(dot11info.Info)
 		bssid := dot11.Address3.String()
 		channel := mhz2chan(int(radiotap.ChannelFrequency))
@@ -343,7 +342,7 @@ func (w *WiFiRecon) discoverAccessPoints(radiotapLayer gopacket.Layer, dot11 *la
 	}
 }
 
-func (w *WiFiRecon) discoverClients(radiotapLayer gopacket.Layer, dot11 *layers.Dot11, ap net.HardwareAddr, packet gopacket.Packet) {
+func (w *WiFiRecon) discoverClients(radiotap *layers.RadioTap, dot11 *layers.Dot11, ap net.HardwareAddr, packet gopacket.Packet) {
 	// only check data packets of connected stations
 	if dot11.Type.MainType() != layers.Dot11TypeData {
 		return
@@ -352,7 +351,6 @@ func (w *WiFiRecon) discoverClients(radiotapLayer gopacket.Layer, dot11 *layers.
 	bssid := dot11.Address1
 	// packet going to this specific BSSID?
 	if bytes.Compare(bssid, ap) == 0 {
-		radiotap, _ := radiotapLayer.(*layers.RadioTap)
 		src := dot11.Address2
 		channel := mhz2chan(int(radiotap.ChannelFrequency))
 		w.wifi.AddIfNew("", src.String(), false, channel)
@@ -386,29 +384,17 @@ func (w *WiFiRecon) Start() error {
 				break
 			}
 
-			radiotapLayer := packet.Layer(layers.LayerTypeRadioTap)
-			if radiotapLayer == nil {
-				continue
-			}
-
-			dot11Layer := packet.Layer(layers.LayerTypeDot11)
-			if dot11Layer == nil {
-				continue
-			}
-
-			dot11, ok := dot11Layer.(*layers.Dot11)
-			if ok == false {
-				continue
-			}
-
-			w.updateStats(dot11, packet)
-
-			if w.isApSelected() == false {
+			// perform initial dot11 parsing and layers validation
+			if ok, radiotap, dot11 := packets.Dot11Parse(packet); ok == true {
+				// keep collecting traffic statistics
+				w.updateStats(dot11, packet)
 				// no access point bssid selected, keep scanning for other aps
-				w.discoverAccessPoints(radiotapLayer, dot11, packet)
-			} else {
-				// discover stations connected to the selected access point bssid
-				w.discoverClients(radiotapLayer, dot11, w.accessPoint, packet)
+				if w.isApSelected() == false {
+					w.discoverAccessPoints(radiotap, dot11, packet)
+				} else {
+					// discover stations connected to the selected access point bssid
+					w.discoverClients(radiotap, dot11, w.accessPoint, packet)
+				}
 			}
 		}
 	})
diff --git a/packets/dot11.go b/packets/dot11.go
index 20c90586..3b8e773f 100644
--- a/packets/dot11.go
+++ b/packets/dot11.go
@@ -3,6 +3,7 @@ package packets
 import (
 	"net"
 
+	"github.com/google/gopacket"
 	"github.com/google/gopacket/layers"
 )
 
@@ -27,3 +28,27 @@ func NewDot11Deauth(a1 net.HardwareAddr, a2 net.HardwareAddr, a3 net.HardwareAdd
 		&deauth,
 	)
 }
+
+func Dot11Parse(packet gopacket.Packet) (ok bool, radiotap *layers.RadioTap, dot11 *layers.Dot11) {
+	ok = false
+	radiotap = nil
+	dot11 = nil
+
+	radiotapLayer := packet.Layer(layers.LayerTypeRadioTap)
+	if radiotapLayer == nil {
+		return
+	}
+	radiotap, ok = radiotapLayer.(*layers.RadioTap)
+	if ok == false || radiotap == nil {
+		return
+	}
+
+	dot11Layer := packet.Layer(layers.LayerTypeDot11)
+	if dot11Layer == nil {
+		ok = false
+		return
+	}
+
+	dot11, ok = dot11Layer.(*layers.Dot11)
+	return
+}