diff options
Diffstat (limited to 'main.py')
| -rwxr-xr-x | main.py | 89 |
1 files changed, 48 insertions, 41 deletions
@@ -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: |