#!/usr/bin/env python3
#
# Run a powstream test
#

import atexit
import datetime
import fcntl
import json
import pscheduler
import sys
import os
import signal
import time
import pytz
from powstream_defaults import *
from powstream_utils import get_config, parse_raw_owamp_output, cleanup_dir, cleanup_file, handle_run_error, sleep_or_end, graceful_exit
from subprocess import Popen, PIPE, call

#track when this run starts - make sure it is aware that it is UTC
start_time = datetime.datetime.utcnow().replace(tzinfo=pytz.utc)
#create tag to include in data directory based on current time
#prevents old processes from stomping on new during restart
time_tag = start_time.strftime("%Y%b%dT%H%M%S%f")

#Init logging
log = pscheduler.Log(prefix="tool-powstream", quiet=True)

#DEBUGGING: Set static values below
# task_uuid = 'ABC123'
# participant = 0
# participant_data = [{}, {'ctrl-port': 861}]
# test_spec = {'source': '10.0.1.28', 'dest': '10.0.1.25'}
# duration = pscheduler.iso8601_as_timedelta('PT2M')
# input = { 'schedule': {'until': "2016-09-15T15:22:40Z" }}

#parse JSON input
input = pscheduler.json_load(exit_on_error=True)
try:
    task_uuid = input['task-uuid']
    participant = input['participant']
    participant_data = input['participant-data']
    test_spec = input['test']['spec']
    duration = pscheduler.iso8601_as_timedelta(input['schedule']['duration'])
except KeyError as ex:
    pscheduler.fail("Missing required key in run input: %s" % ex)
except:
    pscheduler.fail("Error parsing run input: %s" % sys.exc_info()[0])

#check if given start time
if ('schedule' in input) and ('start' in input['schedule']):
    start_time = pscheduler.iso8601_as_datetime(input['schedule']['start'])
    
#set end time - use lesser of now + duration and until time
#NOTE: I don't think we are given until time anymore, leaving this here anyways in case it gets added back
#      Assuming it doesn't we are always given a start and duration, so should be sufficient
end_time = start_time + duration
if ('schedule' in input) and ('until' in input['schedule']):
    until_time = pscheduler.iso8601_as_datetime(input['schedule']['until'])
    if until_time < end_time:
        end_time = until_time
log.debug("Powstream run ends at %s" % end_time)
#determine whether this is a reverse test
flip = test_spec.get('flip', False)

#setup infrastructure for killing processes on normal exit
proc = None
def terminate_proc():
    try:
        if proc:
            proc.terminate()
            log.debug("Terminated powstream process %s" % proc.pid)
    except:
        pass
atexit.register(terminate_proc)

#constants
ADDR_FORMAT = "[%s]:%d"
DEFAULT_POWSTREAM_CMD = '/usr/bin/powstream'
DEFAULT_OWSTATS_CMD = '/usr/bin/owstats'
DEFAULT_PKILL_CMD = '/usr/bin/pkill'
DEFAULT_DATA_DIR = '/var/lib/pscheduler/tool/powstream'
DEFAULT_BUCKET_WIDTH = TIME_SCALE #convert to ms
DEFAULT_RAW_OUTPUT = False #don't display raw packets by default
POWSTREAM_RANGE_ARGS = [
    ('data-ports', '-P'),
]
POWSTREAM_VAL_ARGS = [
    ('ip-tos', '-D'),    
    ('packet-padding', '-s')
]

#read config file
config = get_config()

powstream_cmd = DEFAULT_POWSTREAM_CMD
if config and config.has_option(CONFIG_SECTION, CONFIG_OPT_POWSTREAM_CMD):
    powstream_cmd = config.get(CONFIG_SECTION, CONFIG_OPT_POWSTREAM_CMD)
    
owstats_cmd = DEFAULT_OWSTATS_CMD
if config and config.has_option(CONFIG_SECTION, CONFIG_OPT_OWSTATS_CMD):
    owstats_cmd = config.get(CONFIG_SECTION, CONFIG_OPT_OWSTATS_CMD)

pkill_cmd = DEFAULT_PKILL_CMD
if config and config.has_option(CONFIG_SECTION, CONFIG_OPT_PKILL_CMD):
    pkill_cmd = config.get(CONFIG_SECTION, CONFIG_OPT_PKILL_CMD)
    
keep_data_files = False
if config and config.has_option(CONFIG_SECTION, CONFIG_OPT_KEEP_DATA_FILES):
    keep_data_files = config.getboolean(CONFIG_SECTION, CONFIG_OPT_KEEP_DATA_FILES)


# Determine data directory in this order of preference:
#  1. Config file's data_dir
#  2. Environment's TMPDIR
#  3. Contents of DEFAULT_DATA_DIR

if config and config.has_option(CONFIG_SECTION, CONFIG_OPT_DATA_DIR):
    parent_data_dir = config.get(CONFIG_SECTION, CONFIG_OPT_DATA_DIR)
    log.debug("Using config's temporary space %s", parent_data_dir)
else:
    try:
        parent_data_dir = os.environ['TMPDIR']
        log.debug("Using TMPDIR temporary space %s", parent_data_dir)
    except KeyError:
        parent_data_dir = DEFAULT_DATA_DIR
        log.debug("Using default temporary space %s", parent_data_dir)

if not parent_data_dir.endswith("/"):
    parent_data_dir += "/"

