#!/usr/bin/python
#
# Run an iperf3 test
#

import datetime
import logging
import json
import pscheduler
import re
import shutil
import sys
import time
import iperf3_parser
import traceback
import ipaddr
import iperf3_utils
from iperf3_defaults import *

# track when this run starts
start_time = datetime.datetime.now()

logger = pscheduler.Log(prefix='tool-iperf3', quiet=True)
logger.debug("starting iperf3 tool")

# parse JSON input
input = pscheduler.json_load(exit_on_error=True)

logger.debug("Input is %s" % input)

try:
    participant = input['participant']
    participant_data = input['participant-data']
    test_spec = input['test']['spec']    
    limits_passed = input.get('limits-passed', [])
except KeyError as e:
    pscheduler.fail("Missing required key in run input: %s" % e)
except Exception:
    pscheduler.fail("Error parsing run input: %s" % sys.exc_info()[0])

if len(participant_data) != 2:
    pscheduler.fail("iperf3 requires exactly 2 participants, got %s" % (len(participant_data)))

config = iperf3_utils.get_config()

# look up our local iperf3 command path
iperf3_cmd  = config["iperf3_cmd"]

# grab the server port from the test spec
server_port = participant_data[1]['server_port']   

# need this to pad out the wait timeout,
# convert to seconds
omit_secs = test_spec.get('omit', 0)
if omit_secs:
    omit_secs = pscheduler.timedelta_as_seconds(pscheduler.iso8601_as_timedelta(omit_secs))         
omit_secs = int(omit_secs)

# convert from ISO to seconds for test duration
test_duration = test_spec.get('duration')
if test_duration:
    delta = pscheduler.iso8601_as_timedelta(test_duration)
    test_duration = int(pscheduler.timedelta_as_seconds(delta))
else:
    test_duration = DEFAULT_DURATION

# if we're doing a reverse test, have to change a few things
reverse = test_spec.get('reverse', False)

