infix/test/infamy/dhcp.py

"""Start a DHCP server in the background"""
import os
import tempfile as tf
import subprocess


class Server:
    config_file = '/tmp/udhcpd.conf'
    leases_file = '/tmp/udhcpd.leases'

    def __init__(self, netns, start='192.168.0.100', end='192.168.0.110',
                 netmask='255.255.255.0', ip=None, router=None, prefix=None,
                 hostname=None, iface="iface"):
        self.process = None
        self.netns = netns
        self.iface = iface
        self._create_files(start, end, netmask, ip, router, prefix, hostname)

    def __del__(self):
        """Clean up config and lease files"""
        try:
            if os.path.exists(self.config_file):
                os.unlink(self.config_file)
            if os.path.exists(self.leases_file):
                os.unlink(self.leases_file)
        except:
            pass

    def __enter__(self):
        self.start()

    def __exit__(self, _, __, ___):
        self.stop()

    def _create_files(self, start, end, netmask, ip, router, prefix, hostname):
        f = open(self.leases_file, "w")
        f.close()

        with open(self.config_file, 'w') as f:
            if ip:
                start=ip
                end=ip
            f.write(f'''# Generated by Infamy DHCP
lease_file {self.leases_file}
interface {self.iface}
start {start}
end {end}
max_leases 1
option subnet {netmask}
option lease 864000
''')
            if router:
                f.write(f"option router {router}\n")
            if prefix and router:
                f.write(f"option staticroutes {prefix} {router}\n")
            if hostname:
                f.write(f"option hostname {hostname}\n")

    def get_pid(self):
        return self.process.pid

    def start(self):
        if not os.path.exists(self.config_file):
            raise Exception("Config file does not exist. Please create it first.")
        cmd = f"udhcpd -f {self.config_file}"
        self.process = self.netns.popen(cmd.split(" "))

    def stop(self):
        if self.process:
            self.process.terminate()
            self.process.wait()
            self.process = None


class Server6Dnsmasq:
    """DHCPv6 server using dnsmasq"""

    def __init__(self, netns, start=None, end=None, dns=None, domain=None,
                 iface="iface", address=None):
        """Initialize DHCPv6 server

        Args:
            netns:  Network namespace to run server in
            start:  Starting IPv6 address for range (e.g., "2001:db8::100")
            end:    Ending IPv6 address for range (e.g., "2001:db8::200")
            dns:    DNS server addresses (list or string)
            domain: DNS search domain
            iface:  Interface to listen on
        """
        self.process = None
        self.netns = netns
        self.iface = iface
        self.config_file = tf.NamedTemporaryFile(mode='w', prefix='dnsmasq6_',
                                                 suffix='.conf', delete=False)
        self.config_path = self.config_file.name
        self.leases_file = tf.NamedTemporaryFile(mode='w', prefix='dnsmasq6_',
                                                 suffix='.leases', delete=False)
        self.leases_path = self.leases_file.name
        self.leases_file.close()

        self.hosts_file = tf.NamedTemporaryFile(mode='w', delete=False)
        self.hosts_path = self.hosts_file.name

        self._create_config(start, end, dns, domain, address)

    def __del__(self):
        try:
            if os.path.exists(self.config_path):
                os.unlink(self.config_path)
            if os.path.exists(self.leases_path):
                os.unlink(self.leases_path)
            if os.path.exists(self.hosts_path):
                os.unlink(self.hosts_path)
        except:
            pass

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, _, __, ___):
        self.stop()

    def _create_config(self, start, end, dns, domain, address):
        """Create dnsmasq configuration for DHCPv6"""
        with self.hosts_file:
            # NOTE: most basic tooling expect the server to reply with
            #       and A record, not just a AAAA record, so a client
            #       like ping, in dhcpv6_basic, will fail even if -6
            self.hosts_file.write("# Generated by Infamy DHCPv6\n")
            self.hosts_file.write(f"10.10.0.1 server.{domain}\n")  # dummy
            self.hosts_file.write(f"{address} server.{domain}\n")

        with self.config_file:
            self.config_file.write("# Generated by Infamy DHCPv6\n")
            self.config_file.write(f"interface={self.iface}\n")
            self.config_file.write(f"addn-hosts={self.hosts_path}\n")
            self.config_file.write(f"domain={domain}\n")
            self.config_file.write("enable-ra\n")
            self.config_file.write(f"dhcp-leasefile={self.leases_path}\n")
            self.config_file.write(f"dhcp-range={self.iface},")

            if start and end:
                # Stateful DHCPv6 - assign addresses
                self.config_file.write(f"{start},{end},64,3600\n")
            else:
                # Stateless DHCPv6 - only provide options
                # TODO: needs more testing with ra-names + slaac
                self.config_file.write(f"::10,::fff,constructor:{self.iface},")
                self.config_file.write("64,3600\n")

            if dns:
                if isinstance(dns, list):
                    dns_str = ','.join(dns)
                else:
                    dns_str = dns
                self.config_file.write(f"dhcp-option=option6:dns-server,{dns_str}\n")

            if domain:
                self.config_file.write(f"dhcp-option=option6:domain-search,{domain}\n")

            self.config_file.write("log-debug\n")
            self.config_file.write("log-dhcp\n")

    def start(self):
        """Start the DHCPv6 server"""
        if not os.path.exists(self.config_path):
            raise Exception("Config file does not exist")

        # Drop DEVNULL redirec to debug
        cmd = f"dnsmasq --conf-file={self.config_path} --no-daemon"
        self.process = self.netns.popen(cmd.split(" "),
                                        stdout=subprocess.DEVNULL,
                                        stderr=subprocess.DEVNULL)

    def stop(self):
        """Stop the DHCPv6 server"""
        if self.process:
            self.process.terminate()
            self.process.wait()
            self.process = None

