#
# Emulab support for Rhumba
#
#    Sander Vrijders  <sander.vrijders@intec.ugent.be>
#    Wouter Tavernier <wouter.tavernier@intec.ugent.be>
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
# MA  02110-1301  USA

import socket
import paramiko
import time
import os
import re
from ast import literal_eval
import configparser

import warnings
warnings.filterwarnings("ignore")

def get_ssh_client():
    ssh_client = paramiko.SSHClient()
    ssh_client.load_system_host_keys()
    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    return ssh_client

def ops_server(testbed):
    '''
    Return server name of the ops-server (is testbed specific)

    @param testbed: testbed info

    @return: server name of the ops-server
    '''
    return 'ops.' + testbed.url

def full_name(testbed, node_name):
    '''
    Return server name of a node

    @param node_name: name of the node
    @param testbed: testbed info

    @return: server name of the node
    '''
    return node_name + '.' + testbed.exp_name + '.' + \
        testbed.proj_name + '.' + testbed.url

def execute_command(testbed, hostname, command, time_out = 3):
    '''
    Remote execution of a list of shell command on hostname. By
    default this function will exit (timeout) after 3 seconds.

    @param testbed: testbed info
    @param hostname: host name or ip address of the node
    @param command: *nix shell command
    @param time_out: time_out value in seconds, error will be generated if
    no result received in given number of seconds, the value None can
    be used when no timeout is needed

    @return: stdout resulting from the command
    '''
    ssh_client = get_ssh_client()

    try:
        ssh_client.connect(hostname, 22,
                           testbed.username, testbed.password,
                           look_for_keys = True, timeout = time_out)
        stdin, stdout, stderr = ssh_client.exec_command(command)
        err = str(stderr.read()).strip('b\'\"\\n')
        if err != "":
            print(err)
        output = str(stdout.read()).strip('b\'\"\\n')
        ssh_client.close()

        return output

    except Exception as e:
        print(str(e))
        return

def copy_file_to_testbed(testbed, hostname, text, file_name):
    '''
    Write a string to a given remote file.
    Overwrite the complete file if it already exists!

    @param testbed: testbed info
    @param hostname: host name or ip address of the node
    @param text: string to be written in file
    @param file_name: file name (including full path) on the host
    '''
    ssh_client = get_ssh_client()

    try:
        ssh_client.connect(hostname, 22,
                           testbed.username,
                           testbed.password,
                           look_for_keys=True)

        cmd = "touch " + file_name + \
              "; chmod a+rwx " + file_name

        stdin, stdout, stderr = ssh_client.exec_command(cmd)
        err = str(stderr.read()).strip('b\'\"\\n')
        if err != "":
            print(err)

        sftp_client = ssh_client.open_sftp()
        remote_file = sftp_client.open(file_name, 'w')

        remote_file.write(text)
        remote_file.close()

    except Exception as e:
        print(str(e))

def get_experiment_list(testbed, project_name = None):
    '''
    Get list of made emulab experiments accessible with your credentials

    @param testbed: testbed info
    @param project_name: optional filter on project

    @return: list of created experiments (strings)
    '''
    cmd = '/usr/testbed/bin/sslxmlrpc_client.py -m experiment getlist'
    out = execute_command(testbed, ops_server(testbed), cmd)

    try:
        if project_name != None:
            return  literal_eval(out)[project_name][project_name]
        else:
            return  literal_eval(out)
    except:
        return { project_name: { project_name: [] }}

def swap_exp_in(testbed):
    '''
    Swaps experiment in

    @param testbed: testbed info
    '''
    cmd = '/usr/testbed/bin/sslxmlrpc_client.py swapexp proj=' + \
          testbed.proj_name + \
          ' exp=' + \
          testbed.exp_name + \
          ' direction=in'

    output = execute_command(testbed, ops_server(testbed), cmd)

    return output

def create_experiment(testbed, nodes, links):
    '''
    Creates an emulab experiment

    @param testbed: testbed info
    @param nodes: holds the nodes in the experiment
    @param links: holds the links in the experiment
    '''
    proj_name = testbed.proj_name
    exp_name = testbed.exp_name

    exp_list = get_experiment_list(testbed)

    try:
        if exp_name in exp_list[proj_name][proj_name]:
            print("Experiment already exists.")
            return
    except:
        print("First experiment to be created for that project.")

    ns = generate_ns_script(testbed, nodes, links)
    dest_file_name = '/users/'+ testbed.username + \
                     '/temp_ns_file.%s.ns' % os.getpid()
    copy_file_to_testbed(testbed, ops_server(testbed), ns, dest_file_name)

    cmd = '/usr/testbed/bin/sslxmlrpc_client.py startexp ' + \
          'batch=false wait=true proj="' + proj_name + \
          '" exp="' + exp_name + '" noswapin=true ' + \
          'nsfilepath="' + dest_file_name + '"'

    execute_command(testbed, ops_server(testbed), cmd, time_out = None)
    execute_command(testbed, ops_server(testbed),'rm ' + dest_file_name)
    print("New experiment succesfully created.")

