#
# Rumba - Namespace testbed (network + mount namespace per node)
#
#    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 os
import subprocess
import tempfile
import time

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

from rumba.executors.nsenter import NsenterExecutor
from rumba.executors.local import LocalExecutor

logger = log.get_logger(__name__)


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


def _try(cmd):
    """Best-effort host-side teardown command; never raises."""
    logger.debug('executing >> %s', cmd)
    try:
        subprocess.check_call(cmd.split())
    except subprocess.CalledProcessError:
        logger.debug('command failed (ignored): %s', cmd)


class Testbed(mod.Testbed):
    """
    Local testbed that gives each node its own Linux network + mount
    namespace, wired together with veth pairs and bridges.

    Unlike :class:`LocalNetworkTestbed` (single IRMd in the default
    namespace), each node here is an isolated namespace, so an independent
    IRMd runs per node.  This is what lets a single host place clients and
    servers on distinct nodes and test real multi-hop forwarding.

    The one-IRMd-per-node guarantee comes entirely from the per-node mount
    namespace: Ouroboros uses compile-time-fixed paths (the lockfile and
    shm objects under ``/dev/shm``, the sockets under ``/var/run/ouroboros``
    -- or ``/tmp`` in no-/var/run builds -- and the FUSE RIB under
    ``/tmp/ouroboros``), so each node gets a fresh private tmpfs mounted
    over each of those trees.  The network namespace alone does NOT separate
    them; they are tmpfs/socket/FUSE-backed, not network state.

    Requires CAP_SYS_ADMIN (root / sudo), e.g. a privileged container.
    Uses host-installed Ouroboros binaries (``needs_install = False``).
    """

    # Uses host-installed binaries; one IRMd per node (not single_irmd);
    # swap_out destroys the namespaces, so the per-node IRMds must be reaped
    # (terminate_prototype) before that -- see rumba.utils.ExperimentManager.
    needs_install = False
    terminate_before_swapout = True

    # tc link quality is supported: HTB rate limiting runs on the IPCP
    # interface inside the node namespace (via the nsenter executor), while
    # netem impairments run on the bridge-side (host_ifname) veth in the
    # default namespace (via each node's host_executor -- see _start_node).
    # This preserves the send back-pressure fix from the veth testbeds.

    # The three filesystem trees Ouroboros pins at compile time.  Each gets a
    # private tmpfs per node so exactly one IRMd owns each per node.
    _STATE_TREES = ('/dev/shm', '/var/run/ouroboros', '/tmp')

    def __init__(self, exp_name='nstb', username='',
                 proj_name='rumba', password=''):
        """
        :param exp_name: The experiment name (also seeds veth name prefix).
        :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 = NsenterExecutor(self)
        # Host-side executor for link shaping: netem lives on the bridge-side
        # veth in the default namespace, so it runs on the host, not via
        # nsenter.  Assigned to each node in _start_node.
        self.host_executor = LocalExecutor(self)

        self.holders = {}        # node.name -> ns-owning pid
        self._popens = {}        # node.name -> Popen (for reaping)
        self.netns_links = []    # /var/run/netns/<pid> symlinks we created
        self.active_bridges = set()
        self.active_veths = []   # host-side veth names to delete
        self._hosts_backup = None

        # Per-experiment prefix so concurrent experiments don't collide on
        # host-visible veth names (IFNAMSIZ = 15 chars).
        self._prefix = hashlib.md5(exp_name.encode()).hexdigest()[:4]
        self._veth_idx = 0

    # ---- swap in ---------------------------------------------------------

    def _swap_in(self, experiment):
        # Prime sudo once so a later prompt can't abort us mid-sequence.
        subprocess.check_call('sudo -v'.split())

        if not os.path.exists('/var/run/netns'):
            _run('sudo mkdir -p /var/run/netns')

        # 1. Namespace holders + private tmpfs (before any wiring or IRMd).
        for node in experiment.nodes:
            self._start_node(node)

        # 2. A bridge per Ethernet layer.
        for layer in experiment.layer_ordering:
            if not layer.is_eth:
                continue
            bridge = 'br-%s' % layer.name
            if bridge in self.active_bridges:
                continue
            _run('sudo ip link add %s type bridge' % bridge)
            _run('sudo ip link set dev %s up' % bridge)
            self.active_bridges.add(bridge)

        # 3. A veth pair per Ethernet IPCP: host end on the bridge, peer end
        #    moved into the node's namespace as the IPCP interface.
        for node in experiment.nodes:
            pid = self.holders[node.name]
            eth_idx = 0
            for ipcp in node.ipcps:
                if not isinstance(ipcp, mod.EthIPCP):
                    continue
                bridge = 'br-%s' % ipcp.layer.name
                ifname = 'e%d' % eth_idx
                eth_idx += 1
                host_end, peer = self._veth_names()

                ipcp.ifname = ifname
                ipcp.host_ifname = host_end

                _run('sudo ip link add %s type veth peer name %s'
                     % (host_end, peer))
                _run('sudo ip link set %s master %s' % (host_end, bridge))
                _run('sudo ip link set dev %s up' % host_end)
                _run('sudo ip link set %s netns %d name %s'
                     % (peer, pid, ifname))
                node.execute_command('ip link set dev %s up' % ifname,
                                     as_root=True)
                self.active_veths.append(host_end)

        # 4. UDP layers: a point-to-point veth pair with IP addresses, each
        #    end moved into its node's namespace.
        self._wire_udp(experiment)

        logger.info("Experiment swapped in (%d nodes, %d bridges, %d veths)",
                     len(self.holders), len(self.active_bridges),
                     len(self.active_veths))

    def _veth_names(self):
        host_end = '%sv%d' % (self._prefix, self._veth_idx)
        peer = '%sp%d' % (self._prefix, self._veth_idx)
        self._veth_idx += 1
        return host_end, peer

    def _start_node(self, node):
        """Create the node's net+mount namespace holder and its tmpfs."""
        # unshare execs the shell in place (no --fork), so the shell -- and,
        # after `exec sleep`, the sleep -- owns the new namespaces.  We read
        # its own pid from stdout ($$) rather than Popen.pid, because Popen.pid
        # is sudo's pid, not the namespace owner's.
        popen = subprocess.Popen(
            ['sudo', 'unshare', '--net', '--mount', '--propagation', 'private',
             'sh', '-c', 'echo $$; exec sleep infinity'],
            stdout=subprocess.PIPE, universal_newlines=True)
        # Track the Popen immediately, before anything that can raise, so a
        # failed startup is still reaped by _swap_out.
        self._popens[node.name] = popen

        line = popen.stdout.readline()
        if not line.strip():
            # unshare printed no pid: missing CAP_SYS_ADMIN, denied sudo, or a
            # kernel without the requested namespaces.  Bail with a clear error.
            raise RuntimeError(
                "Namespace holder for node %s failed to start (unshare "
                "produced no pid; CAP_SYS_ADMIN / sudo required)" % node.name)
        pid = int(line.strip())

        # Wait for the holder to be alive before using its namespaces.
        # Poll /proc/<pid> (world-traversable), not /proc/<pid>/ns/net:
        # stat of the ns symlink needs ptrace access, which an unprivileged
        # runner lacks for the root-owned holder.  The sudo'd ln below is
        # the only consumer that dereferences ns/net.
        for _ in range(50):
            if os.path.exists('/proc/%d' % pid):
                break
            time.sleep(0.1)
        else:
            raise RuntimeError(
                "Namespace holder for node %s did not come up" % node.name)

        self.holders[node.name] = pid

        # netem on the bridge-side veth must run on the host, not in the ns.
        node.host_executor = self.host_executor

        # Expose the netns for `ip link set ... netns <pid>`; overwrite any
        # stale symlink left by a crashed run with the same pid.
        link = '/var/run/netns/%d' % pid
        _run('sudo ln -sf /proc/%d/ns/net %s' % (pid, link))
        self.netns_links.append(link)

        # Private tmpfs over each Ouroboros state tree, inside this node's
        # (private-propagation) mount namespace, before any IRMd starts.
        node.execute_command('mkdir -p %s' % ' '.join(self._STATE_TREES),
                             as_root=True)
        for tree in self._STATE_TREES:
            node.execute_command('mount -t tmpfs tmpfs %s' % tree,
                                 as_root=True)

        # Bring loopback up (UDP shim / local traffic).
        node.execute_command('ip link set dev lo up', as_root=True)

    def _wire_udp(self, experiment):
        udp_layers_done = set()
        udp_layer_idx = 0
        hosts_entries = []  # (ip, name) pairs for /etc/hosts

        for layer in experiment.layer_ordering:
            if not layer.is_udp or layer.name in udp_layers_done:
                continue
            udp_layers_done.add(layer.name)
            udp_layer_idx += 1

            members = [(node, ipcp)
                       for node in experiment.nodes
                       for ipcp in node.ipcps
                       if isinstance(ipcp, mod.UdpIPCP) and ipcp.layer == layer]
            if len(members) != 2:
                raise RuntimeError(
                    "UDP layer %s needs exactly 2 members, got %d"
                    % (layer.name, len(members)))

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

            for end, (node, ipcp), host in ((end_a, members[0], 1),
                                            (end_b, members[1], 2)):
                pid = self.holders[node.name]
                ifname = 'u%d' % host
                ip = '10.%d.0.%d' % (udp_layer_idx, host)

                _run('sudo ip link set %s netns %d name %s'
                     % (end, pid, ifname))
                node.execute_command('ip addr add %s/24 dev %s' % (ip, ifname),
                                     as_root=True)
                node.execute_command('ip link set dev %s up' % ifname,
                                     as_root=True)

                ipcp.ifname = ifname
                ipcp.ip_addr = ip

                hosts_entries.extend(self._udp_hosts_entries(node, ipcp,
                                                             layer, ip))

        self._write_hosts(hosts_entries)

    @staticmethod
    def _udp_hosts_entries(node, ipcp, udp_layer, ip):
        """/etc/hosts (ip, name) pairs for UDP name resolution (see localnet):
        the UDP shim resolves names by their md5 hex digest."""
        entries = [(ip, ipcp.name)]
        for other in node.ipcps:
            if other is ipcp:
                continue
            if other.layer in node.registrations \
               and udp_layer in node.registrations[other.layer]:
                entries.append((ip, other.name))
                entries.append((ip, other.layer.name))
        return entries

    def _write_hosts(self, hosts_entries):
        """Append md5-hashed name->IP entries to the shared /etc/hosts.

        The mount namespaces do not remount /etc, so /etc/hosts is the single
        host file, visible identically in every node -- one global edit is
        enough (and matches LocalNetworkTestbed)."""
        if not hosts_entries:
            return

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

        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 nstb UDP ---\n' + '\n'.join(lines) + '\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)

    @staticmethod
    def _kill_ns_processes(holder_pid):
        """Best-effort SIGKILL of every process sharing the holder's network
        namespace (except the holder itself).  Identified by net-ns inode, so
        it is prototype-agnostic and does not depend on process names."""
        try:
            target = os.stat('/proc/%d/ns/net' % holder_pid).st_ino
        except OSError:
            return
        for entry in os.listdir('/proc'):
            if not entry.isdigit():
                continue
            pid = int(entry)
            if pid == holder_pid:
                continue
            try:
                if os.stat('/proc/%d/ns/net' % pid).st_ino == target:
                    _try('sudo kill -9 %d' % pid)
            except OSError:
                continue

    # ---- swap out --------------------------------------------------------

    def _swap_out(self, experiment):
        # By the time we get here the per-node IRMds have already been reaped
        # (terminate_before_swapout -> terminate_prototype, while the holders
        # were still alive).  Now tear down the namespaces and wiring.  Every
        # step tolerates missing resources so partial swap-in state is cleaned.

        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
            _try('sudo cp %s /etc/hosts' % tmp_path)
            subprocess.run(['rm', '-f', tmp_path], stdout=subprocess.DEVNULL,
                           stderr=subprocess.DEVNULL)
            self._hosts_backup = None

        # Tear down every holder we started (tracked before pid capture, so a
        # holder whose unshare failed is reaped too).  Killing the last process
        # in a namespace drops it, so the kernel tears down the private tmpfs
        # and any FUSE RIB mounts.  We do NOT use a PID namespace, so killing
        # the holder does not by itself cascade to in-ns IRMd/IPCPd: reap those
        # first by net-namespace membership, so swap-out is correct even if
        # terminate_prototype did not run (or partially failed).
        for name, popen in list(self._popens.items()):
            pid = self.holders.get(name)
            if pid is not None:
                self._kill_ns_processes(pid)
                _try('sudo kill %d' % pid)
            else:
                try:
                    popen.kill()
                except OSError:
                    pass
            try:
                popen.wait(timeout=5)
            except subprocess.TimeoutExpired:
                logger.debug('holder %s slow to exit', name)
        self._popens = {}
        self.holders = {}

        # Host-side artifacts the namespace teardown does NOT remove.
        for link in self.netns_links:
            _try('sudo rm -f %s' % link)
        self.netns_links = []

        for host_end in self.active_veths:
            _try('sudo ip link del %s' % host_end)
        self.active_veths = []

        for bridge in list(self.active_bridges):
            _try('sudo ip link del %s' % bridge)
            self.active_bridges.discard(bridge)

        logger.info("Experiment swapped out")
