#!/bin/env python
#
# check_js - job submission/wn testing based on jess
#
# Author: Marian Babik
# Copyright CERN 2021
#
import logging
import pexpect
import os
import re
import nap
import nap.core

log = logging.getLogger()

app = nap.core.Plugin(description='This plugin tests grid job submission with configurable payload.')
app.add_argument('-T', '--token', help='Token to use for submitting the job', required=True)


def run(cmd, cmd_timeout=300, log_file=None):
    log.debug("    Subprocess %s starting" % cmd)
    try:
        cmd_out, ret_code = pexpect.run(' '.join(cmd), timeout=cmd_timeout, withexitstatus=True, env=os.environ,
                                        logfile=log_file)
    except pexpect.ExceptionPexpect as e:
        log.exception("    Subprocess %s failed with (%s)" % (cmd, e))
        return False, None, e

    log.debug("    Subprocess {} ({}) finished".format(cmd, ret_code))
    log.debug("    Subprocess output: %s " % cmd_out)

    return True, ret_code, cmd_out


@app.metric(metric_name="Ping", seq=1)
def test_ce_ping(args, io):
    os.environ['BEARER_TOKEN_FILE'] = args.token
    os.environ['_condor_SEC_CLIENT_AUTHENTICATION_METHODS'] = 'SCITOKENS'
    os.environ['_condor_TOOL_DEBUG'] = '"D_CAT D_ALWAYS:2 D_SECURITY:2"'
    log.debug(os.environ)
    st, ret_code, cmd_out = run(['/bin/condor_ce_ping -verbose -debug -name {0} -pool {0}:9619 WRITE'.format(
                                 args.hostname)])
    if not st:
        io.status = nap.CRITICAL
        io.summary = 'Failed to run condor_ce_ping, exception was caught: {}'.format(cmd_out)
        return

    print(ret_code)
    print(cmd_out)
    match = re.search("Authorized:.*TRUE", cmd_out)
    if match:
        io.summary = "Scitokens authentication succeeded"
        io.status = nap.OK
        return

    io.status = nap.CRITICAL
    io.summary = 'Failed to authenticate using scitokens'


if __name__ == '__main__':
    app.run()