def generate_ns_script(testbed, nodes, p2plinks):
    '''
    Generate ns script based on network graph.
    Enables to customize default node image.

    @param nodes: holds the nodes in the experiment
    @param links: holds the links in the experiment
    @param testbed: testbed info

    @return: ns2 script for Emulab experiment
    '''

    ns2_script = "# ns script generated by Rhumba\n"
    ns2_script += "set ns [new Simulator]\n"
    ns2_script += "source tb_compat.tcl\n"

    for node in nodes:
            ns2_script += "set " + node.name + " [$ns node]\n"
            ns2_script += "tb-set-node-os $" + node.name + " " + \
                          testbed.image + "\n"

    for link in p2plinks:
            ns2_script += "set " + link.name + \
                          " [$ns duplex-link $" + \
                          link.node_a.name + " $" + \
                          link.node_b.name + " 1000Mb 0ms DropTail]\n"

    ns2_script += "$ns run\n"

    return ns2_script

def wait_until_nodes_up(testbed):
    '''
    Checks if nodes are up

    @param testbed: testbed info
    '''
    print("Waiting until all nodes are up")

    cmd = '/usr/testbed/bin/script_wrapper.py expinfo -e' + \
          testbed.proj_name + \
          ',' + \
          testbed.exp_name + \
          ' -a | grep State | cut -f2,2 -d " "'

    res = execute_command(testbed, ops_server(testbed), cmd)
    active = False
    if res == "active":
        active = True
    while active != True:
        res = execute_command(testbed, ops_server(testbed), cmd)
        if res == "active":
            active = True
        print("Still waiting")
        time.sleep(5)

def complete_experiment_graph(testbed, nodes, p2plinks):
    '''
    Gets the interface (ethx) to link mapping

    @param testbed: testbed info
    @param nodes: holds the nodes in the experiment
    @param links: holds the links in the experiment
    '''

    node_full_name = full_name(testbed, nodes[0].name)
    cmd = 'cat /var/emulab/boot/topomap'
    topomap = execute_command(testbed, node_full_name, cmd)
    # Almost as ugly as yo momma
    index = topomap.rfind("# lans")
    topo_array = topomap[:index].split('\\n')[1:-1]
    # Array contains things like 'r2b1,link7:10.1.6.3 link6:10.1.5.3'
    for item in topo_array:
        item_array = re.split(',? ?', item)
        node_name = item_array[0]
        for item2 in item_array[1:]:
            item2 = item2.split(':')
            link_name = item2[0]
            link_ip = item2[1]
            for link in p2plinks:
                if link.name == link_name:
                    if link.node_a.name == node_name:
                        link.int_a.ip = link_ip
                    elif link.node_b.name == node_name:
                        link.int_b.ip = link_ip

    for node in nodes:
        cmd = 'cat /var/emulab/boot/ifmap'
        node_full_name  = full_name(testbed, node.name)
        output = execute_command(testbed, node_full_name, cmd)
        output = re.split('\\\\n', output)
        for item in output:
            item = item.split()
            for link in p2plinks:
                if link.node_a.name == node.name and \
                   link.int_a.ip == item[1]:
                    link.int_a.name = item[0]
                elif link.node_b.name == node.name and \
                     link.int_b.ip == item[1]:
                    link.int_b.name =  item[0]

def setup_vlan(testbed, node_name, vlan_id, int_name):
    '''
    Gets the interface (ethx) to link mapping

    @param testbed: testbed info
    @param node_name: the node to create the VLAN on
    @param vlan_id: the VLAN id
    @param int_name: the name of the interface
    '''
    print("Setting up VLAN on node " + node_name)

    node_full_name = full_name(node_name, testbed)
    cmd = "sudo ip link add link " + \
          str(int_name) + \
          " name " + str(int_name) + \
          "." + str(vlan_id) + \
          " type vlan id " + str(vlan_id)
    execute_command(testbed, node_full_name, cmd)
    cmd = "sudo ifconfig " + \
          str(int_name) + "." + \
          str(vlan_id) + " up"
    execute_command(node_full_name, cmd, testbed)
    cmd = "sudo ethtool -K " + \
          str(int_name) + " rxvlan off"
    execute_command(node_full_name, cmd, testbed)
    cmd = "sudo ethtool -K " + \
          str(int_name) + " txvlan off"
    execute_command(node_full_name, cmd, testbed)