#!/usr/bin/python
#
# Run a test.  Just the test spec is provided on stdin.
#


import datetime
import icmperror
import pscheduler
import re

log = pscheduler.Log(prefix='traceroute', quiet=True)

input = pscheduler.json_load(exit_on_error=True);
#input = {"test": {"type": "trace", "spec": {"dest": "10.0.1.17", "schema": 1}}, "participant-data": [{}], "schedule": {"duration": "PT1M11S", "start": "2016-09-02T14:55:35-04:00"}, "participant": 0, "schema": 1}

# TODO: Validate the input
# TODO: Verify can-run

participant = input['participant']

if participant != 0:
    pscheduler.fail("Invalid participant.")

spec = input['test']['spec']


argv = [
    'bwtraceroute',
    '-T', 'traceroute',
    '-y', 'a' #no lookups
    ]


#
# IP Version
#

# TODO: Determine if the destination address is IPv4 or IPv6 if it's
# not a DNS name.

try:
    ipversion = spec['ip-version']
    # TODO: Make sure this is compatible with the destination address
    argv.append('-' + str(ipversion))
except KeyError:
    pass

#
# First TTL
#

try:
    first_ttl = spec['first-ttl']
    argv.append('-F')
    argv.append(str(first_ttl))
except KeyError:
    pass

#
# Hops
#

try:
    hops = spec['hops']
    argv.append('-M')
    argv.append(str(hops))
except KeyError:
    hops = 30

#
# TOS
#

try:
    tos = spec['tos']
    argv.append('-S')
    argv.append(str(tos))
    need_root = True
except KeyError:
    pass



#
# Packet Length 
#
    
try:
    length = spec['length']
    argv.append('-l')
    argv.append(str(length))
except KeyError:
    pass

#
# Source
#

try:
    source = spec['source']
    # TODO: Needs to be same IP type as dest
    argv.append('-s')
    argv.append(source)
except KeyError:
    pass

#
# Destination 
#

argv.append('-c')
argv.append(spec['dest'])


# Stringify the arguments        
argv = map(lambda x: str(x), argv)

log.debug("Running %s" % ' '.join(argv))

run_timeout_secs = 60

status, stdout, stderr \
    = pscheduler.run_program(argv, timeout = run_timeout_secs)

log.debug("Traceroute exited %d: %s",
          status, stdout if status == 0 else stderr)

if status != 0:
    pscheduler.succeed_json( {
            'succeeded': False,
            'diags': " ".join(argv),
            'error': stderr,
            'result': None
            } )


#
# Dissect the results
#
err_lines = stderr.split("\n")
for err_line in err_lines:
    if err_line.strip().startswith("bwctl: Unable to initiate peer handshake"):
        pscheduler.succeed_json({"succeeded": False, "error": err_line.strip()})
    elif err_line.strip().startswith("bwctl: Unable to connect"):
        pscheduler.succeed_json({"succeeded": False, "error": err_line.strip()})
    elif err_line.strip().startswith("bwctl: local tool did not complete in allocated time frame and was killed"):
            pscheduler.succeed_json({"succeeded": False, "error": err_line.strip()})
traced_hops = []
ips = []
last_hop = 0

for line in stdout.split('\n'):
        
    line = re.sub('\s+', ' ', line).strip()

    matches = re.match('^(\d*)\s+(.*)$', line)
    if matches is None:
        log.debug("Discarding: %s" % line)
        continue

    log.debug("Output Line: %s", line)

    this_hop =  int(matches.group(1))

    # Repeats of a hop replace earlier ones
    # TODO: Does this happen in traceroute or only tracepath?
    if this_hop == len(traced_hops):
        log.debug("This hop is a repeat")
        traced_hops.pop()

    elements = matches.group(2).split()
    log.debug("Line elements: %s", elements)

    hop = {}

    # No reply means no results

    if elements[0] == '*':
        log.debug("No reply for this hop")
        traced_hops.append(hop)
        continue

    # IP
    ip = elements.pop(0)
    hop['ip'] = ip
    ips.append(ip)
    
    # RTT (ms)
    rtt = float(elements.pop(0)[:-2]) / 1000.0
    rtt_delta = datetime.timedelta(seconds=rtt)
    hop['rtt'] = pscheduler.timedelta_as_iso8601(rtt_delta)

    # The "ms" afterward
    elements.pop(0)

    # Anything left will be an error

    if len(elements) > 0:
        error = elements.pop(0)
        if error[0] == '!':
            hop['error'] = icmperror.translate(error)

    traced_hops.append(hop)


# Spit out the results

pscheduler.succeed_json( {
    'schema': 1,
    'succeeded': True,
    'diags': stdout,
    'error': None,
    'result': {
        'schema': 1,
        'succeeded': True,
        'paths': [
            traced_hops
        ]
    }
} )
