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


import datetime
import icmperror
import pscheduler
import re

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']


run_timeout = datetime.timedelta()


#
# 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)))
run_timeout += ( count * 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 = 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)

# 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("Program exited %d: %s",
          status, stdout if status == 0 else stderr)

# Ping exits 1 if the host isn't up, which for us is still a success.
if not status in [0, 1]:
    pscheduler.succeed_json( {
            'succeeded': False,
            'diags': argv_string,
            'error': stderr,
            'result': None
            } )


#
# Dissect the results
#

final_result = {
    'schema': 1,
    'succeeded': True,
    'diags': argv_string + '\n\n' + stdout,
    'error': None
}

result = {
    'schema': 1
    }


roundtrips = []



#
# Matchers for output lines we care about
#

packet_returned = re.compile("""^\
([0-9]+) bytes\
\s+from (.*):\
\s+icmp_[rs]eq=([0-9]+)\
\s+ttl=([0-9]+)\
\s+time=([0-9.]+) ms\
$""")

# Error:   From 5.6.7.8 icmp_seq=7 Destination Net Unreachable

error_returned = re.compile("""^\
From\
\s+([^\s]+)\
\s+icmp_seq=([0-9]+)\
\s+(.*)\
$""")

# Error messages known to be output by ping
error_strings = {
    "Destination Net Unreachable": 'net-unreachable',
    "Destination Host Unreachable": 'host-unreachable',
    "Destination Protocol Unreachable": 'protocol-unreachable',
    "Destination Port Unreachable": 'port-unreachable',
    "Source Route Failed": 'source-route-failed',
    "Packet filtered": 'communication-administratively-prohibited',    
    }

# Same, but require regexps to match
error_matches = [
    ( re.compile("^Frag needed and DF set"), 'fragmentation-needed-and-df-set' )
    ]

# Times:   rtt min/avg/max/mdev = 19.631/24.191/29.874/4.262 ms
times_returned = re.compile("""^\
rtt min/avg/max/mdev\s*=\s*\
([0-9.]+)\
/([0-9.]+)\
/([0-9.]+)\
/([0-9.]+)\
\s+ms\
$""")

ips = []

for line in stdout.split('\n'):
    line = line.strip()
    log.debug("LINE: %s", line)

    # Returned Packet

    matches = packet_returned.match(line)
    if matches is not None:
        length, ip, seq, ttl, rtt = matches.groups()
        length = int(length)
        seq = int(seq)
        ttl = int(ttl)
        rtt = float(rtt)
        if seq > count:
            log.warning("Discarding sequence %s; greater than count %s", seq, count)
            continue

        # Under some conditions, ping6 returns the interface with the
        # IP.  Remove that.
        ip = ip.split("%")[0]

        #append as we may get multiple of same seq (duplicates) or seq out of order (reorders)
        roundtrips.append({
            'seq': seq,
            'length': length,
            'ip': ip,
            'ttl': ttl,
            'rtt': pscheduler.timedelta_as_iso8601(
                datetime.timedelta(seconds=rtt/1000.0) )
            })
        log.debug("Packet returned: %s", roundtrips[-1])
        ips.append(ip)
        continue

    # Error

    matches = error_returned.match(line)
    if matches is not None:
        ip, seq, error = matches.groups()
        seq = int(seq)
        if seq > count:
            log.warning("Discarding sequence %s; greater than count %s", seq, count)
            continue

        if error in error_strings:
            error = error_strings[error]
        else:
            error_str = error
            error = None
            for regex, string in error_matches:
                if regex.match(error_str):
                    error = string
                    break

        roundtrips.append({
            'seq': seq,
            'ip': ip,
            'error': error
            })
        log.debug("Error returned: %s", roundtrips[-1])
        ips.append(ip)
        continue
    
    # Final times
    matches = times_returned.match(line)
    if matches is not None:
        min_, mean, max_, stddev = matches.groups()
        result['min'] = pscheduler.timedelta_as_iso8601(
            datetime.timedelta(seconds=float(min_)/1000.0) )
        result['mean'] = pscheduler.timedelta_as_iso8601(
            datetime.timedelta(seconds=float(mean)/1000.0) )
        result['max'] = pscheduler.timedelta_as_iso8601(
                datetime.timedelta(seconds=float(max_)/1000.0) )
        result['stddev'] = pscheduler.timedelta_as_iso8601(
            datetime.timedelta(seconds=float(stddev)/1000.0) )
        log.debug("Got time summary")
        # This is the last line we care about.
        break

    # Anything else we just ignore.

result['roundtrips'] = roundtrips

result['succeeded'] = True

#calculate some stats
packets_seen = {}
sent = count
recvd = 0
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
    recvd += 1
    #check for reorder
    if prev_seq is not None and rt["seq"] < prev_seq:
        reorders += 1
    prev_seq = rt["seq"]
#save results
result['sent'] = sent
result['received'] = recvd
result['lost'] = sent - recvd
result['loss'] = float((sent - recvd)/float(sent)) if sent > 0 else 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 )