def run_client():    

    logger.debug("Waiting %s sec for server on other side to start" % DEFAULT_WAIT_SLEEP)

    time.sleep(DEFAULT_WAIT_SLEEP) #wait for server to start on other side
        
    iperf3_args = [ iperf3_cmd ]

    iperf3_args.append('-p')
    iperf3_args.append(server_port)
    
    #Determine if we need to bind to an address and have enough info to do so intelligently
    ip_version = test_spec.get('ip-version', None)
    source = test_spec.get('source', None)
    destination = test_spec['dest'] #required
    normalized_dest = destination
    local_address = test_spec.get('local-address', None)
    if ip_version is not None:
        #Set ip version if provided
        if ip_version == 4:
            iperf3_args.append('-4')
        else:
            iperf3_args.append('-6')
        #use whatever was provided as the bind
        if local_address is not None:
            iperf3_args.append('-B')
            iperf3_args.append(local_address)
        elif source is not None:
            iperf3_args.append('-B')
            iperf3_args.append(source)
        iperf3_args.append('-c')
        iperf3_args.append(destination) 
    elif local_address is not None:
        # if we have the local-address and dest, we can determine what ip version they have in common
        # we prefer IPv6 but fallback to IPv4 if they don't both have IPv6. If one end only 
        # has IPv4 and the other only has IPv6 we don't bind at all but don't throw error in 
        # case there is some external factor we don't know about (maybe should change this?)
        local_ip, dest_ip = pscheduler.ip_normalize_version(local_address, destination)
        if local_ip is not None:
            iperf3_args.append('-B')
            iperf3_args.append(local_ip)
        if dest_ip is not None:
            iperf3_args.append('-c')
            iperf3_args.append(dest_ip)
            normalized_dest = dest_ip
            ip_version = pscheduler.ip_addr_version(dest_ip)[0]
        else:
            pscheduler.succeed_json({"succeeded": False,
                                     "error": "Unable to find common IP version between local-address %s and dest %s" % (local_address, destination)})
    elif source is not None:
        #see comment for local-address, but s/local-address/source/g
        source_ip, dest_ip = pscheduler.ip_normalize_version(source, destination)
        if source_ip is not None:
            iperf3_args.append('-B')
            iperf3_args.append(source_ip)
        if dest_ip is not None:
            iperf3_args.append('-c')
            iperf3_args.append(dest_ip)
            normalized_dest = dest_ip
            ip_version = pscheduler.ip_addr_version(dest_ip)[0]
        else:
            pscheduler.succeed_json({"succeeded": False,
                                     "error": "Unable to find common IP version between source %s and dest %s" % (source, destination)})
    else:
        #set destination
        iperf3_args.append('-c')
        iperf3_args.append(destination)       

    #set duration
    iperf3_args.append('-t')
    iperf3_args.append(test_duration)

    # always ask for json output, a lot more consistent to parse
    iperf3_args.append("--json")

    # big list of optional arguments to iperf3, map from test spec
    if test_spec.has_key('interval') and test_spec['interval'] != None:
        iperf3_args.append('-i')
        delta = pscheduler.iso8601_as_timedelta(test_spec['interval'])
        iperf3_args.append(int(pscheduler.timedelta_as_seconds(delta)))

    if test_spec.has_key('parallel') and test_spec['parallel'] != None:
        iperf3_args.append('-P')
        iperf3_args.append(test_spec['parallel'])

    if test_spec.has_key('window-size') and test_spec['window-size'] != None:
        iperf3_args.append('-w')
        iperf3_args.append(test_spec['window-size'])

    if test_spec.has_key('mss') and test_spec['mss'] != None:
        iperf3_args.append('-M')
        iperf3_args.append(test_spec['mss'])

    if omit_secs:
        iperf3_args.append('-O')
        iperf3_args.append(omit_secs)

    if test_spec.has_key('bandwidth') and test_spec['bandwidth'] != None:
        iperf3_args.append('-b')
        iperf3_args.append(test_spec['bandwidth'])
    else:

        # Find the largest bandwidth in each participant's passed
        # limits and take the minimum of those values.  This will be
        # the maximum bandwidth allowed by any participant.

        logger.debug("Finding maximum bandwidth")

        max_bandwidths = []

        for passed_set in limits_passed:
            set_max = -1
            logger.debug("Limit set:")
            for passed in passed_set:
                try:
                    new_max = pscheduler.si_as_number(passed['bandwidth']['range']['upper'])
                    logger.debug("  Limit says %d", new_max)
                    set_max = max(set_max, new_max)
                except KeyError:
                    logger.debug("  No upper bandwidth limit specified")
                    continue

            if set_max > -1:
                max_bandwidths.append(set_max)

        if max_bandwidths:
            max_bandwidth = sorted(max_bandwidths)[0]
            logger.debug("Limiting bandwidth to %d", max_bandwidth)
            iperf3_args.append('-b')
            iperf3_args.append(max_bandwidth)



    if test_spec.get('udp', False):
        iperf3_args.append('-u')

    if test_spec.has_key('buffer-length') and test_spec['buffer-length'] != None:
        iperf3_args.append('-l')
        iperf3_args.append(test_spec['buffer-length'])

    if test_spec.has_key('ip-tos') and test_spec['ip-tos'] != None:
        iperf3_args.append('-S')
        iperf3_args.append(test_spec['ip-tos'])

    if test_spec.has_key('congestion') and test_spec['congestion'] != None:
        iperf3_args.append('-C')
        iperf3_args.append(test_spec['congestion'])

    if test_spec.has_key('zero-copy'):
        iperf3_args.append('-Z')

    if test_spec.has_key('flow-label') and test_spec['flow-label'] != None:
        iperf3_args.append('-L')
        iperf3_args.append(test_spec['flow-label'])
    
    if reverse:
        iperf3_args.append('-R')

    # figure out cpu affinity, either use what was passed in
    # or try to auto detect it
    affinity = test_spec.get('client-cpu-affinity')
    if affinity == None:
        affinity = pscheduler.source_affinity(normalized_dest, ip_version=ip_version)

    logger.debug("Selected CPU affinity %s" % affinity)
    if affinity != None:
        iperf3_args.insert(0, 'numactl')
        iperf3_args.insert(1, '-N')
        iperf3_args.insert(2, affinity)

    # join and run_program want these all to be string types, so
    # just to be safe cast everything in the list to a string
    iperf3_args = map(lambda x: str(x), iperf3_args)

    command_line = " ".join(iperf3_args)
    logger.debug("Running command: %s" % command_line)

    iperf_timeout = test_duration
    iperf_timeout += omit_secs
    iperf_timeout += DEFAULT_FUDGE_FACTOR
    # no need for the sleep, we already did that above

    logger.debug("timeout for client is %d" % iperf_timeout)

    try:
        status, stdout, stderr = pscheduler.run_program(iperf3_args,
                                                        timeout = iperf_timeout)
    except Exception as e:
        logger.error("iperf3 failed to complete execution: %s" % e)
        pscheduler.succeed_json({"succeeded": False,
                                 "error": "The iperf3 command failed during execution. See server logs for more details."})



    return _make_result(command_line, status, stdout, stderr)

    

