From cac366913a123f4526c81468b9032c69f1c4c22f Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Mon, 28 Jan 2019 12:12:58 +0100
Subject: [PATCH] new: parsing wps wfa version2 field

---
 packets/dot11_wps.go | 114 +++++++++++++++++++++++++++++++------------
 1 file changed, 83 insertions(+), 31 deletions(-)

diff --git a/packets/dot11_wps.go b/packets/dot11_wps.go
index cba87fab..e956353b 100644
--- a/packets/dot11_wps.go
+++ b/packets/dot11_wps.go
@@ -28,7 +28,10 @@ type wpsAttr struct {
 
 var (
 	wpsSignatureBytes = []byte{0x00, 0x50, 0xf2, 0x04}
-	wpsAttributes     = map[uint16]wpsAttr{
+	wfaExtensionBytes = []byte{0x00, 0x37, 0x2a}
+	wpsVersion2ID     = uint8(0x00)
+
+	wpsAttributes = map[uint16]wpsAttr{
 		0x104A: wpsAttr{Name: "Version", Desc: map[string]string{
 			"10": "1.0",
 			"11": "1.1",
@@ -37,8 +40,6 @@ var (
 			"01": "Not Configured",
 			"02": "Configured",
 		}},
-		0x1057: wpsAttr{Name: "AP Setup Locked"},
-		0x1041: wpsAttr{Name: "Selected Registrar"},
 		0x1012: wpsAttr{Name: "Device Password ID", Desc: map[string]string{
 			"0000": "Pin",
 			"0004": "PushButton",
@@ -49,6 +50,9 @@ var (
 			"02": "Registrar",
 			"03": "AP",
 		}},
+		0x1049: wpsAttr{Name: "Vendor Extension", Func: dot11ParseWPSVendorExtension},
+		0x1057: wpsAttr{Name: "AP Setup Locked"},
+		0x1041: wpsAttr{Name: "Selected Registrar"},
 		0x1047: wpsAttr{Name: "UUID-E"},
 		0x1021: wpsAttr{Name: "Manufacturer", Type: wpsStr},
 		0x1023: wpsAttr{Name: "Model Name", Type: wpsStr},
@@ -61,7 +65,6 @@ var (
 		0x103C: wpsAttr{Name: "RF Bands", Func: dott11ParseWPSBands},
 		0x1045: wpsAttr{Name: "SSID", Type: wpsStr},
 		0x102D: wpsAttr{Name: "OS Version", Type: wpsStr},
-		0x1049: wpsAttr{Name: "Vendor Extension"},
 	}
 
 	wpsConfigs = map[uint16]string{
@@ -120,41 +123,90 @@ func dot11ParseWPSConfigMethods(data []byte) string {
 	return hex.EncodeToString(data)
 }
 
+func dot11ParseWPSVendorExtension(data []byte) string {
+	if len(data) > 3 && bytes.Equal(data[0:3], wfaExtensionBytes) {
+		size := len(data)
+		for offset := 3; offset < size; {
+			idByte := uint8(data[offset])
+			sizeByte := uint8(data[offset+1])
+			if idByte == wpsVersion2ID {
+				return fmt.Sprintf("version2=0x%x", data[offset+2])
+			}
+			offset += int(sizeByte) + 2
+		}
+	}
+	return hex.EncodeToString(data)
+}
+
+func wpsUint16At(data []byte, size int, offset *int) (bool, uint16) {
+	if *offset <= size-2 {
+		off := *offset
+		*offset += 2
+		return true, binary.BigEndian.Uint16(data[off:])
+	}
+	// fmt.Printf("uint16At( data(%d), off=%d )\n", size, *offset)
+	return false, 0
+}
+
+func wpsDataAt(data []byte, size int, offset *int, num int) (bool, []byte) {
+	max := size - num
+	if *offset <= max {
+		off := *offset
+		*offset += num
+		return true, data[off : off+num]
+	}
+	// fmt.Printf("dataAt( data(%d), off=%d, num=%d )\n", size, *offset, num)
+	return false, nil
+}
+
+func dot11ParseWPSTag(id uint16, size uint16, data []byte, info *map[string]string) {
+	name := ""
+	val := ""
+
+	if attr, found := wpsAttributes[id]; found {
+		name = attr.Name
+		if attr.Type == wpsStr {
+			val = string(data)
+		} else {
+			val = hex.EncodeToString(data)
+		}
+
+		if attr.Desc != nil {
+			if desc, found := attr.Desc[val]; found {
+				val = desc
+			}
+		}
+
+		if attr.Func != nil {
+			val = attr.Func(data)
+		}
+	} else {
+		name = fmt.Sprintf("0x%X", id)
+		val = hex.EncodeToString(data)
+	}
+
+	(*info)[name] = val
+}
+
 func dot11ParseWPSData(data []byte) (ok bool, info map[string]string) {
 	info = map[string]string{}
 	size := len(data)
 
 	for offset := 0; offset < size; {
-		tagId := binary.BigEndian.Uint16(data[offset:])
-		offset += 2
-		tagLen := binary.BigEndian.Uint16(data[offset:])
-		offset += 2
-		tagData := data[offset : offset+int(tagLen)]
+		ok := false
+		tagId := uint16(0)
+		tagLen := uint16(0)
+		tagData := []byte(nil)
 
-		if attr, found := wpsAttributes[tagId]; found {
-			val := ""
-			if attr.Type == wpsStr {
-				val = string(tagData)
-			} else {
-				val = hex.EncodeToString(tagData)
-			}
-
-			if attr.Desc != nil {
-				if desc, found := attr.Desc[val]; found {
-					val = desc
-				}
-			}
-
-			if attr.Func != nil {
-				val = attr.Func(tagData)
-			}
-
-			info[attr.Name] = val
+		if ok, tagId = wpsUint16At(data, size, &offset); !ok {
+			break
+		} else if ok, tagLen = wpsUint16At(data, size, &offset); !ok {
+			break
+		} else if ok, tagData = wpsDataAt(data, size, &offset, int(tagLen)); !ok {
+			break
 		} else {
-			info[fmt.Sprintf("0x%X", tagId)] = hex.EncodeToString(tagData)
+			dot11ParseWPSTag(tagId, tagLen, tagData, &info)
 		}
-
-		offset += int(tagLen)
 	}
 
 	return true, info