move IP generators into an abstract base class

src/warnet/ip_generator.py

@@ -0,0 +1,45 @@
+from abc import ABC, abstractmethod
+import ipaddress
+import random
+
+class IPV4AddressGenerator(ABC):
+    def generate_ipv4_addr(self, subnet):
+        """
+        Generate a valid random IPv4 address within the given subnet.
+
+        :param subnet: Subnet in CIDR notation (e.g., '100.0.0.0/8')
+        :return: Random IP address within the subnet
+        """
+        reserved_ips = [
+            "0.0.0.0/8",
+            "10.0.0.0/8",
+            "100.64.0.0/10",
+            "127.0.0.0/8",
+            "169.254.0.0/16",
+            "172.16.0.0/12",
+            "192.0.0.0/24",
+            "192.0.2.0/24",
+            "192.88.99.0/24",
+            "192.168.0.0/16",
+            "198.18.0.0/15",
+            "198.51.100.0/24",
+            "203.0.113.0/24",
+            "224.0.0.0/4",
+        ]
+
+        def is_public(ip):
+            for reserved in reserved_ips:
+                if ipaddress.ip_address(ip) in ipaddress.ip_network(reserved, strict=False):
+                    return False
+            return True
+
+        network = ipaddress.ip_network(subnet, strict=False)
+
+        # Generate a random IP within the subnet range
+        while True:
+            ip_int = random.randint(
+                int(network.network_address), int(network.broadcast_address)
+            )
+            ip_str = str(ipaddress.ip_address(ip_int))
+            if is_public(ip_str):
+                return ip_str

src/warnet/lnnode.py

@@ -1,21 +1,21 @@
 
 import json
 import os
+from warnet.ip_generator import IPV4AddressGenerator
 from warnet.utils import (
     exponential_backoff,
-    generate_ipv4_addr
 )
 from backends import BackendInterface, ServiceType
 from .status import RunningStatus
 
-class LNNode:
+class LNNode(IPV4AddressGenerator):
     def __init__(self, warnet, tank, impl, backend: BackendInterface):
         self.warnet = warnet
         self.tank = tank
         assert impl == "lnd"
         self.impl = impl
         self.backend = backend
-        self.ipv4 = generate_ipv4_addr(self.warnet.subnet)
+        self.ipv4 = super().generate_ipv4_addr(self.warnet.subnet)
         self.rpc_port = 10009
 
     @property

src/warnet/tank.py

@@ -6,9 +6,9 @@
 
 from pathlib import Path
 from warnet.lnnode import LNNode
+from warnet.ip_generator import IPV4AddressGenerator
 from warnet.utils import (
     exponential_backoff,
-    generate_ipv4_addr,
     sanitize_tc_netem_command,
     SUPPORTED_TAGS,
 )
@@ -22,7 +22,7 @@
 logger = logging.getLogger("tank")
 
 
-class Tank:
+class Tank(IPV4AddressGenerator):
     DEFAULT_BUILD_ARGS = "--disable-tests --with-incompatible-bdb --without-gui --disable-bench --disable-fuzz-binary --enable-suppress-external-warnings --enable-debug "
 
     def __init__(self, index:int, config_dir: Path, warnet):
@@ -97,7 +97,7 @@ def suffix(self):
     @property
     def ipv4(self):
         if self._ipv4 is None:
-            self._ipv4 = generate_ipv4_addr(self.warnet.subnet)
+            self._ipv4 = self.generate_ipv4_addr(self.warnet.subnet)
         return self._ipv4
 
     @property

src/warnet/utils.py

@@ -78,48 +78,6 @@ def get_architecture():
     return arch
 
 
-def generate_ipv4_addr(subnet):
-    """
-    Generate a valid random IPv4 address within the given subnet.
-
-    :param subnet: Subnet in CIDR notation (e.g., '100.0.0.0/8')
-    :return: Random IP address within the subnet
-    """
-    reserved_ips = [
-        "0.0.0.0/8",
-        "10.0.0.0/8",
-        "100.64.0.0/10",
-        "127.0.0.0/8",
-        "169.254.0.0/16",
-        "172.16.0.0/12",
-        "192.0.0.0/24",
-        "192.0.2.0/24",
-        "192.88.99.0/24",
-        "192.168.0.0/16",
-        "198.18.0.0/15",
-        "198.51.100.0/24",
-        "203.0.113.0/24",
-        "224.0.0.0/4",
-    ]
-
-    def is_public(ip):
-        for reserved in reserved_ips:
-            if ipaddress.ip_address(ip) in ipaddress.ip_network(reserved, strict=False):
-                return False
-        return True
-
-    network = ipaddress.ip_network(subnet, strict=False)
-
-    # Generate a random IP within the subnet range
-    while True:
-        ip_int = random.randint(
-            int(network.network_address), int(network.broadcast_address)
-        )
-        ip_str = str(ipaddress.ip_address(ip_int))
-        if is_public(ip_str):
-            return ip_str
-
-
 def sanitize_tc_netem_command(command: str) -> bool:
     """
     Sanitize the tc-netem command to ensure it's valid and safe to execute, as we run it as root on a container.