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

import datetime
import json
import sys
import time

import pscheduler
import dns.reversename
import dns.resolver

from dnspy_defaults import DEFAULT_TIMEOUT

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

input = pscheduler.json_load(exit_on_error=True);
resolver = dns.resolver.Resolver()


try:
    resolver.nameservers = [ input['test']['spec']['nameserver'] ]
except KeyError:
    pass  # Not there?  Don't care.


try:
    query = input['test']['spec']['query']
    record = input['test']['spec']['record']
except KeyError:
    pscheduler.fail('missing data in input')

timeout_iso = input['test']['spec'].get("timeout", DEFAULT_TIMEOUT )
timeout = pscheduler.timedelta_as_seconds( pscheduler.iso8601_as_timedelta(timeout_iso) )

# Timeout governs the entire query.  Because we don't know how many
# nameservers there will be until runtime, the only sane way to deal
# with individual timeouts is for each to get a fraction of the whole.
resolver.lifetime = timeout
resolver.timeout = timeout / len(resolver.nameservers)


# Perform the test

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

start_time = datetime.datetime.now()
error = None
succeeded = False
answers = None
try:
    answers = resolver.query( query, record )
    succeeded = True
except (dns.exception.Timeout):
    error = 'Timeout'
except (dns.resolver.NoAnswer ):
    error = 'No answer'
except ( dns.resolver.NXDOMAIN ):
    error = 'Domain does not exist'

end_time = datetime.datetime.now()

if answers is None:
    result_rec = None
elif record in [ "a", "aaaa", "ns", "txt" ]:
    result_rec = [ str( ip ) for ip in answers ]
elif record in [ "cname", "ptr" ]:
    result_rec = str( answers[ 0 ])
elif record in [ "mx" ]:
    result_rec = [ {
        "pref" : mx.preference,
        "mx" : str( mx.exchange )
    } for mx in answers ]
elif record in [ "soa" ]:
    soa = answers[ 0 ]
    result_rec = {
        "nameserver" : str( soa.mname ),
        "owner" : str( soa.rname ),
        "serial" : soa.serial,
        "refresh" : soa.refresh,
        "retry" : soa.retry,
        "expire" : soa.expire,
        "minimum" : soa.minimum
    }

#
# Produce results
#

results = {
    'succeeded': succeeded,
    'diags': '',
    'error': error,
    'result': {
        'schema': 1,
        'time': pscheduler.timedelta_as_iso8601( end_time - start_time ),
        'succeeded' : succeeded
    }
}


if result_rec is not None:
    results [ 'result' ][ 'record' ] = result_rec

pscheduler.succeed_json(results)
