From 7edf2432d8df985146c96c37fabb4333bd66aecc Mon Sep 17 00:00:00 2001
From: =?utf8?q?Ale=C5=A1=20Mr=C3=A1zek?= <ales.mrazek@nic.cz>
Date: Tue, 18 Jul 2023 11:31:33 +0200
Subject: [PATCH] manager: datamodel: custom type for IP address with '!'
 suffix

---
 .../datamodel/network_schema.py               |  3 +-
 .../datamodel/types/__init__.py               |  2 +
 .../datamodel/types/types.py                  | 57 +++++++++++++++++++
 3 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/manager/knot_resolver_manager/datamodel/network_schema.py b/manager/knot_resolver_manager/datamodel/network_schema.py
index aacb2c469..e59fe4137 100644
--- a/manager/knot_resolver_manager/datamodel/network_schema.py
+++ b/manager/knot_resolver_manager/datamodel/network_schema.py
@@ -10,6 +10,7 @@ from knot_resolver_manager.datamodel.types import (
     Int0_65535,
     InterfaceOptionalPort,
     IPAddress,
+    IPAddressEM,
     IPNetwork,
     IPv4Address,
     IPv6Address,
@@ -45,7 +46,7 @@ class AddressRenumberingSchema(ConfigSchema):
     """
 
     source: IPNetwork
-    destination: IPAddress
+    destination: Union[IPAddressEM, IPAddress]
 
 
 class TLSSchema(ConfigSchema):
diff --git a/manager/knot_resolver_manager/datamodel/types/__init__.py b/manager/knot_resolver_manager/datamodel/types/__init__.py
index c100472be..350cf2133 100644
--- a/manager/knot_resolver_manager/datamodel/types/__init__.py
+++ b/manager/knot_resolver_manager/datamodel/types/__init__.py
@@ -14,6 +14,7 @@ from .types import (
     IntNonNegative,
     IntPositive,
     IPAddress,
+    IPAddressEM,
     IPAddressOptionalPort,
     IPAddressPort,
     IPNetwork,
@@ -44,6 +45,7 @@ __all__ = [
     "IntNonNegative",
     "IntPositive",
     "IPAddress",
+    "IPAddressEM",
     "IPAddressOptionalPort",
     "IPAddressPort",
     "IPNetwork",
diff --git a/manager/knot_resolver_manager/datamodel/types/types.py b/manager/knot_resolver_manager/datamodel/types/types.py
index f558da0de..19bc75bb2 100644
--- a/manager/knot_resolver_manager/datamodel/types/types.py
+++ b/manager/knot_resolver_manager/datamodel/types/types.py
@@ -376,6 +376,63 @@ class IPv6Address(BaseValueType):
 IPAddress = Union[IPv4Address, IPv6Address]
 
 
+class IPAddressEM(BaseValueType):
+    """
+    IP address with exclamation mark suffix, e.g. '127.0.0.1!'.
+    """
+
+    _value: str
+    _addr: Union[ipaddress.IPv4Address, ipaddress.IPv6Address]
+
+    def __init__(self, source_value: Any, object_path: str = "/") -> None:
+        super().__init__(source_value)
+        if isinstance(source_value, str):
+            if source_value.endswith("!"):
+                addr, suff = source_value.split("!", 1)
+                if suff != "":
+                    raise ValueError(f"suffix '{suff}' found after '!'.")
+            else:
+                raise ValueError("string does not end with '!'.")
+            try:
+                self._addr: Union[ipaddress.IPv4Address, ipaddress.IPv6Address] = ipaddress.ip_address(addr)
+                self._value = source_value
+            except ValueError as e:
+                raise ValueError("failed to parse IP address.") from e
+        else:
+            raise ValueError(
+                "Unexpected value for a IPv6 address."
+                f" Expected string, got '{source_value}' with type '{type(source_value)}'",
+                object_path,
+            )
+
+    def to_std(self) -> str:
+        return self._value
+
+    def __str__(self) -> str:
+        return self._value
+
+    def __int__(self) -> int:
+        raise ValueError("Can't convert to an integer")
+
+    def __repr__(self) -> str:
+        return f'{type(self).__name__}("{self._value}")'
+
+    def __eq__(self, o: object) -> bool:
+        """
+        Two instances of IPAddressEM are equal when they represent same string.
+        """
+        return isinstance(o, IPAddressEM) and o._value == self._value
+
+    def serialize(self) -> Any:
+        return self._value
+
+    @classmethod
+    def json_schema(cls: Type["IPAddressEM"]) -> Dict[Any, Any]:
+        return {
+            "type": "string",
+        }
+
+
 class IPNetwork(BaseValueType):
     _value: Union[ipaddress.IPv4Network, ipaddress.IPv6Network]
 
-- 
2.47.2