mknet/net.py

import subprocess
import os

_netns = ["ip", "netns"]
def run_netns(*cmd):
    process = subprocess.run(_netns + list(cmd), stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if process.returncode != 0:
        raise Exception(f"Failed to run command {cmd}:" + process.stderr)
    return process

class ns:
    def name_unconfined():
        if "unconfined" not in ns.list(True):
            run_netns("attach", "unconfined", str(os.getpid()))

    def list(include_unconfined = False):
        try:
            nss = os.listdir("/var/run/netns")
            return [ns for ns in nss if ns.startswith("testnet-") or include_unconfined and ns == "unconfined"]
        except FileNotFoundError:
            return []

    def forget(name):
        run_netns("del", name)

    def kill(name):
        pids = run_netns("pids", name).stdout.split("\n")
        pids = [pid for pid in pids if pid]
        if pids:
            process = subprocess.run(["sudo", "kill", "-9"] + pids)
            if process.returncode != 0:
                raise Exception("Failed to list namespaces: " + process.stderr)
        ns.forget(name)

    def create(name):
        run_netns("add", name)
        run_netns("exec", name, "ip", "link", "set", "dev", "lo", "up")

    def run(name, cmd, env=None):
        return subprocess.Popen(_netns + ["exec", name] + cmd, env=env)

def create_bridge(name, namespace, ports=[]):
    run_netns("exec", namespace, "ip", "link", "add", "name", name, "type", "bridge")
    run_netns("exec", namespace, "ip", "link", "set", "dev", name, "up")
    for port in ports:
        run_netns("exec", namespace, "ip", "link", "set", "dev", port, "master", name)
    pass

def create_veth(name1, ns1, name2, ns2, ip = None, subnet=0, link=None):
    run_netns("exec", ns1, "ip", "link", "add", "name", name1, "type", "veth",
            "peer", "name", name2, "netns", ns2)
    if ip:
        ip = f"{ip}/{subnet}"
        run_netns("exec", ns1, "ip", "addr", "add", "dev", name1, ip)
    run_netns("exec", ns1, "ip", "link", "set", "dev", name1, "up")
    run_netns("exec", ns2, "ip", "link", "set", "dev", name2, "up")
    
    if link:
        tc(ns1, name1, link)
        tc(ns2, name2, link, True)

def tc(namespace, name, link, invert=False):
    options = []
    if invert:
        options += ["delay", str(link.latency.latency_us), str(link.jitter.latency_us)]
        options += ["rate", str(link.bandwidth.down)]
    else:
        options += ["rate", str(link.bandwidth.up)]
    if link.limit:
        options += ["limit", str(link.limit)]
    run_netns("exec", namespace, "tc", "qdisc", "add", "dev", name, "root", "netem", *options)