class Server6Dhcpd:
    """DHCPv6 server using ISC dhcpd with prefix delegation support"""

    def __init__(self, netns, start=None, end=None, prefix=None, prefix_len=64,
                 dns=None, domain=None, iface="iface", address=None, subnet="2001:db8::/48"):
        """Initialize DHCPv6 server with ISC dhcpd

        Args:
            netns:      Network namespace to run server in
            start:      Starting IPv6 address for IA_NA range (e.g., "2001:db8::100")
            end:        Ending IPv6 address for IA_NA range (e.g., "2001:db8::200")
            prefix:     IPv6 prefix pool for prefix delegation (e.g., "2001:db8:100::")
            prefix_len: Length of delegated prefixes (e.g., 64 for /64 prefixes)
            dns:        DNS server addresses (list or string)
            domain:     DNS search domain
            iface:      Interface to listen on
            address:    Server address on the interface (for subnet config)
            subnet:     Subnet declaration (e.g., "2001:db8::/48")
        """
        self.process = None
        self.netns = netns
        self.iface = iface
        self.subnet = subnet

        self.config_file = tf.NamedTemporaryFile(mode='w', prefix='dhcpd6_',
                                                 suffix='.conf', delete=False)
        self.config_path = self.config_file.name

        self.leases_file = tf.NamedTemporaryFile(mode='w', prefix='dhcpd6_',
                                                 suffix='.leases', delete=False)
        self.leases_path = self.leases_file.name
        self.leases_file.close()

        self.pid_file = tf.NamedTemporaryFile(mode='w', prefix='dhcpd6_',
                                              suffix='.pid', delete=False)
        self.pid_path = self.pid_file.name
        self.pid_file.close()

        self._create_config(start, end, prefix, prefix_len, dns, domain)

    def __del__(self):
        try:
            if os.path.exists(self.config_path):
                os.unlink(self.config_path)
            if os.path.exists(self.leases_path):
                os.unlink(self.leases_path)
            if os.path.exists(self.pid_path):
                os.unlink(self.pid_path)
        except:
            pass

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, _, __, ___):
        self.stop()

    def _create_config(self, start, end, prefix, prefix_len, dns, domain):
        """Create ISC dhcpd configuration for DHCPv6 with prefix delegation"""
        with self.config_file:
            self.config_file.write("# Generated by Infamy DHCPv6\n")

            # Global options
            self.config_file.write("default-lease-time 3600;\n")
            self.config_file.write("max-lease-time 7200;\n")
            self.config_file.write("log-facility local7;\n")
            self.config_file.write("\n")

            # DNS options
            if dns:
                if isinstance(dns, list):
                    dns_str = ', '.join(dns)
                else:
                    dns_str = dns
                self.config_file.write(f"option dhcp6.name-servers {dns_str};\n")

            if domain:
                self.config_file.write(f"option dhcp6.domain-search \"{domain}\";\n")

            self.config_file.write("\n")

            # Subnet declaration
            self.config_file.write(f"subnet6 {self.subnet} {{\n")

            # Address range for IA_NA (regular address assignment)
            if start and end:
                self.config_file.write(f"    range6 {start} {end};\n")

            # Prefix delegation pool
            if prefix:
                # Calculate end of prefix pool based on prefix_len
                # For simplicity, we'll create a reasonable range
                # e.g., if prefix is 2001:db8:100:: and prefix_len is 64,
                # delegate /64 prefixes from 2001:db8:100:: to 2001:db8:1ff::
                prefix_base = prefix.rstrip(':')
                # Simple approach: if prefix ends with ::, add range
                if prefix.endswith('::'):
                    prefix_start = prefix
                    # Create end address by incrementing last hex digit before ::
                    # For 2001:db8:100::, end would be 2001:db8:1ff::
                    parts = prefix_base.split(':')
                    if len(parts) >= 2:
                        last_part = parts[-2] if parts[-2] else parts[-3]
                        try:
                            last_val = int(last_part, 16)
                            end_val = last_val + 0xff
                            parts_copy = parts[:-1]
                            parts_copy[-1] = f"{end_val:x}"
                            prefix_end = ':'.join(parts_copy) + '::'
                        except:
                            prefix_end = prefix  # Fallback
                    else:
                        prefix_end = prefix
                else:
                    prefix_start = prefix
                    prefix_end = prefix

                self.config_file.write(f"    prefix6 {prefix_start} {prefix_end} /{prefix_len};\n")

            self.config_file.write("}\n")

    def start(self):
        """Start the DHCPv6 server"""
        if not os.path.exists(self.config_path):
            raise Exception("Config file does not exist")

        # Debug: show config and interface status
        # self.netns.popen(f"cat {self.config_path}".split(" "))
        # self.netns.popen("ifconfig")

        # ISC dhcpd command for IPv6
        # -6: IPv6 mode
        # -f: Stay in foreground (no daemon)
        # -d: Debug/log to stderr
        # -cf: Config file
        # -lf: Lease file
        # -pf: PID file
        cmd = [
            "dhcpd",
            "-6",           # IPv6 mode
            "-f",           # Foreground
            "-d",           # Debug output
            "-cf", self.config_path,
            "-lf", self.leases_path,
            "-pf", self.pid_path,
            self.iface      # Interface to listen on
        ]

        self.process = self.netns.popen(cmd,
                                        stdout=subprocess.DEVNULL,
                                        stderr=subprocess.DEVNULL)

    def stop(self):
        """Stop the DHCPv6 server"""
        if self.process:
            self.process.terminate()
            try:
                self.process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.process.kill()
                self.process.wait()
            self.process = None