Add support for partial address search

Author: LucasG0

backend/infrahub/core/utils.py

@@ -191,10 +191,66 @@ def convert_ip_to_binary_str(
     return ip_bin.zfill(obj.max_prefixlen)
 
 
-def collapse_ipv6_address_or_network(address_or_network: str) -> str:
-    if "/" in address_or_network:
-        return ipaddress.IPv6Network(address_or_network).with_prefixlen
-    return str(ipaddress.IPv6Address(address_or_network))
+def collapse_ipv6(s: str) -> str:
+    """Collapse an ipv6 address, ipv6 network, or a partial ipv6 address. Raises an error if input does not resemble an IPv6 address."""
+
+    try:
+        return str(ipaddress.IPv6Address(s))
+    except ipaddress.AddressValueError:
+        pass
+
+    try:
+        return ipaddress.IPv6Network(s).with_prefixlen
+    except ipaddress.AddressValueError:
+        pass
+
+    # Input string might be an incomplete address in IPv6 format,
+    # in which case we would like the collapsed form equivalent of this incomplete address for matching purposes.
+    # To get it, we first try to pad the incomplete address with zeros, then we retrieve the collapsed form
+    # of the full address, and we remove extra "::" or ":0" at the end of it.
+
+    error_message = "Input string does not match IPv6 format"
+
+    if "::" in s:
+        raise ValueError(error_message)
+
+    # Add padding to complete the address if needed
+    segments = s.split(":")
+
+    if len(segments) == 0:
+        raise ValueError(error_message)
+
+    # If any of the non-last segments has less than 4 characters it means we deal with
+    # a IPv6 collapsed form or an invalid address
+    for segment in segments[:-1]:
+        if len(segment) != 4:
+            raise ValueError(error_message)
+
+    # Add 0 padding to last segment
+    if len(segments[-1]) > 4:
+        raise ValueError(error_message)
+
+    segments[-1] += "0" * (4 - len(segments[-1]))
+
+    # Complete the address to have 8 segments by padding with zeros
+    while len(segments) < 8:
+        segments.append("0000")
+
+    # Create a full IPv6 address from the partial input
+    full_address = ":".join(segments)
+
+    # Create an IPv6Address object for validation and to build IPv6 collapsed form.
+    ipv6_address = ipaddress.IPv6Address(full_address)
+
+    compressed_address = ipv6_address.compressed
+
+    # We padded with zeros so address might endswith "::" or ":0".
+    if compressed_address.endswith(("::", ":0")):
+        return compressed_address[:-2]
+
+    # Otherwise, it means 8th segment of ipv6 address was not full and not composed of 0 only
+    # e.g. 2001:0db8:0000:0000:0000:0000:03
+    return compressed_address
 
 
 # --------------------------------------------------------------------------------

backend/infrahub/graphql/queries/search.py

@@ -7,7 +7,7 @@
 
 from infrahub.core.constants import InfrahubKind
 from infrahub.core.manager import NodeManager
-from infrahub.core.utils import collapse_ipv6_address_or_network
+from infrahub.core.utils import collapse_ipv6
 
 if TYPE_CHECKING:
     from graphql import GraphQLResolveInfo
@@ -51,8 +51,8 @@ async def search_resolver(
             result.append(matching)
     else:
         try:
-            # Convert any IPv6 address/network to collapsed format as it might be stored in db.
-            q = collapse_ipv6_address_or_network(q)
+            # Convert any IPv6 address, network or partial address to collapsed format as it might be stored in db.
+            q = collapse_ipv6(q)
         except ValueError:
             pass
 

backend/tests/unit/graphql/queries/test_search.py

@@ -1,7 +1,9 @@
+import pytest
 from graphql import graphql
 
 from infrahub.core.branch import Branch
 from infrahub.core.node import Node
+from infrahub.core.utils import collapse_ipv6
 from infrahub.database import InfrahubDatabase
 from infrahub.graphql.initialization import prepare_graphql_params
 
@@ -170,11 +172,60 @@ async def test_search_ipv6_network_extended_format(
     )
 
 
+async def test_search_ipv6_partial_address(
+    db: InfrahubDatabase,
+    ip_dataset_01,
+    branch: Branch,
+):
+    gql_params = prepare_graphql_params(db=db, include_subscription=False, branch=branch)
+
+    res_two_segments = await graphql(
+        schema=gql_params.schema,
+        source=SEARCH_QUERY,
+        context_value=gql_params.context,
+        root_value=None,
+        variable_values={"search": "2001:0db8"},
+    )
+
+    res_partial_segment_1 = await graphql(
+        schema=gql_params.schema,
+        source=SEARCH_QUERY,
+        context_value=gql_params.context,
+        root_value=None,
+        variable_values={"search": "2001:0db8:0"},
+    )
+
+    res_partial_segment_2 = await graphql(
+        schema=gql_params.schema,
+        source=SEARCH_QUERY,
+        context_value=gql_params.context,
+        root_value=None,
+        variable_values={"search": "2001:0db8:0000:0"},
+    )
+
+    assert (
+        res_two_segments.data["InfrahubSearchAnywhere"]["count"]
+        == res_partial_segment_1.data["InfrahubSearchAnywhere"]["count"]
+        == res_partial_segment_2.data["InfrahubSearchAnywhere"]["count"]
+        == 2
+    )
+
+    assert (
+        res_two_segments.data["InfrahubSearchAnywhere"]["edges"][0]["node"]["id"]
+        == res_partial_segment_1.data["InfrahubSearchAnywhere"]["edges"][0]["node"]["id"]
+        == res_partial_segment_2.data["InfrahubSearchAnywhere"]["edges"][0]["node"]["id"]
+    )
+
+
 async def test_search_ipv4(
     db: InfrahubDatabase,
     ip_dataset_01,
     branch: Branch,
 ):
+    """
+    This only tests that ipv6 search specific behavior does not break ipv4 search.
+    """
+
     gql_params = prepare_graphql_params(db=db, include_subscription=False, branch=branch)
 
     result_address = await graphql(
@@ -203,3 +254,37 @@ async def test_search_ipv4(
         result_address.data["InfrahubSearchAnywhere"]["edges"][0]["node"]["id"]
         == result_network.data["InfrahubSearchAnywhere"]["edges"][0]["node"]["id"]
     )
+
+
+@pytest.mark.parametrize(
+    "query,expected",
+    [
+        ("2001:0db8:0000:0000:0000:0000:0000:0000/48", "2001:db8::/48"),
+        ("2001:0db8:0000:0000:0000:0000:0000:0000", "2001:db8::"),
+        ("2001:0db8", "2001:db8"),
+        ("2001:0db8:0", "2001:db8"),
+        ("2001:0db8:0000", "2001:db8"),
+        ("2001:0db8:0000:0", "2001:db8"),
+        ("2001:0db8:0000:0000:00", "2001:db8"),
+        ("2001:0db8:0000:0001:00", "2001:db8:0:1"),
+        ("2001:0db8:0001:0002:00", "2001:db8:1:2"),
+        ("2001:0db8:0001:0000:0002:0000:0003", "2001:db8:1:0:2:0:3"),
+    ],
+)
+def test_collapse_ipv6_address_or_network(query, expected):
+    assert collapse_ipv6(query) == expected
+
+
+@pytest.mark.parametrize(
+    "query",
+    [
+        "invalid",
+        "invalid:case",
+        "2001:invalid",
+        "2001:0db81:0000",
+        "10.0.0.0",
+    ],
+)
+def test_collapse_ipv6_address_or_network_invalid_cases(query):
+    with pytest.raises(ValueError):
+        collapse_ipv6(query)