Home | History | Annotate | Download | only in python 
      1 import os
      2 import subprocess
      3 import pyroute2
      4 from pyroute2 import IPRoute, NetNS, IPDB, NSPopen
      5 
      6 class Simulation(object):
      7     """
      8     Helper class for controlling multiple namespaces. Inherit from
      9     this class and setup your namespaces.
     10     """
     11 
     12     def __init__(self, ipdb):
     13         self.ipdb = ipdb
     14         self.ipdbs = {}
     15         self.namespaces = []
     16         self.processes = []
     17         self.released = False
     18 
     19     # helper function to add additional ifc to namespace
     20     # if called directly outside Simulation class, "ifc_base_name" should be
     21     # different from "name", the "ifc_base_name" and "name" are the same for
     22     # the first ifc created by namespace
     23     def _ns_add_ifc(self, name, ns_ifc, ifc_base_name=None, in_ifc=None,
     24                     out_ifc=None, ipaddr=None, macaddr=None, fn=None, cmd=None,
     25                     action="ok", disable_ipv6=False):
     26         if name in self.ipdbs:
     27             ns_ipdb = self.ipdbs[name]
     28         else:
     29             try:
     30                 nl=NetNS(name)
     31                 self.namespaces.append(nl)
     32             except KeyboardInterrupt:
     33                 # remove the namespace if it has been created
     34                 pyroute2.netns.remove(name)
     35                 raise
     36             ns_ipdb = IPDB(nl)
     37             self.ipdbs[nl.netns] = ns_ipdb
     38             if disable_ipv6:
     39                 cmd1 = ["sysctl", "-q", "-w",
     40                        "net.ipv6.conf.default.disable_ipv6=1"]
     41                 nsp = NSPopen(ns_ipdb.nl.netns, cmd1)
     42                 nsp.wait(); nsp.release()
     43             ns_ipdb.interfaces.lo.up().commit()
     44         if in_ifc:
     45             in_ifname = in_ifc.ifname
     46             with in_ifc as v:
     47                 # move half of veth into namespace
     48                 v.net_ns_fd = ns_ipdb.nl.netns
     49         else:
     50             # delete the potentially leaf-over veth interfaces
     51             ipr = IPRoute()
     52             for i in ipr.link_lookup(ifname='%sa' % ifc_base_name): ipr.link_remove(i)
     53             ipr.close()
     54             try:
     55                 out_ifc = self.ipdb.create(ifname="%sa" % ifc_base_name, kind="veth",
     56                                            peer="%sb" % ifc_base_name).commit()
     57                 in_ifc = self.ipdb.interfaces[out_ifc.peer]
     58                 in_ifname = in_ifc.ifname
     59                 with in_ifc as v:
     60                     v.net_ns_fd = ns_ipdb.nl.netns
     61             except KeyboardInterrupt:
     62                 # explicitly remove the interface
     63                 out_ifname = "%sa" % ifc_base_name
     64                 if out_ifname in self.ipdb.interfaces: self.ipdb.interfaces[out_ifname].remove().commit()
     65                 raise
     66 
     67         if out_ifc: out_ifc.up().commit()
     68         ns_ipdb.interfaces.lo.up().commit()
     69         ns_ipdb.initdb()
     70         in_ifc = ns_ipdb.interfaces[in_ifname]
     71         with in_ifc as v:
     72             v.ifname = ns_ifc
     73             if ipaddr: v.add_ip("%s" % ipaddr)
     74             if macaddr: v.address = macaddr
     75             v.up()
     76         if disable_ipv6:
     77             cmd1 = ["sysctl", "-q", "-w",
     78                    "net.ipv6.conf.%s.disable_ipv6=1" % out_ifc.ifname]
     79             subprocess.call(cmd1)
     80         if fn and out_ifc:
     81             self.ipdb.nl.tc("add", "ingress", out_ifc["index"], "ffff:")
     82             self.ipdb.nl.tc("add-filter", "bpf", out_ifc["index"], ":1",
     83                             fd=fn.fd, name=fn.name, parent="ffff:",
     84                             action=action, classid=1)
     85         if cmd:
     86             self.processes.append(NSPopen(ns_ipdb.nl.netns, cmd))
     87         return (ns_ipdb, out_ifc, in_ifc)
     88 
     89     # helper function to create a namespace and a veth connecting it
     90     def _create_ns(self, name, in_ifc=None, out_ifc=None, ipaddr=None,
     91                    macaddr=None, fn=None, cmd=None, action="ok", disable_ipv6=False):
     92         (ns_ipdb, out_ifc, in_ifc) = self._ns_add_ifc(name, "eth0", name, in_ifc, out_ifc,
     93                                                       ipaddr, macaddr, fn, cmd, action,
     94                                                       disable_ipv6)
     95         return (ns_ipdb, out_ifc, in_ifc)
     96 
     97     def release(self):
     98         if self.released: return
     99         self.released = True
    100         for p in self.processes:
    101             if p.released: continue
    102             try:
    103                 p.kill()
    104                 p.wait()
    105             except:
    106                 pass
    107             finally:
    108                 p.release()
    109         for name, db in self.ipdbs.items(): db.release()
    110         for ns in self.namespaces: ns.remove()
    111 
    112