#!/usr/bin/env python3
#
# Run an iperf2 test
#

# NOTE (see invocations below): Iperf2 has a reputation for not
# returning a nonzero exit code for certain failures.  The only way to
# detect this condition is to look for a non-empty stderr.  This
# causes problems with the server side, which doesn't have a one-shot
# mode and therefore must be left to time out.  The run_program()
# function in the pScheduler library returns 2 on a timeout, and an
# examination of the iperf2 source says nothing exits with that
# status.  Therefore, this has to be taken into account when
# evaluating the exit status.


import datetime
import logging
import json
import pscheduler
import re
import shutil
import sys
import time
import threading
import iperf2_parser
import traceback
import ipaddress
import iperf2_utils
from iperf2_defaults import *

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

logger = pscheduler.Log(prefix='tool-iperf2', quiet=True)

logger.debug("starting iperf2 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']
    duration = pscheduler.iso8601_as_timedelta(input['schedule']['duration'])
except KeyError as ex:
    pscheduler.fail("Missing required key in run input: %s" % str(ex))
except:
    pscheduler.fail("Error parsing run input: %s" % sys.exc_info()[0])

single_ended = test_spec.get('single-ended', False)
loopback = test_spec.get('loopback', False)
participants = len(participant_data)
if not(participants == 2 or (participants == 1 and (single_ended or loopback))):
    pscheduler.fail("iperf2 requires exactly 2 participants, got %s" % (len(participant_data)))


config = iperf2_utils.get_config()

# look up our local iperf2 command path
iperf2_cmd  = config["iperf2_cmd"]

# grab the server port from the test spec if there, otherwise use default
if single_ended:
    server_port = test_spec.get('single-ended-port', DEFAULT_SERVER_PORT)
else:
    if loopback:
        server_port = DEFAULT_SERVER_PORT
    else:
        server_port = participant_data[1].get('server_port', DEFAULT_SERVER_PORT)

# 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



def run_client():    

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

    iperf2_args = [
        iperf2_cmd,
        '-e',
        '--format', 'b'
    ]

    iperf2_args.append('-p')
    iperf2_args.append(server_port)

    if 'mss' in test_spec:
        iperf2_args.append('--mss')
        iperf2_args.append(test_spec['mss'])

    # who to connect to
    destination = test_spec['dest']

    try:
        ipaddress.ip_address(str(destination))
        is_ip_address = True
    except ValueError:
        is_ip_address = False
    
    # iperf2 has no inherent ability to force v4 or v6, so we do it here but
    # don't bother doing lookups on just a straight IP address, seems like
    # a silly combination of flags but whatever
    if test_spec.get('ip-version') is not None:
        if not is_ip_address and test_spec['ip-version'] == 4:
            destination = pscheduler.dns_resolve(destination, ip_version=4)
        
            logger.debug("Resolved %s to %s" % (test_spec["dest"], destination))
                    
            if destination == None:
                pscheduler.succeed_json({"succeeded": False,
                                         "error": "Unable to resolve %s to an IPv4 address" % test_spec["dest"]
                                         })

        elif not is_ip_address and test_spec['ip-version'] == 6:
            destination = pscheduler.dns_resolve(destination, ip_version=6)
                    
            logger.debug("Resolved %s to %s" % (test_spec["dest"], destination))

            if destination == None:
                pscheduler.succeed_json({"succeeded": False,
                                         "error": "Unable to resolve %s to an IPv6 address" % test_spec["dest"]
                                         })



    iperf2_args.append('-c')
    iperf2_args.append(destination)

    # 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

    iperf2_args.append('-t')
    iperf2_args.append(test_duration)

    # always show mss because why not
    iperf2_args.append('-m')

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

    if 'interval' in test_spec and test_spec['interval'] != None:
        iperf2_args.append('-i')
        delta = pscheduler.iso8601_as_timedelta(test_spec['interval'])
        iperf2_args.append(int(pscheduler.timedelta_as_seconds(delta)))

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

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

    if 'bandwidth' in test_spec and test_spec['bandwidth'] != None:
        # Bandwidth throttling is only supported for UDP.  This will
        # have been checked by can_run,
        iperf2_args.append('-b')
        iperf2_args.append(test_spec['bandwidth'])

    if 'fq-rate' in test_spec and test_spec['fq-rate'] is not None:
        iperf2_args.append('--fq-rate')
        iperf2_args.append(test_spec['fq-rate'])

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

    # local-address is deprecated and is overridden by source-bind.

    if test_spec.get('source-bind') is not None:
        iperf2_args.append('-B')
        iperf2_args.append(test_spec['source-bind'])
    else:
        if test_spec.get('local-address') is not None:
            iperf2_args.append('-B')
            iperf2_args.append(test_spec['local-address'])
        

    if 'congestion' in test_spec and test_spec['congestion'] != None:
        iperf2_args.append('-Z')
        iperf2_args.append(test_spec['congestion'])

    
    # join and run_program want these all to be string types, so
    # just to be safe cast everything in the list to a string
    iperf2_args = [str(x) for x in iperf2_args]

    command_line = " ".join(iperf2_args)

    try:
        start_at = input['schedule']['start']
        logger.debug("Sleeping until %s", start_at)
        pscheduler.sleep_until(start_at)
        logger.debug("Starting")
    except KeyError:
        pscheduler.fail("Unable to find start time in input")

    logger.debug("Waiting for server to start")
    time.sleep(DEFAULT_WAIT_SLEEP)

    logger.debug("Running command: %s" % command_line)      
    try:
        status, stdout, stderr = pscheduler.run_program(iperf2_args)
    except Exception:
        pscheduler.succeed_json({"succeeded": False,
                                 "diags": " ".join(iperf2_args),
                                 "error": "%s\n%s" % (stdout, stderr)
                                 })
    
    #see if command completed successfully
    logger.debug("iperf2 returned status %d" % status)

    # See note at the top of this file about evaluating the exit
    # status.
    if status == 2:
        status = 0
        stderr = ""

    if status or stderr:
        pscheduler.succeed_json({"succeeded": False,
                                 "diags": " ".join(iperf2_args),
                                 "error": "iperf2 returned an error: %s" % stderr
                                 })

    logger.debug("Stdout = %s" % stdout)
    logger.debug("Stderr = %s" % stderr)

    lines = stdout.split("\n")    

    logger.debug("Lines are %s " % lines)

    results = iperf2_parser.parse_output(
        lines,
        expect_udp=test_spec.get('udp', False),
        logger=logger
    )
    results['diags'] = "%s\n\n%s\n%s" % (command_line, stdout, stderr)

    return results


def run_server():

    #init command
    iperf2_args = [ iperf2_cmd, '-s' ]

    if 'mss' in test_spec:
        iperf2_args.append('--mss')
        iperf2_args.append(test_spec['mss'])

    iperf2_args.append('-p')
    iperf2_args.append(server_port)

    if test_spec.get('dest-bind') is not None:
        iperf2_args.append('-B')
        iperf2_args.append(test_spec['dest-bind'])

    if test_spec.get('udp'):
        iperf2_args.append('-u')

    # join and run_program want these all to be string types, so
    # just to be safe cast everything in the list to a string
    iperf2_args = [str(x) for x in iperf2_args]


    iperf_timeout = test_duration
    iperf_timeout += iperf2_utils.setup_time(test_spec.get("link-rtt", None))


    try:
        start_at = input['schedule']['start']
        logger.debug("Sleeping until %s", start_at)
        pscheduler.sleep_until(start_at)
        logger.debug("Starting")
    except KeyError:
        pscheduler.fail("Unable to find start time in input")

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

    # See note at the top of this file about evaluating the exit
    # status.
    if status == 2:
        status = 0
        stderr = ""

    if status or stderr:
        pscheduler.succeed_json({"succeeded": False,
                                 "diags": " ".join(iperf2_args),
                                 "error": "iperf2 returned an error: %s" % stderr
                                 })
    
    #log stdout in debug mode
    for line in stdout:
        logger.debug(line)
          
    return {"succeeded": True}



#determine whether we are the client or server mode for iperf2
results = {}
try:
    if participant == 0:
        if loopback:
            server_thread = threading.Thread(target=run_server)
            server_thread.start()
            results = run_client()
            server_thread.join() #Wait until the server thread terminates
        else:
            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)