def run_server():

    #init command
    iperf3_args = [ iperf3_cmd, '-s', '-1', '--json']

    iperf3_args.append("-p")
    iperf3_args.append(server_port)

    #Determine if we need to bind to an address and have enough info to do so intelligently
    ip_version = test_spec.get('ip-version', None)
    source = test_spec.get('source', None)
    normalized_dest = test_spec['dest']
    if ip_version is not None:
        #if the ip version is given we just use the dest as the bind address and have 
        #iperf3 force ipv4 or ipv6
        iperf3_args.append('-B')
        iperf3_args.append(test_spec['dest'])
        if test_spec['ip-version'] == 4:
            iperf3_args.append('-4')
        else:
            iperf3_args.append('-6')
    elif source is not None:
        # if we have the source and dest, we can determine what ip version they have in common
        # we prefer IPv6 but fallback to IPv4 if they don't both have IPv6. If one end only 
        # has IPv4 and the other only has IPv6 we don't bind at all but don't throw error in 
        # case there is some external factor we don't know about (maybe should change this?)
        source_ip, bind_address = pscheduler.ip_normalize_version(source, test_spec['dest'])
        if bind_address is not None:
            iperf3_args.append('-B')
            iperf3_args.append(bind_address)
            normalized_dest = bind_address
            ip_version = pscheduler.ip_addr_version(bind_address)[0]
        else:
            logger.debug("Skipping bind. Source %s and destination %s don't have any common IPs of the same version as far as tool can tell." % (source, test_spec['dest']))
    elif test_spec.get('udp', False):
        # We really shouldn't do this but I think we have to because of iperf3 issue #505. 
        # Basically not doing this causes UDP to bind to the first interface instead of all 
        # interfaces which is probably more likely to break things than not specifying. This
        # is imperfect though and we should really advise people to specify the source and/or 
        # ip-version when doing a UDP test.
        iperf3_args.append('-B')
        iperf3_args.append(test_spec['dest'])
        logger.debug("Binding UDP test to %s. If %s is a hostname with both AAAA and A records, it is highly recommended you modify this test to specify ip-version and/or the source address" % (test_spec['dest'], test_spec['dest']))
    else:
        #If we are here its because we have a TCP test where only the dest was specified 
        # and no ip-version provided. In this case we can't determine what we're supposed 
        # to bind to, so don't specify and iperf3 will bind to all interfaces
        logger.debug("Skipping bind. No source or ip-version given so can't make an intelligent decision about where to bind. iperf3 will listen on all interfaces")
        pass
    
    # try to grab our default affinity if one wasn't passed in
    affinity = test_spec.get('server-cpu-affinity')
    if affinity == None:
        # look up what interface we're going to be receiving on        
        interface = pscheduler.address_interface(normalized_dest, ip_version=ip_version)
        if interface:
            affinity = pscheduler.interface_affinity(interface)

    if affinity != None:
        iperf3_args.insert(0, 'numactl')
        iperf3_args.insert(1, '-N')
        iperf3_args.insert(2, affinity)


    iperf3_args = map(lambda x: str(x), iperf3_args)
    logger.debug("Running command: %s" % " ".join(iperf3_args))

    stdout = ""
    stderr = ""
    status = 0

    try:

        iperf_timeout = test_duration
        iperf_timeout += omit_secs
        iperf_timeout += DEFAULT_FUDGE_FACTOR
        iperf_timeout += DEFAULT_WAIT_SLEEP

        logger.debug("timeout for server is %d" % iperf_timeout)

        status, stdout, stderr = pscheduler.run_program(iperf3_args,
                                                        timeout = iperf_timeout)

    except Exception as e:
        logger.error("iperf3 failed to complete execution: %s" % e);
        pscheduler.succeed_json({"succeeded": False,
                                 "error": "The iperf3 command failed during execution. See server logs for more details."})


    return _make_result(" ".join(iperf3_args), status, stdout, stderr)

    

def _make_result(command, status, stdout, stderr):
    if status:
        error_text = ""
        try:
            json_stdout = pscheduler.json_load(stdout)
            error_text = json_stdout["error"]            
        except (ValueError, KeyError):
            pass

        if not error_text:
            error_text = "%s\n%s" % (stdout, stderr)

        pscheduler.succeed_json({"succeeded": False,
                                 "error": "iperf3 returned an error: %s" % (error_text)})

    lines = stdout.split("\n")    
    logger.debug("Lines are %s " % lines)

    results = iperf3_parser.parse_output(lines)
    results['diags'] = "%s\n\n%s" % (command, stdout)

    return results


#determine whether we are the client or server mode for iperf
results = {}
try:
    if participant == 0:
        results = run_client()
    elif participant == 1:
        results = run_server()
    else:
        pscheduler.fail("Invalid participant.")
except Exception as ex:
    _, _, ex_traceback = sys.exc_info()
    if ex_traceback is None:
        ex_traceback = ex.__traceback__
    tb_lines = [ line.rstrip('\n') for line in
                 traceback.format_exception(ex.__class__, ex, ex_traceback)]
    logger.debug(tb_lines)
    logger.error("Exception %s" % ex);

logger.debug("Results: %s" % results)

pscheduler.succeed_json(results)
