mvp

4 files changed, 114 insertions, 43 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..fdc002f
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,3 @@
+.current_state.yml
+__pycache__
+*.swp
diff --git a/config.yml b/config.yml
index 3ff5b80..4b1cea3 100644
--- a/config.yml
+++ b/config.yml
@@ -1,10 +1,10 @@
 links:
   - &100
     bandwidth: 100M
-    latency: 2ms
+    latency: 500us
   - &1000
     bandwidth: 1000M
-    latency: 1ms
+    latency: 100us
   - &vdsl
     bandwidth:
       up: 3M
diff --git a/main.py b/main.py
index e50e088..8e236d6 100755
--- a/main.py
+++ b/main.py
@@ -5,7 +5,7 @@ import shutil
 import subprocess
 import sys
 import yaml
-from pyroute2 import NDB, netns
+import net
 
 class SubnetManager:
     def __init__(self, config):
@@ -32,24 +32,24 @@ class SubnetManager:
 class Latency:
     def __init__(self, latency, offset = None):
         if type(latency) is int:
-            self.latency_ms = latency
+            self.latency_us = latency
         else:
-            for suffix, factor in [("ms", 1), ("s", 1000), ("m", 60 * 1000)]:
+            for suffix, factor in [("us", 1),("ms", 1000), ("s", 1000000), ("m", 60 * 1000000)]:
                 if latency.endswith(suffix):
-                    self.latency_ms = int(float(latency[:-len(suffix)]) * factor)
+                    self.latency_us = int(float(latency[:-len(suffix)]) * factor)
                     break
             else:
-                self.latency_ms = int(latency)
+                self.latency_us = int(latency) * 1000
         if offset:
-            self.latency_ms -= Latency(offset).latency_ms
-        if self.latency_ms < 0:
-            self.latency_ms = 0
+            self.latency_us -= Latency(offset).latency_us
+        if self.latency_us < 0:
+            self.latency_us = 0
 
     def __eq__(self, o):
-        return isinstance(o, Latency) and o.latency_ms == self.latency_ms
+        return isinstance(o, Latency) and o.latency_us == self.latency_us
 
     def __str__(self):
-        return f'{self.latency_ms}ms'
+        return f'{self.latency_us}ms'
 
 class Bandwidth:
     def __init__(self, bw):
@@ -190,64 +190,53 @@ class Network:
 
 class NamespaceManager:
     def __init__(self):
-        self.ndb = NDB(log='debug')
         self.namespaces = set()
+        self.prefixlen = 0
 
     def make_namespace(self, name):
         if not name in self.namespaces:
+            net.ns.create(name)
             self.namespaces.add(name)
-            self.ndb.sources.add(netns=name)
 
-    def make_bridge(self, name, namespace):
+    def make_bridge(self, name, namespace, ports):
         self.make_namespace(namespace)
-        return (self.ndb.interfaces.create(ifname=name, kind='bridge', target=namespace)
-            .set('state', 'up'))
+        net.create_bridge(name, namespace, ports)
         
     def make_veth(self, name1, name2, space1, space2, ip, link):
         self.make_namespace(space1)
         self.make_namespace(space2)
-        veth = (self.ndb.interfaces.create(ifname=name1, target=space1, kind='veth',
-                peer={'ifname': name2, 'net_ns_fd': space2})
-                .set('state', 'up'))
-        if ip:
-            veth.add_ip(address=str(ip), prefixlen=self.prefixlen)
-        veth.commit()
-        #NSPopen(space1, ['tc', 'qdisc', 'add', 'dev', nam1, 'root', 'netem', 'delay', '1000ms', '0ms'])
-        #add the other way too
+        net.create_veth(name1, space1, name2, space2, ip, self.prefixlen, link)
 
     def build_network(self, network):
         self.prefixlen = network.subnet_manager.prefix
         netns = "testnet-core"
-        bridge = self.make_bridge("br0", netns)
+        ports = []
         for zone in network.zones.values():
             if zone.is_zone():
                 self.build_zone(zone)
             else:
                 self.build_server(zone)
-            bridge.add_port({'ifname': 'veth-' + zone.name, 'target': netns})
-        bridge.commit()
+            ports.append('veth-' + zone.name)
+        self.make_bridge("br0", netns, ports)
 
     def build_zone(self, zone):
         netns = "testnet-" + zone.name
-        bridge = self.make_bridge("br-" + zone.name, netns)
         self.make_veth("veth-" + zone.name, "veth-" + zone.name, netns, "testnet-core", None, zone.link)
-        bridge.add_port({'ifname': 'veth-' + zone.name, 'target': netns})
+        ports = ['veth-' + zone.name]
         for server in zone.servers.values():
             self.build_server(server, zone)
-            bridge.add_port({'ifname': 'veth-' + server.name, 'target': netns})
-        bridge.commit()
+            ports.append('veth-' + server.name)
+        self.make_bridge("br-" + zone.name, netns, ports)
 
     def build_server(self, server, zone = None):
         if zone:
             zone_name = "testnet-" + zone.name
+            namespace = zone_name + "-" + server.name
         else:
             zone_name = "testnet-core"
-        namespace = zone_name + "-" + server.name
+            namespace = "testnet-" + server.name + "-" + server.name
         self.make_veth("veth", "veth-" + server.name, namespace, zone_name, server.ip, server.link)
 
-    def close(self):
-        self.ndb.close()
-
 def parse(yaml):
     server_list = yaml["servers"]
     global_conf = yaml.get("global", {})
@@ -269,7 +258,6 @@ def create(config_path):
     network = parse(config)
     nsm = NamespaceManager()
     nsm.build_network(network)
-    nsm.close()
 
 def run(netns, cmd):
     if ":" in netns:
@@ -309,12 +297,31 @@ def run(netns, cmd):
         cmd = [os.getenv("SHELL") or "/bin/sh"]
     os.execve("/bin/env", ["/bin/env", "ip", "netns" , "exec", name ] + cmd, env)
 
+def runall(cmd):
+    with open(".current_state.yml", "r") as file:
+        config = yaml.safe_load(file)
+    zones = parse(config).zones
+
+    for zone in zones.values():
+        if zone.is_zone():
+            for server in zone.servers.values():
+                env = os.environ.copy()
+                env["ZONE"] = zone.name
+                env["HOST"] = server.name
+                env["IP"] = str(server.ip)
+                name = f'testnet-{zone.name}-{server.name}'
+                net.ns.run(name, cmd, env)
+        else:
+            env = os.environ.copy()
+            env["HOST"] = zone.name
+            env["IP"] = str(zone.ip)
+            name = f'testnet-{zone.name}-{zone.name}'
+            net.ns.run(name, cmd, env)
+
 
 def destroy():
-    for ns in netns.listnetns():
-        if ns.startswith("testnet-"):
-            subprocess.run(f"ip netns pids {ns} | xargs -r kill", check=True, shell=True)
-            netns.remove(ns)
+    for ns in net.ns.list():
+        net.ns.kill(ns)
     os.remove(".current_state.yml")
 
 if __name__ == "__main__":
@@ -328,10 +335,10 @@ if __name__ == "__main__":
     cmd = sys.argv[1]
     if cmd == "create":
         create(sys.argv[2] if len(sys.argv) > 2 else "config.yml")
-    elif cmd == "run-all":
-        runall(sys.argv[2:])
     elif cmd == "run":
         run(sys.argv[2], sys.argv[3:])
+    elif cmd == "run-all":
+        runall(sys.argv[2:])
     elif cmd == "destroy":
         destroy()    
     else:
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)