From 339c8f319d2df1e64322d48c691c56912887aff8 Mon Sep 17 00:00:00 2001
From: chris marget <5470536+chrismarget@users.noreply.github.com>
Date: Fri, 20 Sep 2019 15:45:46 -0400
Subject: [PATCH] added functions to pad mac strings for net.ParseMAC()

---
 arp_linux.go   |  2 +-
 arp_unix.go    |  3 +--
 arp_windows.go |  2 +-
 util.go        | 59 ++++++++++++++++++++++++++++++++++++++++++++++++++
 util_test.go   | 56 +++++++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 118 insertions(+), 4 deletions(-)
 create mode 100644 util.go
 create mode 100644 util_test.go

diff --git a/arp_linux.go b/arp_linux.go
index 63cfff6..f709b42 100644
--- a/arp_linux.go
+++ b/arp_linux.go
@@ -34,7 +34,7 @@ func Table() ArpTable {
 	for s.Scan() {
 		line := s.Text()
 		fields := strings.Fields(line)
-		table[fields[f_IPAddr]] = fields[f_HWAddr]
+		table[fields[f_IPAddr]] = padMacString(fields[f_HWAddr])
 	}
 
 	return table
diff --git a/arp_unix.go b/arp_unix.go
index 51b11e5..5823338 100644
--- a/arp_unix.go
+++ b/arp_unix.go
@@ -30,8 +30,7 @@ func Table() ArpTable {
 		// strip brackets around IP
 		ip := strings.Replace(fields[1], "(", "", -1)
 		ip = strings.Replace(ip, ")", "", -1)
-
-		table[ip] = fields[3]
+		table[ip] = padMacString(fields[3])
 	}
 
 	return table
diff --git a/arp_windows.go b/arp_windows.go
index 14fc65f..6478d4c 100644
--- a/arp_windows.go
+++ b/arp_windows.go
@@ -42,7 +42,7 @@ func Table() ArpTable {
 
 		ip := fields[0]
 		// Normalize MAC address to colon-separated format
-		table[ip] = strings.Replace(fields[1], "-", ":", -1)
+		table[ip] = padMacString(strings.Replace(fields[1], "-", ":", -1))
 	}
 
 	return table
diff --git a/util.go b/util.go
new file mode 100644
index 0000000..fe15072
--- /dev/null
+++ b/util.go
@@ -0,0 +1,59 @@
+package arp
+
+import (
+	"fmt"
+	"strings"
+)
+
+const numberChars = "ABCDEFabcdef0123456789"
+
+// onlyValidChars returns true if string "test" consists entirely of
+// characters from string "expected". If any other characters it returns
+// false.
+func onlyValidChars(test string, expected string) bool {
+	for _, char := range test {
+		if !strings.Contains(expected, string(char)) {
+			return false
+		}
+	}
+	return true
+}
+
+// padMacString takes MAC addresses in string form, pads them to comply with
+// the expectations of the standard library's net.ParseMAC(). For example,
+// "0:0:c:7:ac:0" becomes "00:00:0c:07:ac:00". If input string cannot be
+// understood/padded, then the string is returned without modification.
+func padMacString(in string) string {
+	var sep string
+	var pad string
+
+	switch {
+	case strings.Contains(in, ":"):
+		sep = ":"
+		pad = "%2s"
+		if !onlyValidChars(in, numberChars+sep) {
+			return in
+		}
+	case strings.Contains(in, "-"):
+		sep = "-"
+		pad = "%2s"
+		if !onlyValidChars(in, numberChars+sep) {
+			return in
+		}
+	case strings.Contains(in, "."):
+		sep = "."
+		pad = "%4s"
+		if !onlyValidChars(in, numberChars+sep) {
+			return in
+		}
+	default:
+		return in
+	}
+
+	s := strings.Split(in, sep)
+	for i := range s {
+		s[i] = strings.Replace(fmt.Sprintf(pad, s[i]), " ", "0", -1)
+	}
+
+	return strings.Join(s, sep)
+}
diff --git a/util_test.go b/util_test.go
new file mode 100644
index 0000000..c669a95
--- /dev/null
+++ b/util_test.go
@@ -0,0 +1,56 @@
+package arp
+
+import "testing"
+
+func TestPadMacString(t *testing.T) {
+	var validInput []string
+	var expectedResult []string
+
+	validInput = append(validInput,"0:0:0:0:0:0")
+	expectedResult = append(expectedResult, "00:00:00:00:00:00")
+	validInput = append(validInput,"00:00:00:00:00:00")
+	expectedResult = append(expectedResult, "00:00:00:00:00:00")
+	validInput = append(validInput,"a:b:c:D:E:F")
+	expectedResult = append(expectedResult, "0a:0b:0c:0D:0E:0F")
+	validInput = append(validInput,"0-0-0-0-0-0")
+	expectedResult = append(expectedResult, "00-00-00-00-00-00")
+	validInput = append(validInput,"00-00-00-00-00-00")
+	expectedResult = append(expectedResult, "00-00-00-00-00-00")
+	validInput = append(validInput,"a-b-c-D-E-F")
+	expectedResult = append(expectedResult, "0a-0b-0c-0D-0E-0F")
+	validInput = append(validInput,"0.00.000")
+	expectedResult = append(expectedResult, "0000.0000.0000")
+	validInput = append(validInput,"0000.0000.0000")
+	expectedResult = append(expectedResult, "0000.0000.0000")
+	validInput = append(validInput,"0.0a.abc")
+	expectedResult = append(expectedResult, "0000.000a.0abc")
+
+	for i, _ := range validInput {
+		result := padMacString(validInput[i])
+		if result != expectedResult[i] {
+			t.Fatalf("expectedResult %s, got %s", expectedResult[i], result)
+		}
+	}
+
+	var invalidInput []string
+
+	invalidInput = append(invalidInput, "0000.0000.00:00")
+	invalidInput = append(invalidInput, "0000.0000.00-00")
+	invalidInput = append(invalidInput, "0000.0000:0000")
+	invalidInput = append(invalidInput, "0000.0000-0000")
+	invalidInput = append(invalidInput, "0000.0000.000g")
+	invalidInput = append(invalidInput, "00:00:00:00:00.00")
+	invalidInput = append(invalidInput, "00:00:00:00:00-00")
+	invalidInput = append(invalidInput, "00:00:00:00:00:0g")
+	invalidInput = append(invalidInput, "00-00-00-00-00:00")
+	invalidInput = append(invalidInput, "00-00-00-00-00.00")
+	invalidInput = append(invalidInput, "00-00-00-00-00-0g")
+
+	for i, _ := range invalidInput {
+		result := padMacString(invalidInput[i])
+		if result != invalidInput[i] {
+			t.Fatalf("invalid input %s should have been unmodified but got %s",
+				invalidInput[i], result)
+		}
+	}
+}
\ No newline at end of file