mvp

1 files changed, 48 insertions, 41 deletions

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: