tn

#!/usr/bin/env python3

import sys
import getopt
import argparse
import pprint
import pydot
import yaml

class AppInfo:
  verbose = False
  host = None
  name = ''
  specfile = ''
  version = 'TiNet version 0.0 alpha'
  command = ''
  namespace = ''
  data = []
  def __init__(self):
    self.specfile = 'spec.yaml'
    self.verbose = False
    self.host = None
    self.name = ''
    self.command = 'version'
    if self.host != None:
      sys.stderr.write('remote host isn\'t supported.\n')
      sys.exit(1)
    if self.name != '':
      sys.stderr.write('project-name isn\'t supported.\n')
      sys.exit(1)
  def read_yaml(self, filename):
    filename = appinfo.specfile
    f = open(filename, "r+")
    appinfo.data = yaml.load(f, Loader=yaml.SafeLoader)
    print('## Check META config')
    if 'meta' in appinfo.data:
        if 'namespace' in appinfo.data['meta']:
          appinfo.namespace = appinfo.data['meta']['namespace']
          print('# namespace was configured as {}'.format(appinfo.namespace))

appinfo = AppInfo()

template_spec='''

postinit:
  - cmds:
    - cmd: echo slankdev

meta:
  namespace: ns_

nodes:
  - name: C0
    type: docker #optional (default:docker)
    image: slankdev/ubuntu:18.04
    interfaces:
      - { name: net0, type: direct, args: C1#net0 }
      - { name: net1, type: bridge, args: B0 }
      - { name: net2, type: veth  , args: peer0 }
      - { name: net3, type: phys }
  - name: C1
    type: netns #optional (default:docker)
    interfaces:
      - { name: net0, type: direct, args: C0#net0 }
      - { name: net1, type: bridge, args: B0 }

switches:
  - name: B0
    interfaces:
      - { name: net0, type: docker, args: C0 }
      - { name: net0, type: netns, args: C0 }

node_configs:
  - name: C0
    cmds:
      - cmd: ip link set dev net0 up
  - name: C1
    cmds:
      - cmd: echo slankdev slankdev
      - cmd: >-
          echo slankdev &&
          echo slnakdev

test:
  - name: p2p
    cmds:
    - cmd: docker exec C0 ping -c2 10.0.0.2
    - cmd: echo slankdev slankdev
  - name: lo
    cmds:
    - cmd: docker exec C0 ping -c2 10.255.0.1
    - cmd: echo slankdev slankdev
'''

text__mount_docker_netns='''
mount_docker_netns () {
  if [ $# -ne 2 ]; then
    echo "Usage: $0 <container> <netns>"
    exit 1
  fi
  mkdir -p /var/run/netns
  PID=`docker inspect $1 -f "{{.State.Pid}}"`
  ln -s /proc/$PID/ns/net /var/run/netns/$2
}
'''

text__kokobr='''
kokobr () {
  if [ $# -ne 3 ]; then
    echo "Usage: $0 <bridge> <container> <ifname>"
    exit 1
  fi
  mount_docker_netns $2 $2
  ip link add name $3 type veth peer name $2$3
  ip link set dev $3 netns $2
  ip link set $2$3 up
  ip netns exec $2 ip link set $3 up
  ip netns del $2
  ovs-vsctl add-port $1 $2$3
}
'''

text__kokobr_netns='''
kokobr_netns () {
  if [ $# -ne 3 ]; then
    echo "Usage: $0 <bridge> <netns> <ifname>"
    exit 1
  fi
  ip link add name $3 type veth peer name $2$3
  ip link set dev $3 netns $2
  ip link set $2$3 up
  ip netns exec $2 ip link set $3 up
  ovs-vsctl add-port $1 $2$3
}
'''

text__koko_physnet ='''
koko_physnet () {
if [ $# -ne 2 ]; then
  echo "Usage: $0 <container> <netif>"
  exit 1
fi
mount_docker_netns $1 $1
ip link set dev $2 netns $1
ip netns exec $1 ip link set $2 up
ip netns del $1
}
'''

