mvp

1 files changed, 61 insertions, 0 deletions

diff --git a/net.py b/net.py
new file mode 100644
index 0000000..bc20e4d
--- /dev/null
+++ b/net.py
@@ -0,0 +1,61 @@
+import subprocess
+
+_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 list():
+        process = run_netns("list")
+        return [ns for ns in process.stdout.split("\n") if ns.startswith("testnet-")]
+
+    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):
+        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)]
+    run_netns("exec", namespace, "tc", "qdisc", "add", "dev", name, "root", "netem", *options)