Extend checks for IPv6 addresses

- Make sure enclosing brackets match.
- Use the built-in classes/checks for the IPv6/IPv4 address format.
- Extend tests to bad cases

Signed-off-by: Ivan Kanakarakis <ivan.kanak@gmail.com>

2 files changed, 33 insertions, 40 deletions

diff --git a/src/saml2/validate.py b/src/saml2/validate.py
index c6caf47d..26de5472 100644
--- a/src/saml2/validate.py
+++ b/src/saml2/validate.py
@@ -4,6 +4,9 @@ import re
 import struct
 import base64
 import time
+from ipaddress import AddressValueError
+from ipaddress import IPv4Address
+from ipaddress import IPv6Address
 
 from saml2 import time_util
 
@@ -112,57 +115,30 @@ def validate_before(not_before, slack):
 
 
 def valid_address(address):
+    """Validate IPv4/IPv6 addresses."""
     if not (valid_ipv4(address) or valid_ipv6(address)):
         raise NotValid("address")
     return True
 
 
 def valid_ipv4(address):
-    parts = address.split(".")
-    if len(parts) != 4:
+    """Validate IPv4 addresses."""
+    try:
+        IPv4Address(address)
+    except AddressValueError:
         return False
-    for item in parts:
-        try:
-            if not 0 <= int(item) <= 255:
-                raise NotValid("ipv4")
-        except ValueError:
-            return False
     return True
 
-#
-IPV6_PATTERN = re.compile(r"""
-    ^
-    \s*                         # Leading whitespace
-    \[?                         # See https://tools.ietf.org/html/rfc4038#section-5.1
-    (?!.*::.*::)                # Only a single wildcard allowed
-    (?:(?!:)|:(?=:))            # Colon iff it would be part of a wildcard
-    (?:                         # Repeat 6 times:
-        [0-9a-f]{0,4}           #   A group of at most four hexadecimal digits
-        (?:(?<=::)|(?<!::):)    #   Colon unless preceeded by wildcard
-    ){6}                        #
-    (?:                         # Either
-        [0-9a-f]{0,4}           #   Another group
-        (?:(?<=::)|(?<!::):)    #   Colon unless preceeded by wildcard
-        [0-9a-f]{0,4}           #   Last group
-        (?: (?<=::)             #   Colon iff preceeded by exacly one colon
-         |  (?<!:)              #
-         |  (?<=:) (?<!::) :    #
-         )                      # OR
-     |                          #   A v4 address with NO leading zeros
-        (?:25[0-4]|2[0-4]\d|1\d\d|[1-9]?\d)
-        (?: \.
-            (?:25[0-4]|2[0-4]\d|1\d\d|[1-9]?\d)
-        ){3}
-    )
-    \]?                         # See https://tools.ietf.org/html/rfc4038#section-5.1
-    \s*                         # Trailing whitespace
-    $
-""", re.VERBOSE | re.IGNORECASE | re.DOTALL)
-
 
 def valid_ipv6(address):
-    """Validates IPv6 addresses. """
-    return IPV6_PATTERN.match(address) is not None
+    """Validate IPv6 addresses."""
+    is_enclosed_in_brackets = address.startswith("[") and address.endswith("]")
+    address_raw = address[1:-1] if is_enclosed_in_brackets else address
+    try:
+        IPv6Address(address_raw)
+    except AddressValueError:
+        return False
+    return True
 
 
 def valid_boolean(val):
diff --git a/tests/test_13_validate.py b/tests/test_13_validate.py
index 9c91299a..3c13e2c1 100644
--- a/tests/test_13_validate.py
+++ b/tests/test_13_validate.py
@@ -121,10 +121,27 @@ def test_valid_anytype():
     assert valid_anytype("P1Y2M3DT10H30M")
     assert valid_anytype("urn:oasis:names:tc:SAML:2.0:attrname-format:uri")
 
+
 def test_valid_address():
     assert valid_address("130.239.16.3")
     assert valid_address("2001:8003:5555:9999:555a:5555:c77:d5c5")
+    assert valid_address("2001:8003:5555::555a:5555:c77:d5c5")
 
     # See https://tools.ietf.org/html/rfc4038#section-5.1 regarding
     # the inclusion of brackets in the ipv6 address below.
     assert valid_address("[2001:8003:5555:9999:555a:5555:c77:d5c5]")
+
+    with raises(NotValid):
+        assert valid_address("127.0.0.256")
+    with raises(NotValid):
+        assert valid_address("127.0.0.")
+    with raises(NotValid):
+        assert valid_address("127.0.0")
+    with raises(NotValid):
+        assert valid_address("2001::5555:9999::5555:c77:d5c5]")
+    with raises(NotValid):
+        assert valid_address("2001:8003:5555:9999:555a:5555:c77:d5c5]")
+    with raises(NotValid):
+        assert valid_address("[2001:8003:5555:9999:555a:5555:c77:d5c5")
+    with raises(NotValid):
+        assert valid_address("[[2001:8003:5555:9999:555a:5555:c77:d5c5]")