2 files changed, 72 insertions, 14 deletions

diff --git a/dhcpv6/ztpv6/parse_vendor_options.go b/dhcpv6/ztpv6/parse_vendor_options.go
index 63a6a05..212733c 100644
--- a/dhcpv6/ztpv6/parse_vendor_options.go
+++ b/dhcpv6/ztpv6/parse_vendor_options.go
@@ -16,24 +16,36 @@ type VendorData struct {
 	VendorName, Model, Serial string
 }
 
-// ParseVendorData will try to parse dhcp6 Vendor Specific Information options data
-// looking for more specific vendor data (like model, serial number, etc).
-// If the options are missing we will just return nil
+// ParseVendorData will try to parse dhcp6 Vendor Specific Information options
+// ( 16 and 17) data looking for more specific vendor data (like model, serial
+// number, etc). If the options are missing we will just return nil
 func ParseVendorData(packet dhcpv6.DHCPv6) (*VendorData, error) {
-	opt := packet.GetOneOption(dhcpv6.OptionVendorOpts)
-	if opt == nil {
-		return nil, errors.New("vendor options not found")
+	// check for both options 16 and 17 if both are present will use opt 17
+	opt16 := packet.GetOneOption(dhcpv6.OptionVendorClass)
+	opt17 := packet.GetOneOption(dhcpv6.OptionVendorOpts)
+	if (opt16 == nil) && (opt17 == nil) {
+		return nil, errors.New("no vendor options or vendor class found")
 	}
 
 	vd := VendorData{}
-	vo := opt.(*dhcpv6.OptVendorOpts).VendorOpts
+	vData := []string{}
 
-	for _, opt := range vo {
-		optData := string(opt.(*dhcpv6.OptionGeneric).OptionData)
+	if opt17 != nil {
+		vo := opt17.(*dhcpv6.OptVendorOpts).VendorOpts
+		for _, opt := range vo {
+			vData = append(vData, string(opt.(*dhcpv6.OptionGeneric).OptionData))
+		}
+	} else {
+		data := opt16.(*dhcpv6.OptVendorClass).Data
+		for _, d := range data {
+			vData = append(vData, string(d))
+		}
+	}
+	for _, d := range vData {
 		switch {
 		// Arista;DCS-0000;00.00;ZZZ00000000
-		case strings.HasPrefix(optData, "Arista;"):
-			p := strings.Split(optData, ";")
+		case strings.HasPrefix(d, "Arista;"):
+			p := strings.Split(d, ";")
 			if len(p) < 4 {
 				return nil, errVendorOptionMalformed
 			}
@@ -44,8 +56,8 @@ func ParseVendorData(packet dhcpv6.DHCPv6) (*VendorData, error) {
 			return &vd, nil
 
 		// ZPESystems:NSC:000000000
-		case strings.HasPrefix(optData, "ZPESystems:"):
-			p := strings.Split(optData, ":")
+		case strings.HasPrefix(d, "ZPESystems:"):
+			p := strings.Split(d, ":")
 			if len(p) < 3 {
 				return nil, errVendorOptionMalformed
 			}
diff --git a/dhcpv6/ztpv6/parse_vendor_options_test.go b/dhcpv6/ztpv6/parse_vendor_options_test.go
index 083866c..45b42d6 100644
--- a/dhcpv6/ztpv6/parse_vendor_options_test.go
+++ b/dhcpv6/ztpv6/parse_vendor_options_test.go
@@ -7,7 +7,7 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestParseVendorData(t *testing.T) {
+func TestParseVendorDataWithVendorOpts(t *testing.T) {
 	tt := []struct {
 		name string
 		vc   string
@@ -53,3 +53,49 @@ func TestParseVendorData(t *testing.T) {
 		})
 	}
 }
+
+func TestParseVendorDataWithVendorClass(t *testing.T) {
+	tt := []struct {
+		name string
+		vc   string
+		want *VendorData
+		fail bool
+	}{
+		{name: "empty", fail: true},
+		{name: "unknownVendor", vc: "VendorX;BFR10K;XX12345", fail: true, want: nil},
+		{name: "truncatedArista", vc: "Arista;1234", fail: true, want: nil},
+		{name: "truncatedZPE", vc: "ZPESystems:1234", fail: true, want: nil},
+		{
+			name: "arista",
+			vc:   "Arista;DCS-7050S-64;01.23;JPE12345678",
+			want: &VendorData{VendorName: "Arista", Model: "DCS-7050S-64", Serial: "JPE12345678"},
+		}, {
+			name: "zpe",
+			vc:   "ZPESystems:NSC:001234567",
+			want: &VendorData{VendorName: "ZPESystems", Model: "NSC", Serial: "001234567"},
+		},
+	}
+
+	for _, tc := range tt {
+		t.Run(tc.name, func(t *testing.T) {
+			packet, err := dhcpv6.NewMessage()
+			if err != nil {
+				t.Fatalf("failed to creat dhcpv6 packet object: %v", err)
+			}
+
+			packet.AddOption(&dhcpv6.OptVendorClass{
+				EnterpriseNumber: 0000, Data: [][]byte{[]byte(tc.vc)}})
+
+			vd, err := ParseVendorData(packet)
+			if err != nil && !tc.fail {
+				t.Errorf("unexpected failure: %v", err)
+			}
+
+			if vd != nil {
+				require.Equal(t, *tc.want, *vd, "comparing vendor class data")
+			} else {
+				require.Equal(t, tc.want, vd, "comparing vendor class data")
+			}
+		})
+	}
+}