data_dir = parent_data_dir + task_uuid + '-' + time_tag
log.debug("Data directory is %s", data_dir)
os.makedirs(data_dir)



#Always print files (-p)
powstream_args = [powstream_cmd, '-p', '-d', data_dir]

#register various handlers that make sure data dir is removed on exit
atexit.register(cleanup_dir, data_dir, keep_data_files=keep_data_files)
cleanup_handler = lambda signum, frame: graceful_exit(data_dir, keep_data_files=keep_data_files, pkill_cmd=pkill_cmd)
signal.signal(signal.SIGTERM, cleanup_handler)

# set log level if needed
if config and config.has_option(CONFIG_SECTION, CONFIG_OPT_LOG_LEVEL):
    log_level = config.get(CONFIG_SECTION, CONFIG_OPT_LOG_LEVEL)
    powstream_args.append('-g')
    powstream_args.append(log_level)
    
#build basic arguments
for arg in POWSTREAM_VAL_ARGS:
    if arg[0] in test_spec:
        powstream_args.append(arg[1])
        powstream_args.append(str(test_spec[arg[0]]))
for rarg in POWSTREAM_RANGE_ARGS:
    if rarg[0] in test_spec:
        powstream_args.append(rarg[1])
        powstream_args.append("%d-%d" % (test_spec[rarg[0]]['lower'], test_spec[rarg[0]]['upper']))
        
#set interval,count and timeout to ensure consistent with duration
count = test_spec.get('packet-count', DEFAULT_PACKET_COUNT)
interval = test_spec.get('packet-interval', DEFAULT_PACKET_INTERVAL)
packet_timeout = test_spec.get('packet-timeout', 0)
powstream_args.append('-c')
powstream_args.append(str(count))
powstream_args.append('-i')
powstream_args.append(str(interval))
if packet_timeout > 0:
    powstream_args.append('-L')
    powstream_args.append(str(packet_timeout))
#calculate min time between results
result_sleep = count * interval + packet_timeout

#set if ipv4 only or ipv6 only
ip_version = test_spec.get('ip-version', None)
if ip_version is not None:
    powstream_args.append(f'-{ip_version}')

#bucket width is used for rounding delay values used as buckets for histogram
bucket_width = test_spec.get('bucket-width', DEFAULT_BUCKET_WIDTH)

#determine whether we will return raw packets
raw_output = test_spec.get('output-raw', DEFAULT_RAW_OUTPUT)

#determine control port
control_port = int(test_spec.get('ctrl-port', DEFAULT_OWAMPD_PORT))

# Normalize the source and destination if both are present

if ('source' in test_spec) and ('dest' in test_spec):

    source_ip, dest_ip = pscheduler.ip_normalize_version(
        test_spec['source'],
        test_spec['dest'],
        ip_version=ip_version
    )
    log.debug(f'Normalized source/dest to {source_ip} -> {dest_ip}')
    if source_ip is None or dest_ip is None:
        pscheduler.succeed_json({
            "succeeded": False,
            "error": 'Unable to find common IP version between the two hosts.'
        })
    test_spec['source'] = source_ip
    test_spec['dest'] = dest_ip
    
#finally, set the addresses and packet flow direction
if flip:
    #reverse test
    if 'dest' in test_spec:
        powstream_args.append('-S')
        powstream_args.append(test_spec['dest'])
    powstream_args.append(ADDR_FORMAT % (test_spec['source'], control_port))
else:
    #forward test
    powstream_args.append('-t')
    if 'source' in test_spec:
        powstream_args.append('-S')
        powstream_args.append(test_spec['source'])
    powstream_args.append(ADDR_FORMAT % (test_spec['dest'], control_port))


#Run the process

emitter = pscheduler.RFC7464Emitter(sys.stdout)

def powstream_line(line, keep_data_files=False):

    log.debug("Processing line '%s'", line)

    try:

        # Any line we get should be a path to a readable file.
        if line.endswith('.owp'):

            log.debug("Processing .owp file %s" % line)

            owstats_args = [owstats_cmd, '-R', line]
            log.debug("Running owstats command: %s" % " ".join(owstats_args))
            try:
                # TODO: The timeout (inherited from the old code) seems long.
                status, out, error = pscheduler.run_program(owstats_args, timeout=30)
                log.debug("owstats returned status %s", status)

                if status == 0:
                    result = parse_raw_owamp_output(out, raw_output=raw_output, bucket_width=bucket_width)
                    log.debug("Posting result: %s", result)
                    emitter(result)
                else:
                    message = "Powstream call to owstats failed: %s" % (error)
                    log.error(message)
                    emitter({
                        "succeeded": False,
                        "error": message
                    })

            except Exception:
                handle_run_error(emitter, "While executing owstats", exception=True)

        else:

            log.debug("Don't care about this file")

    finally:

        # Whatever happened, tidy up
        cleanup_file(line, keep_data_files=keep_data_files)



log.debug("Will run %s", " ".join(powstream_args))

duration_secs = pscheduler.timedelta_as_seconds(duration)
log.debug("Timeout for process is %s seconds", duration_secs)

# Wait for the start time
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, _, err = pscheduler.run_program(powstream_args,
                                        line_call=powstream_line,
                                        timeout=duration_secs,
                                        timeout_ok=True)

if status != 0:
    log.error("Powstream exited %d: %s", status, err)

log.debug("Powestream error output: %s", err)

pscheduler.succeed()
