#
# Rumba - Local Network Testbed (veth + bridges)
#
#    Copyright (C) 2017-2026 imec
#
#    Dimitri Staessens <dimitri.staessens@ugent.be>
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., http://www.fsf.org/about/contact/.
#

import hashlib
import subprocess
import tempfile

import rumba.model as mod
import rumba.log as log

from rumba.executors.local import LocalExecutor

logger = log.get_logger(__name__)


def _run(cmd):
    """Run a command with sudo, logging it first."""
    logger.debug('executing >> %s', cmd)
    subprocess.check_call(cmd.split())


class LocalNetworkTestbed(mod.Testbed):
    """
    Local testbed using Linux bridges and veth pairs.

    Each Ethernet layer gets a Linux bridge. Each node's Ethernet IPCP
    gets a veth pair: one end is attached to the layer's bridge, the
    other end is the interface that the IPCP binds to. All processes
    (IRMd, IPCPs, applications) run in the default network namespace
    under a single IRMd instance.

    This allows testing the real eth-dix / eth-llc stack locally with
    netem support on the veth interfaces.
    """

    def __init__(self, exp_name='localnet', username='',
                 proj_name='rumba', password=''):
        """
        :param exp_name: The experiment name.
        :param username: User of the experiment.
        :param proj_name: Project name of the experiment.
        :param password: Password of the user.
        """
        mod.Testbed.__init__(self, exp_name, username, password, proj_name)

        self.executor = LocalExecutor(self)
        self.active_bridges = set()
        self.active_veths = []   # list of (host_end, ipcp_end) tuples
        self._hosts_backup = None  # original /etc/hosts content

    def _swap_in(self, experiment):
        # Ensure sudo is primed
        subprocess.check_call('sudo -v'.split())

        # Create a bridge for each Ethernet layer
        for layer in experiment.layer_ordering:
            if not layer.is_eth:
                continue

            bridge = 'br-%s' % layer.name
            _run('sudo ip link add %s type bridge' % bridge)
            _run('sudo ip link set dev %s up' % bridge)
            self.active_bridges.add(bridge)

        # Create veth pairs for each Ethernet IPCP
        veth_idx = 0
        for node in experiment.nodes:
            for ipcp in node.ipcps:
                if not isinstance(ipcp, mod.EthIPCP):
                    continue

                bridge = 'br-%s' % ipcp.layer.name

                # Generate unique interface names (max 15 chars)
                ipcp_end = 'veth%d' % veth_idx
                host_end = 'vb%d' % veth_idx
                veth_idx += 1

                # Assign the IPCP-facing interface name
                ipcp.ifname = ipcp_end

                # Create veth pair
                _run('sudo ip link add %s type veth peer name %s'
                     % (host_end, ipcp_end))

                # Attach bridge end to the layer bridge
                _run('sudo ip link set %s master %s' % (host_end, bridge))

                # Bring both ends up
                _run('sudo ip link set dev %s up' % host_end)
                _run('sudo ip link set dev %s up' % ipcp_end)

                self.active_veths.append((host_end, ipcp_end))

        # Create veth pairs with IP addresses for each UDP layer.
        # UDP layers are point-to-point (max 2 members), so we create
        # a single veth pair and assign an IP to each end.
        udp_layers_done = set()
        udp_layer_idx = 0
        for layer in experiment.layer_ordering:
            if not layer.is_udp:
                continue
            if layer.name in udp_layers_done:
                continue
            udp_layers_done.add(layer.name)
            udp_layer_idx += 1

            # Collect the (at most 2) UdpIPCPs for this layer
            udp_ipcps = [ipcp for node in experiment.nodes
                         for ipcp in node.ipcps
                         if isinstance(ipcp, mod.UdpIPCP)
                         and ipcp.layer == layer]

            if len(udp_ipcps) != 2:
                raise Exception(
                    "UDP layer %s needs exactly 2 members, got %d"
                    % (layer.name, len(udp_ipcps)))

            end_a = 'veth%d' % veth_idx
            end_b = 'veth%d' % (veth_idx + 1)
            veth_idx += 2

            _run('sudo ip link add %s type veth peer name %s'
                 % (end_a, end_b))

            # Assign IP addresses (10.{layer_idx}.0.{1,2}/24)
            ip_a = '10.%d.0.1/24' % udp_layer_idx
            ip_b = '10.%d.0.2/24' % udp_layer_idx

            _run('sudo ip addr add %s dev %s' % (ip_a, end_a))
            _run('sudo ip addr add %s dev %s' % (ip_b, end_b))
            _run('sudo ip link set dev %s up' % end_a)
            _run('sudo ip link set dev %s up' % end_b)

            udp_ipcps[0].ifname = end_a
            udp_ipcps[0].ip_addr = '10.%d.0.1' % udp_layer_idx
            udp_ipcps[1].ifname = end_b
            udp_ipcps[1].ip_addr = '10.%d.0.2' % udp_layer_idx

            self.active_veths.append((end_a, end_b))

        # ---- /etc/hosts entries for UDP name resolution ----
        # The UDP shim IPCP resolves names via DNS (falling back to
        # /etc/hosts when dns_addr is 0.0.0.0).  We add entries so
        # that IPCP names and overlay names registered in UDP layers
        # resolve to the correct veth IP address.
        hosts_entries = []  # list of (ip, md5_hash) tuples
        for node in experiment.nodes:
            # Find this node's UDP IPCP(s) and their IPs
            for ipcp in node.ipcps:
                if not isinstance(ipcp, mod.UdpIPCP):
                    continue
                if ipcp.ip_addr is None or ipcp.ip_addr == '0.0.0.0':
                    continue
                udp_ip = ipcp.ip_addr
                udp_layer = ipcp.layer

                # 1) The UDP IPCP's own name (MD5-hashed)
                hosts_entries.append((udp_ip, ipcp.name))

                # 2) Names of overlay IPCPs registered in this UDP layer,
                #    plus their layer names (needed for enrollment)
                for other_ipcp in node.ipcps:
                    if other_ipcp == ipcp:
                        continue
                    if other_ipcp.layer in node.registrations \
                       and udp_layer in node.registrations[
                           other_ipcp.layer]:
                        hosts_entries.append((udp_ip, other_ipcp.name))
                        hosts_entries.append((udp_ip,
                                             other_ipcp.layer.name))

        if hosts_entries:
            result = subprocess.run(
                ['cat', '/etc/hosts'], capture_output=True, text=True)
            self._hosts_backup = result.stdout

            # Deduplicate and append.
            # The UDP IPCP resolves names by their MD5 hex digest,
            # so /etc/hosts must map the hash to the IP address.
            seen = set()
            lines = []
            for ip, name in hosts_entries:
                md5 = hashlib.md5(name.encode()).hexdigest()
                key = (ip, md5)
                if key not in seen:
                    seen.add(key)
                    lines.append('%s %s  # %s' % (ip, md5, name))

            block = '\n# --- rumba localnet UDP ---\n'
            for line in lines:
                block += line + '\n'

            with tempfile.NamedTemporaryFile(
                    mode='w', suffix='.hosts', delete=False) as tmp:
                tmp.write(self._hosts_backup + block)
                tmp_path = tmp.name
            _run('sudo cp %s /etc/hosts' % tmp_path)
            subprocess.run(['rm', '-f', tmp_path],
                           stdout=subprocess.DEVNULL,
                           stderr=subprocess.DEVNULL)

            logger.debug('/etc/hosts entries added: %s', lines)

        logger.info("Experiment swapped in (%d bridges, %d veth pairs)",
                     len(self.active_bridges), len(self.active_veths))

    def _swap_out(self, experiment):
        # Restore /etc/hosts if we modified it
        if self._hosts_backup is not None:
            with tempfile.NamedTemporaryFile(
                    mode='w', suffix='.hosts', delete=False) as tmp:
                tmp.write(self._hosts_backup)
                tmp_path = tmp.name
            _run('sudo cp %s /etc/hosts' % tmp_path)
            subprocess.run(['rm', '-f', tmp_path],
                           stdout=subprocess.DEVNULL,
                           stderr=subprocess.DEVNULL)
            self._hosts_backup = None
            logger.debug('/etc/hosts restored')

        # Deleting one end of a veth pair automatically removes the other
        for host_end, _ in self.active_veths:
            try:
                _run('sudo ip link del %s' % host_end)
            except subprocess.CalledProcessError:
                logger.debug('veth %s already gone', host_end)

        self.active_veths = []

        for bridge in list(self.active_bridges):
            try:
                _run('sudo ip link del %s' % bridge)
            except subprocess.CalledProcessError:
                logger.debug('bridge %s already gone', bridge)
            self.active_bridges.discard(bridge)

        logger.info("Experiment swapped out")