text__detach_physif_from_docker='''
detach_physif_from_docker () {
if [ $# -ne 2 ]; then
  echo "Usage: $0 <container> <netif>"
  exit 1
fi
mount_docker_netns $1 $1
ip netns exec $1 ip link set $2 netns 1
ip netns del $1
}
'''

def generate_mount_docker_netns():
  print(text__mount_docker_netns)

def generate_kokobr():
  print(text__kokobr)

def generate_kokobr_netns():
  print(text__kokobr_netns)

def generate_koko_physnet():
  print(text__koko_physnet)

def generate_detach_physif_from_docker():
  print(text__detach_physif_from_docker)

def precmd():
    if 'precmd' in appinfo.data:
      for node in appinfo.data['precmd']:
        for cmd in node['cmds']:
          print('{}'.format(cmd['cmd']))

def command_ps(args):
  print('docker ps --format "table {{.ID}}\\t{{.Names}}"')

def postinit():
  print('\n#####################')
  print('# POST-INIT COMMANDS #')
  print('#####################')
  if 'postinit' in appinfo.data:
    for node in appinfo.data['postinit']:
      for cmd in node['cmds']:
        print('{}'.format(cmd['cmd']))


def command_up(args):
  def preinit():
    print('\n#####################')
    print('# PRE-INIT COMMANDS #')
    print('#####################')
    if 'preinit' in appinfo.data:
      for node in appinfo.data['preinit']:
        for cmd in node['cmds']:
          print('{}'.format(cmd['cmd']))

  print("set -ue")
  precmd()
  generate_mount_docker_netns()
  generate_kokobr()
  generate_kokobr_netns()
  generate_koko_physnet()
  preinit()

  print('\n#####################')
  print('# INIT COMMANDS #')
  print('#####################')

  # CREATE NODES
  print('\n# generate nodes')
  if 'nodes' in appinfo.data:
    for node in appinfo.data['nodes']:
      nodetype = node.get('type', 'docker')
      net_base = node.get('net_base', None)

      volume_base = node.get('volume', '/tmp/tinet')
      tinet_volume_option = '-v {}:/tinet'.format(volume_base)

      docker_run_extra_args = node.get('docker_run_extra_args', '')

      entrypoint = node.get('entrypoint', None)
      entrypoint_option = ''
      if entrypoint is not None:
        entrypoint_option = '--entrypoint {}'.format(entrypoint)

      if nodetype == 'docker':
        if net_base == None:
          net_base = 'none'
        print('docker run -td --hostname {} --net {} '
          '--name {}{} --rm --privileged {} {} {} {}'
          .format(node['name'], net_base, appinfo.namespace, node['name'],
          entrypoint_option, tinet_volume_option, docker_run_extra_args, node['image']))
      elif nodetype == 'netns':
        if net_base != None:
          print('You cannot use net_base with type: netns')
          sys.exit(1)
        print('ip netns add {}{}'.format(appinfo.namespace, node['name']))
        print('ip netns exec {}{} ip link set lo up'.format(appinfo.namespace, node['name']))
      else:
        print('unknown node-type {}'.format(node['type']))
        sys.exit(1)

  # CREATE SWITCHES
  if 'switches' in appinfo.data:
    for sw in appinfo.data['switches']:
      print('ovs-vsctl add-br {}{} && ip link set {} up'.format(appinfo.namespace, sw['name'], sw['name']));

  links = []
  # CREATE Node to Node LINKS
  class n2n_link:
    left_node_name = None
    left_iface_name = None
    left_isset = False
    left_node_type = None
    right_node_name = None
    right_iface_name = None
    right_isset = False
    right_node_type = None
    def __init__(self, ln_name, ln_iface, rn_name=None, rn_iface=None):
      self.left_node_name = ln_name
      self.left_iface_name = ln_iface
      self.right_node_name = rn_name
      self.right_iface_name = rn_iface
      self.left_isset = True
    def koko(self):
      # EXAMPLE:
      #    sudo ip link add net0 netns C0 type veth \
      #        peer name net0 netns C1
      print('\n# connect {} to {}'.format(self.left_node_name, self.right_node_name))
      if self.left_node_type == 'docker':
        print('mount_docker_netns {}{} {}{}'.format(
          appinfo.namespace,
          self.left_node_name,
          appinfo.namespace,
          self.left_node_name))
      if self.right_node_type == 'docker':
        print('mount_docker_netns {}{} {}{}'.format(
          appinfo.namespace,
          self.right_node_name,
          appinfo.namespace,
          self.right_node_name))
      print('ip link add {} netns {}{} type veth peer name {} netns {}{}'.format(
        self.left_iface_name,
        appinfo.namespace,
        self.left_node_name,
        self.right_iface_name,
        appinfo.namespace,
        self.right_node_name))
      print('ip netns exec {}{} ip link set {} up'.format(
        appinfo.namespace,
        self.left_node_name, self.left_iface_name))
      print('ip netns exec {}{} ip link set {} up'.format(
        appinfo.namespace,
        self.right_node_name, self.right_iface_name))
      if self.left_node_type == 'docker':
        print('ip netns del {}{}'.format(appinfo.namespace, self.left_node_name))
      if self.right_node_type == 'docker':
        print('ip netns del {}{}'.format(appinfo.namespace, self.right_node_name))
    def check(self):
      if not (self.right_isset and self.left_isset):
        Exception('n2n_link state is not correct')

  if 'nodes' in appinfo.data:
    for node in appinfo.data['nodes']:
      for i in range(len(node['interfaces'])):
        iface = node['interfaces'][i]
        if (iface['type']=='direct'):
          r_nname = iface['args'].split('#')[0]
          r_ifname = iface['args'].split('#')[1]
          peer_found = False
          for link in links:
            if link.right_isset == False:
              b0 = link.left_node_name == r_nname
              b1 = link.left_iface_name == r_ifname
              if b0 and b1: # found peer
                link.right_node_name = node['name']
                link.right_iface_name = iface['name']
                link.right_node_type = node.get('type', 'docker')
                link.right_isset = True
                peer_found = True
          if peer_found == False:
            new_link = n2n_link(node['name'], iface['name'], r_nname, r_ifname)
            new_link.left_node_type = node.get('type', 'docker')
            links.append(new_link)
    for link in links:
      link.check()
      link.koko()

  links = []
  # CREATE Node to Switch LINKS
  class s2n_link:
    node_name = None
    node_iface = None
    node_isset = False
    node_type = None
    switch_name = None
    switch_isset = False
    def __init__(self, n_name, n_iface, n_type, s_name):
      self.node_name = n_name
      self.node_iface = n_iface
      self.node_type = n_type
      self.switch_name = s_name

    def koko(self):
      if self.node_type == 'docker':
        print('kokobr {} {} {}'.format(self.switch_name, self.node_name, self.node_iface))
      elif self.node_type == 'netns':
        print('kokobr_netns {} {} {}'.format(self.switch_name, self.node_name, self.node_iface))
      else:
        print('unknown node-type {}'.format(self.node_type))

    def check(self):
      if not (self.node_isset and self.switch_isset):
        Exception('s2n_link state is not correct')

  if ('nodes' in appinfo.data) and ('switches' in appinfo.data):
    for node in appinfo.data['nodes']:
      for i in range(len(node['interfaces'])):
        iface = node['interfaces'][i]
        if (iface['type']=='bridge'):
          print('\n# connect {} and {}'
              .format(node['name'], iface['args']))
          new_link = s2n_link(node['name'], iface['name'],
                  node.get('type', 'docker'), iface['args'])
          new_link.node_isset = True
          new_link.switch_isset = False
          links.append(new_link)
    for sw in appinfo.data['switches']:
      for i in range(len(sw['interfaces'])):
        iface = sw['interfaces'][i]
        for link in links:
          if link.switch_isset == False:
            b0 = iface['args'] == link.node_name
            b1 = iface['name'] == link.node_iface
            b2 = sw['name'] == link.switch_name
            if b0 and b1 and b2:
              link.switch_isset = True
    for link in links:
      link.check()
      link.koko()

  # Attach PhysNetIf to Conitaner
  if 'nodes' in appinfo.data:
    for node in appinfo.data['nodes']:
      for i in range(len(node['interfaces'])):
        iface = node['interfaces'][i]
        if (iface['type']=='phys'):
          print('koko_physnet {}{} {}'.format(appinfo.namespace, node['name'], iface['name']))

  # Attach raw veth to Conitaner
  if 'nodes' in appinfo.data:
    for node in appinfo.data['nodes']:
      for i in range(len(node['interfaces'])):
        iface = node['interfaces'][i]
        if (iface['type']=='veth'):
          print('ip link add {} type veth peer name {}'.format(
              iface['name'], iface['args']))
          print('koko_physnet {}{} {}'.format(appinfo.namespace, node['name'], iface['name']))
          print('ip link set {} up'.format(iface['args']))
  postinit()

def command_down(args):
  def prefini():
    print('\n#####################')
    print('# PRE-FINI COMMANDS #')
    print('#####################')
    if 'prefini' in appinfo.data:
      for node in appinfo.data['prefini']:
        for cmd in node['cmds']:
          print('{}'.format(cmd['cmd']))

  def postfini():
    print('\n#####################')
    print('# POST-FINI COMMANDS #')
    print('#####################')
    if 'postfini' in appinfo.data:
      for node in appinfo.data['postfini']:
        for cmd in node['cmds']:
          print('{}'.format(cmd['cmd']))

  precmd()
  generate_mount_docker_netns()
  generate_detach_physif_from_docker()

  print('\n################')
  print('# FINI Commands  #')
  print('################')
  if 'nodes' in appinfo.data:
    for node in appinfo.data['nodes']:
      for i in range(len(node['interfaces'])):
        iface = node['interfaces'][i]
        if (iface['type']=='phys'):
          print('detach_physif_from_docker {}{} {}'.format(appinfo.namespace, node['name'], iface['name']))
    for node in appinfo.data['nodes']:
      nodetype = node.get('type', 'docker')
      if nodetype=='docker':
        print('docker rm -f {}{}'.format(appinfo.namespace, node['name']))
      elif nodetype=='netns':
        print('ip netns del {}{}'.format(appinfo.namespace, node['name']))
      else:
        print('unknown nodetype {}'.format(nodetype))
        sys.exit(1)
  if 'switches' in appinfo.data:
    for sw in appinfo.data['switches']:
      print('ovs-vsctl del-br {}{}'.format(appinfo.namespace, sw['name']));
  postfini()

def command_pull(args):
  precmd()
  print('{}()'.format(sys._getframe().f_code.co_name))
  print('\n######################')
  print('# PULL DOCKER IMAGES #')
  print('######################')
  if 'nodes' in appinfo.data:
    for node in appinfo.data['nodes']:
      if node['type'] == 'docker':
        print('docker pull {}'.format(node['image']))

def command_exec(args):
  node = None
  for n in appinfo.data['nodes']:
    if n['name'] == args.NODE:
      node = n
      break
  if node == None:
    sys.stderr.write('no such node...\n')
    sys.exit(1)
  cmdstr = ''
  node_type = node.get('type', 'docker')
  if node_type == 'docker':
    cmdstr = 'docker exec -it {}{}'.format(appinfo.namespace, args.NODE)
  elif node_type == 'netns':
    cmdstr = 'ip netns exec {}{}'.format(appinfo.namespace, args.NODE)
  else:
    sys.stderr.write('no such node type...\n')
    sys.exit(1)
  for s in args.COMMAND:
      cmdstr = '{} {}'.format(cmdstr, s)
  print('{}'.format(cmdstr))

def command_build(args):
  print('{}()'.format(sys._getframe().f_code.co_name))
  sys.stderr.write('ATODE\n')
  sys.exit(1)

def command_conf(args):
  def preconf(data):
    print('\n#####################')
    print('# PRE-CONF COMMANDS #')
    print('#####################')
    if 'preconf' in appinfo.data:
      for node in appinfo.data['preconf']:
        for cmd in node['cmds']:
          print('{}'.format(cmd['cmd']))

  def postconf(data):
    print('\n#####################')
    print('# POST-CONF COMMANDS #')
    print('#####################')
    if 'postconf' in appinfo.data:
      for node in appinfo.data['postconf']:
        for cmd in node['cmds']:
          print('{}'.format(cmd['cmd']))

  def get_node_type_from_nodes(name, default):
    if 'nodes' in appinfo.data:
      for node in appinfo.data['nodes']:
        if node['name'] == name:
          return node.get('type', default)
    print('not found node {}'.format(name))
    sys.exit(1)

  print("set -ue")
  precmd()
  preconf(appinfo.data)
  print('\n################')
  print('# NODE CONFIG  #')
  print('################')
  if 'node_configs' in appinfo.data:
    for node in appinfo.data['node_configs']:
      print('echo {}::[config]::start'.format(node['name']))
      nodetype = get_node_type_from_nodes(node['name'], 'docker')
      for cmd in node['cmds']:
        if nodetype=='docker':
          print('docker exec {}{} {} > /dev/null'.format(appinfo.namespace, node['name'], cmd['cmd']))
        elif nodetype=='netns':
          print('ip netns exec {}{} {} > /dev/null'.format(appinfo.namespace, node['name'], cmd['cmd']))
        else:
          print('unknown node-type {}'.format(nodetype))
      print('echo {}::[config]::fin'.format(node['name']))
  postconf(appinfo.data)

def command_upconf(args):
  precmd()
  command_up(args)
  command_conf(args)

def command_reconf(args):
  precmd()
  command_reup(args)
  command_conf(args)

def command_reup(args):
  precmd()
  command_down(args)
  command_up(args)

def command_version(args):
  print(appinfo.version)

def command_test(args):
  precmd()
  test_name = args.NAME
  print('\n#####################')
  print('# INTEGRATION TEST  #')
  print('#####################')
  if 'test' in appinfo.data:
    print('set -eu')
    for node in appinfo.data['test']:
      if (node['name'] == test_name):
        for cmd in node['cmds']:
          cmd_str = cmd['cmd']
          print('echo -n execute [{}]...'.format(cmd_str))
          print('{} >> /dev/null'.format(cmd_str))
          print('echo done\n')
        return
    print('test \'{}\' was not found'.format(test_name))

def command_postinit(args):
  print("set -ue")
  precmd()
  postinit()

def command_init(args):
  print(template_spec)

def command_img(args):
  def is_exist(graph, me):
    edge_list = graph.get_edge_list()
    for edge in edge_list:
      pair = edge.obj_dict['points']
      me_pair = me.obj_dict['points']
      b0 = (pair[0]==me_pair[0] and pair[1]==me_pair[1])
      b1 = (pair[1]==me_pair[0] and pair[0]==me_pair[1])
      if (b0 or b1):
        return True
    return False

  graph = pydot.Dot(graph_type='graph', overlap=False, splines=True)
  for i in range(len(appinfo.data['nodes'])):
    node_name = appinfo.data["nodes"][i]["name"]
    graph.add_node(pydot.Node(node_name, label=node_name))
    iface_info = appinfo.data["nodes"][i]["interfaces"]
    for j in range(len(iface_info)):
      args_name = iface_info[j]['args'].split("#")[0]
      new_edge = pydot.Edge(node_name, args_name)
      new_edge = pydot.Edge(node_name, args_name)
      iface_name_label = iface_info[j]['name']
      new_edge.set_headlabel(iface_info[j]['args'].split("#")[1])
      new_edge.set_taillabel(iface_info[j]['name'])
      new_edge.set_fontsize(8)
      if not is_exist(graph, new_edge):
        graph.add_edge(new_edge)
  graph.write_png('topo.png', prog='sfdp')

