support decoding for different vendor formats

tests/test_parse_header.py

@@ -3,6 +3,7 @@
 
 import struct
 from ipaddress import IPv4Address, IPv4Network, IPv6Address, IPv6Network
+from itertools import product
 
 import pytest
 
@@ -114,6 +115,34 @@ TAGGED_ATTRIBUTES = [
     (b"\x71\x07\x01\x00\x40\x20\x03", IPv6Network("2003::0/64")),
 ]
 
+VENDOR_FORMAT_COMBINATIONS = [
+    (1234, 1, 0),
+    (1235, 1, 1),
+    (1236, 1, 2),
+    (1237, 2, 0),
+    (1238, 2, 1),
+    (1239, 2, 2),
+    (1240, 4, 0),
+    (1241, 4, 1),
+    (1242, 4, 2),
+]
+
+
+def flattened_product(l1, l2):
+    result = []
+    for combination in product(l1, l2):
+        left, right = combination
+        result.append(list(left) + list(right))
+    return result
+
+
+VENDOR_TEST_ATTRIBUTES = flattened_product(
+    VENDOR_FORMAT_COMBINATIONS, TEST_ATTRIBUTES
+)
+VENDOR_TAGGED_ATTRIBUTES = flattened_product(
+    VENDOR_FORMAT_COMBINATIONS, TAGGED_ATTRIBUTES
+)
+
 
 @pytest.fixture
 def radius_dictionary():
@@ -143,18 +172,6 @@ def test_decode_attribute_rfc(radius_dictionary, attr_bytes, expected):
     assert attrs[0].tag == 0
 
 
-@pytest.mark.parametrize("attr_bytes, expected", TEST_ATTRIBUTES)
-def test_decode_attribute_vsa(radius_dictionary, attr_bytes, expected):
-    vsa_length = (6 + len(attr_bytes)).to_bytes(1, "big")
-    raw_packet = (
-        bytes(20) + b"\x1a" + vsa_length + b"\x00\x00\x04\xd2" + attr_bytes
-    )
-    attrs = utils.decode_attributes(radius_dictionary, raw_packet)
-    assert len(attrs) == 1
-    assert attrs[0].value == expected
-    assert attrs[0].tag == 0
-
-
 @pytest.mark.parametrize("attr_bytes, expected", TAGGED_ATTRIBUTES)
 def test_decode_attribute_rfc_tagged(radius_dictionary, attr_bytes, expected):
     raw_packet = bytes(20) + attr_bytes
@@ -164,12 +181,50 @@ def test_decode_attribute_rfc_tagged(radius_dictionary, attr_bytes, expected):
     assert attrs[0].tag == 1
 
 
-@pytest.mark.parametrize("attr_bytes, expected", TAGGED_ATTRIBUTES)
-def test_decode_attribute_vsa_tagged(radius_dictionary, attr_bytes, expected):
-    vsa_length = (6 + len(attr_bytes)).to_bytes(1, "big")
-    raw_packet = (
-        bytes(20) + b"\x1a" + vsa_length + b"\x00\x00\x04\xd2" + attr_bytes
+def generate_vendor_attribute(vendor_id, tlen, llen, attr_bytes):
+    vendor_id = vendor_id.to_bytes(4, "big")
+    vsa_length = (4 + tlen + llen + len(attr_bytes)).to_bytes(1, "big")
+    attr_type = attr_bytes[0].to_bytes(tlen, "big")
+
+    attr_len = attr_bytes[1] + tlen
+    if llen == 0:
+        attr_len = b""
+    else:
+        attr_len = attr_len.to_bytes(llen, "big")
+
+    attr_bytes = attr_bytes[2:]
+    packet = (
+        bytes(20)
+        + b"\x1a"
+        + vsa_length
+        + vendor_id
+        + attr_type
+        + attr_len
+        + attr_bytes
     )
+    return packet
+
+
+@pytest.mark.parametrize(
+    "vendor_id, tlen, llen, attr_bytes, expected", VENDOR_TEST_ATTRIBUTES
+)
+def test_decode_attribute_vsa(
+    radius_dictionary, vendor_id, tlen, llen, attr_bytes, expected
+):
+    raw_packet = generate_vendor_attribute(vendor_id, tlen, llen, attr_bytes)
+    attrs = utils.decode_attributes(radius_dictionary, raw_packet)
+    assert len(attrs) == 1
+    assert attrs[0].value == expected
+    assert attrs[0].tag == 0
+
+
+@pytest.mark.parametrize(
+    "vendor_id, tlen, llen, attr_bytes, expected", VENDOR_TAGGED_ATTRIBUTES
+)
+def test_decode_attribute_vsa_tagged(
+    radius_dictionary, vendor_id, tlen, llen, attr_bytes, expected
+):
+    raw_packet = generate_vendor_attribute(vendor_id, tlen, llen, attr_bytes)
     attrs = utils.decode_attributes(radius_dictionary, raw_packet)
     assert len(attrs) == 1
     assert attrs[0].value == expected