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


import datetime
import icmperror
import pscheduler
import re

from ping_utils import *

log = pscheduler.Log(prefix='tool-ping', quiet=True)

input = pscheduler.json_load(exit_on_error=True);

# TODO: Validate the input

participant = input['participant']

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

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


#
# Figure out how to invoke the program
#

need_root = False

argv = []


#
# Get address related parameters
#
## Source
source = spec.get('source', None)

## Destination (required)
try:
    dest = spec['dest']
except KeyError:
    pscheduler.fail("Dest is a required parameter")

## IP Version
ipversion = spec.get('ip-version', None)

# Lookup ip address so ping does not have to and determine ip version

source_ip = None
dest_ip = None

if source is not None:
    (source_ip, dest_ip) = pscheduler.ip_normalize_version(source, dest, ip_version=ipversion)
    if source_ip is None and dest_ip is None:
        pscheduler.fail("Can't find same-family IPs for %s and %s." % (source, dest))
    if dest_ip is not None:
        (ipversion, dest_ip) = pscheduler.ip_addr_version(dest_ip)
elif ipversion is None:
    (ipversion, dest_ip) = pscheduler.ip_addr_version(dest)
    if ipversion is None:
        dest_ip = None

# Make a best guess at the IP if we didn't find one above.
if dest_ip is None:
    dest_ip_map = pscheduler.dns_bulk_resolve([dest], ip_version=ipversion)
    # For failed resolution, just try the hostname and let ping deal with it.
    dest_ip = dest_ip_map.get(dest, dest)

#Determine ping program
if (ipversion is not None) and (ipversion == 6):
    program = 'ping6'
else:
    program = 'ping'
argv.append(program)


# Turn off DNS; we'll do that afterward.
argv.append('-n')


# Count

log.debug("SPEC: %s", spec)
try:
    count = spec['count']
except KeyError:
    count = 5

argv.append('-c')
argv.append(str(count))


# Flow Label  (IPv6 Only)

try:
    flowlabel = spec['flow-label']
    if ipversion == 6:
        argv.append('-F')
        argv.append(str(flowlabel))
    # TODO: Raise an exception if V4?  Can-run should have vetoed running this.
except KeyError:
    pass  # Nothing's just fine.


# Fragmentation
try:
    fragment = spec['fragment']
    argv.append('-M')
    argv.append('want' if fragment else 'do')
except KeyError:
    pass  # Nothing's just fine.

# Interval

try:
    interval = pscheduler.iso8601_as_timedelta(spec['interval'])
except KeyError:
    interval=datetime.timedelta(seconds=1)

if interval < datetime.timedelta(seconds=0.2):
    need_root = True


argv.append('-i')
argv.append(str(pscheduler.timedelta_as_seconds(interval)))


# Interface/Source Address
if source_ip is not None:
    argv.append('-I')
    argv.append(source_ip)



# TODO: Add preload to the spec?  This will require adjustments to the
# duration and requires root if more than 3.


# Multicast Loopback Suppression

try:
    suppress_loopback = spec['suppress-loopback']
    if suppress_loopback:
        argv.append('-L')
except KeyError:
    pass  # Whatever.


# TOS

try:
    ip_tos = spec['ip-tos']
    argv.append('-Q')
    # Ping takes the full byte.
    argv.append(str(ip_tos))
except KeyError:
    pass  # Whatever.


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


# Time to Live
    
try:
    ttl = spec['ttl']
    argv.append('-t')
    argv.append(str(ttl))
except KeyError:
    pass


# Deadline

try:
    deadline = int(pscheduler.timedelta_as_seconds(
        pscheduler.iso8601_as_timedelta(spec['deadline']) ))
    argv.append('-w')
    argv.append(str(deadline))
except KeyError:
    pass


# Timeout
    
try:
    timeout = pscheduler.timedelta_as_seconds(
        pscheduler.iso8601_as_timedelta(spec['timeout']) )
except KeyError:
    timeout = 1.0
argv.append('-W')
argv.append(str(timeout))



# Destination (must be last since it's an argument, not a switch)).
argv.append(str(dest_ip or spec['dest']))



#
# Run the test
#

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

# Stringify the arguments        
argv = [str(x) for x in argv]

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

run_timeout = ping_tool_duration(spec)
run_timeout_secs = pscheduler.timedelta_as_seconds(run_timeout)
log.debug("Timeout is %s", 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("Program exited %d: %s",
          status, stdout if status == 0 else stderr)

diags = f'Exited {status}\n\n{argv_string}\n\nStandard output:\n{stdout}\n\nStandard error:\n{stderr}'

# Ping exits 1 if the host isn't up, which for us is still a success,
# but it also exits 1 on error.  The presence of anything in stderr is
# what signifies a "real" error.

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

# At this point there should be output.
if not stdout:
    pscheduler.succeed_json( {
            'succeeded': False,
            'diags': diags,
            'error': 'No output from ping',
            'result': None
            } )

#
# Dissect the results
#

final_result = {
    'schema': 1,
    'succeeded': True,
    'diags': diags,
    'error': None
}

result = {
    'schema': 1
    }

# Use pscheduler library to parse results into dictionary
parsed_results = pscheduler.parse_ping(stdout, count)

try:
    result['min'] = parsed_results['min']
    result['max'] = parsed_results['max']
    result['mean'] = parsed_results['mean']
    result['stddev'] = parsed_results['stddev']
except KeyError:
    # Don't care if these are missing.
    pass

ips = parsed_results['ips']
roundtrips = parsed_results['roundtrips']
result['roundtrips'] = roundtrips


result['succeeded'] = True

#calculate some stats
packets_seen = {}
dups = 0
reorders = 0
prev_seq = None
for rt in roundtrips:
    #duplicate packet
    if rt["seq"] in packets_seen:
        dups += 1
        continue
    
    #skip lost
    if rt.get('error', None):
        continue
    #made it to other end
    packets_seen[rt["seq"]] = True
    #check for reorder
    if prev_seq is not None and rt["seq"] < prev_seq:
        reorders += 1
    prev_seq = rt["seq"]

#save results
sent = parsed_results.get('sent', 0)
result['sent'] = sent
received = parsed_results.get('received', 0)
result['received'] = received
result['lost'] = sent - received
result['loss'] = parsed_results.get('loss', 100.0)
result['duplicates'] = dups
result['reorders'] = reorders

# 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(roundtrips))
    for hop in roundtrips:
        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


# Spit out the results

final_result['result'] = result
pscheduler.succeed_json( final_result )
