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


import icmperror
import pscheduler
import re

input = pscheduler.json_load(exit_on_error=True);

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

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

participant = input['participant']

log.debug("Participant %d", participant)
if participant != 0:
    pscheduler.succeed_json( {
            'succeeded': False,
            'diags': None,
            'error': "Invalid participant %d" % participant,
            'result': None
            } )

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




#
# Figure out how to invoke the program
#

argv = [
    'tracepath',
    # Always run without resolving IPs; we'll do that in parallel after it finishes.
    '-n'
]

try:
    ipversion = spec['ip-version']
except KeyError:
    (ipversion, ip) = pscheduler.ip_addr_version(spec['dest'])

if ipversion is not None:
    argv.append(f'-{ipversion}')

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


# At some point, tracepath changed the way its command line works to
# be more compatible with traceroute's options.  Since it doesn't
# provide a way to determine which scheme is in use, the only way to
# figure it out is to run the program, expect a return code of 255 and
# figure out which style of invocation is expects.  That's quality
# with a capital "K."
#
# Earlier:  tracepath [-n] [-l <len>]                <destination>[/<port>]
# Later:    tracepath [-n] [-l <len>] [-b] [-p port] <destination>

# TODO: Investigate whether we care about the -b switch.

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

status, stdout, stderr = pscheduler.run_program( [ 'tracepath' ], timeout=2)
if status != 255 or '<destination>' not in stderr:
    pscheduler.succeed_json( {
            'succeeded': False,
            'diags': None,
            'error': "Unable to determine version of tracepath installed.",
            'result': None
            } )


if '[-p port]' in stderr:
    traceroute_compatible = True
else:
    traceroute_compatible = False


dest = spec['dest']

try:
    port = str(spec['dest-port'])
    if traceroute_compatible:
        argv.append('-p')
        argv.append(port)
    else:
        dest += '/' + str(port)
except KeyError:
    pass

argv.append(dest)

# Force all args to be strings
argv = [str(x) for x in argv]



#
# Run the test
#



# 94 seconds is the worst case plus a second of slop.
status, stdout, stderr = pscheduler.run_program( argv, timeout = 94 )

diags = "\n".join([ " ".join(argv), "", stdout, "", stderr ])

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


#
# Dissect the results
#

try:
    as_ = spec['as']
except KeyError:
    as_ = False

traced_hops = []
ips = []
last_hop = 0

ttl_re = re.compile(r'^(\d*)\??:');
no_reply_re = re.compile(r'no reply');
reached_re = re.compile(r'reached');
rtt_re = re.compile(r'([0-9]+\.[0-9]+)ms');
mtu_re = re.compile(r'pmtu ([0-9]+)');
error_re = re.compile(r'!(\w+)$');
path_mtu = None
for line in stdout.split('\n'):
    line = re.sub(r'\s+', ' ', line).strip()
    matches = ttl_re.match(line)
    if matches is None:
        continue
    ttl =  int(matches.group(1))

    log.debug("LINE %s", line)

    hop = {}
    # Repeats of a hop usually contain more info for first, but replace any repeat info
    if ttl == len(traced_hops):
        hop = traced_hops.pop()

    # No reply means no results
    if no_reply_re.search(line):
        traced_hops.append(hop)
        continue

    # IP.  We forced tracepath to behave this way.
    line_parts = line.split(' ')
    ip = line_parts[1]
    hop['ip'] = ip
    ips.append(ip)
    log.debug("  IP %s", ip)

    # RTT (ms)
    matches = rtt_re.search(line)
    if matches:
        rtt = float(matches.group(1)) / 1000.0
        hop['rtt'] = 'PT%fS' % rtt

    # Path MTU (bytes) - update if changes, otherwise carry over from prev hop
    matches = mtu_re.search(line)
    if matches:
        path_mtu = int(matches.group(1))
    if path_mtu is not None:
         hop['mtu'] = path_mtu

    # Search for errors
    matches = error_re.search(line)
    if matches:
        hop['error'] = icmperror.translate(matches.group(1))

    if ip == '[LOCALHOST]':
        log.debug('  Skipping LOCALHOST line as hop')
        continue

    traced_hops.append(hop)


# If we're doing hostnames, bulk-resolve them.

try:
    hostnames = spec['hostnames']
except KeyError:
    hostnames = True

if hostnames and len(ips) > 0:
    log.debug("Reverse-resolving IPs: %s", str(ips))
    revmap = pscheduler.dns_bulk_resolve(ips, reverse=True, threads=len(traced_hops))
    for hop in traced_hops:
        try:
            ip = hop['ip']
            if ip in revmap and revmap[ip] is not None:
                hop.update({ 'hostname': revmap[ip] })
        except KeyError:
            # No IP is fine.
            pass


# Figure out ASes if we're doing that

try:
    do_ases = spec['as']
except KeyError:
    do_ases = True

if do_ases:
    ases = pscheduler.as_bulk_resolve(ips, threads=len(ips))
    for index, hop in enumerate(traced_hops):
        try:
            hop_as = ases[hop['ip']]
            if hop_as is None:
                continue

            (asn, owner) = hop_as
            if asn is None:
                continue

            result = { 'number': asn }
            if owner is not None:
                result['owner'] = owner
            traced_hops[index]['as'] = result
        except KeyError:
            pass



# Spit out the results

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