def main():
  parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter)
  parser.add_argument('-f', '--specfile', default='spec.yaml')
  parser.add_argument('-H', '--host', default=None)
  parser.add_argument('--project-name', default='')
  parser.add_argument('--verbose', action='store_true')
  parser.add_argument('--dry-run', action='store_true')
  # usage_text='''
  #   tn [-f <arg>...] [options] [COMMAND] [ARGS...]
  #   tn -h|--help
  #
  # Options:
  #   -f --specfile NAME    Specify specification yaml file
  #   -v, --verbose         Generate verbose shell
  #   --dry-run             Print the recipes that are needed to execute the
  #                         targets up to date, but not actually execute them.
  #   --project-name NAME   Specify an alternate project name
  #                         (default: none)
  #   --host HOST           Daemon socket to connect to
  #   img        Generate topology png file
  # '''

  subparsers = parser.add_subparsers(dest='command')
  subparsers.required = True

  parser_ps = subparsers.add_parser('ps',
          help='List services')
  parser_ps.set_defaults(func=command_ps)

  parser_up = subparsers.add_parser('up',
          help='Create and start containers')
  parser_up.set_defaults(func=command_up)

  parser_down = subparsers.add_parser('down',
          help='Stop and remove containers')
  parser_down.set_defaults(func=command_down)

  parser_pull = subparsers.add_parser('pull',
          help='Pull service images')
  parser_pull.set_defaults(func=command_pull)

  parser_exec = subparsers.add_parser('exec',
          help='Execute a command in a running container')
  parser_exec.add_argument('NODE')
  parser_exec.add_argument('COMMAND', nargs='+')
  parser_exec.set_defaults(func=command_exec)

  parser_build = subparsers.add_parser('build',
          help='Generate a Docker bundle from the spec file')
  parser_build.set_defaults(func=command_build)

  parser_conf = subparsers.add_parser('conf',
          help='Execute config-cmd in a running container')
  parser_conf.set_defaults(func=command_conf)

  parser_reconf = subparsers.add_parser('upconf',
          help='Create, start and config')
  parser_reconf.set_defaults(func=command_upconf)

  parser_reconf = subparsers.add_parser('reconf',
          help='Stop, remove, create, start and config')
  parser_reconf.set_defaults(func=command_reconf)

  parser_reup = subparsers.add_parser('reup',
          help='Stop, remove, create and start')
  parser_reup.set_defaults(func=command_reup)

  parser_version = subparsers.add_parser('version',
          help='Show the tinet version information')
  parser_version.set_defaults(func=command_version)

  parser_test = subparsers.add_parser('test',
          help='Execute tests')
  parser_test.set_defaults(func=command_test)
  parser_test.add_argument('NAME')

  parser_init = subparsers.add_parser('init',
          help='Generate template spec file')
  parser_init.set_defaults(func=command_init)

  parser_img = subparsers.add_parser('img',
          help='Generate topology png file')
  parser_img.set_defaults(func=command_img)

  parser_postinit = subparsers.add_parser('postinit',
          help='Generate post init script')
  parser_postinit.set_defaults(func=command_postinit)

  args = parser.parse_args()

  # print('------------')
  # print('verbose {}'.format(args.verbose))
  # print('specfile {}'.format(args.specfile))
  # print('host {}'.format(args.host))
  # print('pjname {}'.format(args.project_name))
  # print('command {}'.format(args.command))
  # print('------------')

  appinfo.verbose=args.verbose
  appinfo.specfile=args.specfile
  appinfo.host=args.host
  appinfo.name=args.project_name
  appinfo.command=args.command

  need_data = {'ps', 'up', 'down', 'pull', 'exec',
    'build', 'conf', 'upconf', 'reconf', 'reup', 'test', 'img',
    'postinit'}
  if args.command in need_data:
    try:
      appinfo.read_yaml(appinfo.specfile)
    except IOError:
      sys.stderr.write('IOError: file {} not found.\n'.format(appinfo.specfile))
      sys.exit(1)
    except:
      sys.stderr.write('Yaml {} maybe invalid...?\n'.format(appinfo.specfile))
      sys.exit(1)
  args.func(args)

main()