start_testbed.py

#! /usr/bin/env python

import os
import subprocess
import yaml
from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument(
    '-f', '--file',
    dest='file',
    default='example.yml',
    help='YML file containing the testbed configuration'
   )

args = parser.parse_args()

with open(args.file) as cfile:
    conf = yaml.load(cfile)

ntwk = conf['network']
base_addr = reduce(
    lambda a,b : a*256+b,
    [ int(q) for q in ntwk['subnet_addr'].split('.') ]
   )
subnet = ntwk['subnet_addr']+'/'+str(ntwk['subnet_len'])

nodelist = list()

retval = subprocess.call(['docker','pull',conf['image']])

def make_ip(i):
    addr = base_addr + i
    masks = [ 24, 16, 8, 0 ]
    return '.'.join([ str((addr>>mask)&0xFF) for mask in masks ])
    
def make_route(f,net,fdev):
    subprocess.call([
        'sudo', 'ip', 'netns', 'exec', f['name'],
        'ip', 'route', 'add', net,
        'dev', fdev, 'proto', 'static', 'scope', 'global',
        'src', f['addr']
       ])

def connect_nodes(node1, node2):
    ctr1 = nodelist[node1]
    ctr2 = nodelist[node2]
    print ctr1, ctr2
    # Create a virtual ethernet link.
    dev_1 = '{name}_{n1}_{n2}'.format(name=ntwk['name'],n1=node1,n2=node2)
    dev_2 = '{name}_{n2}_{n1}'.format(name=ntwk['name'],n1=node1,n2=node2)
    subprocess.call([
        'sudo',
        'ip', 'link', 'add', dev_1, 'type', 'veth',
        'peer', 'name', dev_2
        ])
    # Add the link endpoints to the container namespaces.
    def add_to_ns(dev,ns):
        subprocess.call([
            'sudo', 'ip', 'link', 'set', dev, 'netns', ns
           ])
    add_to_ns(dev_1,ctr1['name'])
    add_to_ns(dev_2,ctr2['name'])

    # Set the IP addresses.
    def set_addr(dev,ns,addr):
        subprocess.call([
            'sudo', 'ip', 'netns', 'exec', ns,
            'ip', 'addr', 'add', addr+'/32', 'dev', dev
           ])
    set_addr(dev_1,ctr1['name'],ctr1['addr'])
    set_addr(dev_2,ctr2['name'],ctr2['addr'])

    # Bring the link up.
    def link_up(ns,dev):
        subprocess.call([
            'sudo', 'ip', 'netns', 'exec', ns,
            'ip', 'link', 'set', dev, 'up'
           ])
    link_up(ctr1['name'],dev_1)
    link_up(ctr2['name'],dev_2)

    # Set initial routes.
    def route_from_to(f,t,fdev):
        make_route(f,t['addr']+'/32',fdev)
    route_from_to(ctr1,ctr2,dev_1)
    route_from_to(ctr2,ctr1,dev_2)


def make_mesh():
    global nodelist
    num_nodes = len(nodelist)
    [ connect_nodes(n1,n2)
        for n1 in range(num_nodes-1)
        for n2 in range(n1+1,num_nodes) ]

def make_star():
    global nodelist
    def get_dev(i):
        return 'bac0_{i}_0'.format(i=i)
    num_nodes = len(nodelist)
    [ connect_nodes(0,n) for n in range(1,num_nodes) ]
    [ make_route(n,subnet,get_dev(n)) for n in range(1,num_nodes) ]

def start_node(base_name,i):
    global nodelist
    name = base_name + str(i)

    # Start the container.
    command = ['docker',
               'run',
               '-d',
               '--name',
               name,
               conf['image']
              ]
    command.extend(conf['image_cmd'].split())
    retval = subprocess.call(command)

    entry = dict()
    entry['name'] = name
    # Determine the IP address to use for this node.
    entry['addr'] = make_ip(i)
    # Find the container's process ID.
    entry['pid'] = subprocess.check_output([
        'docker',
        'inspect',
        '-f',
        '{{.State.Pid}}',
        name
        ]).strip()
 
    # Set up a network namespace.
    subprocess.call([
        'sudo',
        'ln',
        '-s',
        '/proc/{pid}/ns/net'.format(pid=entry['pid']),
        '/var/run/netns/{name}'.format(name=name)
       ])

    # Add the node's info to our list.
    nodelist.append(entry)

for n in range(ntwk['num_nodes']):
    start_node(ntwk['name'],n)

if ntwk['topology'] == 'mesh':
    make_mesh()
else:
    make_star()