#!/usr/bin/python3
#
# 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);

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

participant = input['participant']

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

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


run_timeout = datetime.timedelta()


#
# Figure out how to invoke the program
#

need_root = False

argv = [
    'traceroute',
    '-q', '1'
    ]


#
# IP Version
#

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

if ipversion is not None:
    argv.append('-' + str(ipversion))


#
# Probe Type
#

probe_types = {
    'udp': None,
    'icmp': '-I',
    'tcp': '-T'
    }

# Hold the probe type for use in TCP and UDP destination ports
try:
    probe_type = spec['probe-type']
    switch = probe_types[probe_type]
    if switch is not None:
        argv.append(switch)
except KeyError:
    probe_type = 'udp'

if probe_type in ['icmp', 'tcp']:
    need_root = True


#
# Fragmentation
#

try:
    if not spec['fragment']:
        argv.append('-F')
except KeyError:
    pass


#
# Flow Label
#

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



#
# First TTL
#

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


#
# Source
#

try:
    source = spec['source']
    # TODO: Needs to be same IP type as dest

    # Traceroute's -s switch doesn't resolve, so that needs to be
    # handled here.
    source_resolved = pscheduler.dns_resolve(source, ip_version=ipversion)
    if source_resolved is None:
        message = "Unable to resolve %s to an IPv%d address" % (source, ipversion)
        pscheduler.succeed_json( {
            'succeeded': False,
            'diags': None,
            'error': message,
            'result': None
            } )

    argv.append('-s')
    argv.append(source_resolved)
except KeyError:
    pass


#
# Hops
#

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

# Do all hops in parallel
argv.append('-N')
argv.append(str(hops))



#
# Hostnames
#

# This is always forced; name resolution is done in bulk after the
# trace completes to make the run time shorter and more predictable.

argv.append('-n')


#
# Destination Port
#

try:
    dest_port = spec['dest-port']
    if probe_type in ['udp', 'tcp']:
        argv.append('-p')
        argv.append(dest_port)
except KeyError:
    pass


#
# TOS
#

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


#
# Wait
#

try:
    wait = pscheduler.iso8601_as_timedelta(spec['wait'])
    argv.append('-w')
    argv.append(pscheduler.timedelta_as_seconds(wait))
except KeyError:
    # Program default
    wait = datetime.timedelta(seconds=5)
    pass
run_timeout += wait


#
# Send Wait
#

try:
    send_wait = spec['sendwait']

    # This has weird behavior, per the manual page:
    #
    # If the value is more than 10, then it specifies a number in
    # milliseconds, else it is a number of seconds (float point values
    # allowed too).  Useful when some routers use rate-limit for icmp
    # messages.
    #
    # That means the rules are:
    #   -- Convert everything to milliseconds
    #   -- If < 10 ms, express as a float

    send_wait_delta = pscheduler.iso8601_as_timedelta(send_wait)
    send_wait_secs = pscheduler.timedelta_as_seconds(send_wait_delta)
    
    if send_wait_secs >= 0.010:
        send_wait_secs *= 1000.0

    argv.append('-z')
    argv.append(str(send_wait_secs))
except KeyError:
    # This will be used in a calculation later.
    send_wait_delta = datetime.timedelta()
    pass

run_timeout += send_wait_delta * hops


#
# AS Resolution
#

try:
    # Dodge a reserved word
    as_ = bool(spec['as'])
    if as_:
        argv.append('-A')
except KeyError:
    pass


#
# Destination (must be last since it's an argument, not a switch))
#


argv.append(spec['dest'])

#
# Packet Length (must be last since it's an argument, not a switch))
#
    
try:
    length = spec['length']
    argv.append(str(length))
except KeyError:
    pass



#
# Run the test
#

if need_root:
    argv.insert(0, 'sudo')

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

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

# Add some run slop
run_timeout += datetime.timedelta(seconds=2)
log.debug("Timeout is %s", run_timeout)

run_timeout_secs = pscheduler.timedelta_as_seconds(run_timeout)

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(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
#

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

traced_hops = []
ips = []
last_hop = 0

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


    matches = re.match(r'^(\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)


    # AS

    if as_:
        astext = elements.pop(0)[1:-1]
        log.debug("AS string %s", astext)
        single_as = astext.split('/')[-1]
        # Should look like ASnnnn or *
        if single_as != '*':
            hop['as'] = int(single_as[2:])


    # 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)



# 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:
                rev = revmap[ip]
                if rev is not None:
                    hop.update({ 'hostname': rev })
        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( {
    'schema': 1,
    'succeeded': True,
    'diags': "%s\n\n%s" % (" ".join(argv), stdout),
    'error': None,
    'result': {
        'schema': 1,
        'succeeded': True,
        'paths': [
            traced_hops
        ]
    }
} )
