#!/usr/bin/python3

# Helper script to invoke SSM to send accounting record to APEL from ARC Jura
# It relies on SSM libraries (https://github.com/apel/ssm) developed by Will Rogers

import sys
import os
import logging
import argparse

# Initialize root logger
rootlog = logging.getLogger()
rootlog.setLevel(logging.INFO)

# Log to stderr (will be written to jura.log by A-REX)
logger = logging.getLogger('ARC.Jura.SSM')
log_handler_stderr = logging.StreamHandler()
log_handler_stderr.setFormatter(
    logging.Formatter('[%(asctime)s] [%(name)s] [%(levelname)s] [%(process)d] [%(message)s]'))
logger.addHandler(log_handler_stderr)

# ARC-prefix path in PYTHONPATH
arc_prefix_pythonpath = '${prefix}/lib/python3.9/site-packages'.replace('${prefix}', '/usr')
if os.path.isdir(arc_prefix_pythonpath):
    if arc_prefix_pythonpath not in sys.path:
        sys.path.insert(1, arc_prefix_pythonpath)

# Try to use third-party SSM libraries if installed or fallback to re-distributed with NorduGrid ARC
try:
    from ssm import __version__ as ssm_version
    from ssm.ssm2 import Ssm2, Ssm2Exception
    from ssm.crypto import CryptoException
except ImportError:
    from arc.ssm import __version__ as ssm_version
    from arc.ssm.ssm2 import Ssm2, Ssm2Exception
    from arc.ssm.crypto import CryptoException
    logger.info('Using NorduGrid ARC redistributed SSM libraries version %s.%s.%s', *ssm_version)
else:
    logger.info('Using system-wide installed SSM libraries version %s.%s.%s', *ssm_version)


def get_parser():
    parser = argparse.ArgumentParser(description='APEL SSM sender for NorduGrid ARC Jura')
    parser.add_argument('-H', '--host', action='store', help='APEL broker hostname', required=True)
    parser.add_argument('-p', '--port', action='store', help='APEL broker port', required=True)
    parser.add_argument('-m', '--message-path', action='store', help='Path to messages to be sent to APEL',
                        required=True)
    parser.add_argument('-t', '--topic', action='store', default='/queue/global.accounting.cpu.central',
                        help='APEL topic (default is %(default)s)')
    parser.add_argument('-l', '--logfile', action='store', default='/var/spool/arc/ssm/ssmsend.log',
                        help='Dedicated SSM log file (default is %(default)s)')
    parser.add_argument('-d', '--debug', action='store', default='WARNING',
                        help='A-REX logging verbosity level (default is %(default)s)',
                        choices=['DEBUG', 'VERBOSE', 'INFO', 'WARNING', 'ERROR', 'FATAL'])
    parser.add_argument('-s', '--ssl', action='store_true', help='Use SSL connection to brocker')
    parser.add_argument('-c', '--cert', action='store', default='/etc/grid-security/hostcert.pem',
                        help='Path to client certificate (default is %(default)s)')
    parser.add_argument('-k', '--key', action='store', default='/etc/grid-security/hostkey.pem',
                        help='Path to client key (default is %(default)s)')
    parser.add_argument('-C', '--ca-dir', action='store', default='/etc/grid-security/certificates',
                        help='Path to CA certificates directory (default is %(default)s)')
    return parser


if __name__ == '__main__':
    # command line arguments parsing
    args_parser = get_parser()
    args = args_parser.parse_args()
    # convert A-REX loglevel to Python logging string loglevel
    if args.debug == 'FATAL':
        args.debug = 'CRITICAL'
    elif args.debug == 'VERBOSE':
        args.debug = 'DEBUG'
    # set loglevel
    loglevel = getattr(logging, args.debug, 30)
    rootlog.setLevel(loglevel)

    # Init all SSM logging to file
    log_handler_file = logging.FileHandler(args.logfile)
    log_handler_file.setFormatter(
        logging.Formatter('[%(asctime)s] [%(name)s] [%(levelname)s] [%(process)d] [%(message)s]'))
    rootlog.addHandler(log_handler_file)

    # Define broker in accordance to command line options
    brokers = [(args.host, int(args.port))]

    logger.info('Initializing APEL SSM sender to publish records in %s to %s', args.message_path, args.host)
    try:
        # create SSM2 object for sender
        sender = Ssm2(brokers,
                      qpath=args.message_path,
                      cert=args.cert,
                      key=args.key,
                      capath=args.ca_dir,
                      dest=args.topic,
                      use_ssl=args.ssl)

        if sender.has_msgs():
            sender.handle_connect()
            sender.send_all()
            logger.info('SSM records sending to %s has finished.', args.host)
        else:
            logger.info('No messages found in %s to be sent.', args.message_path)

    except (Ssm2Exception, CryptoException) as e:
        logger.error('SSM failed to send messages successfully. Error: %s', e)
    except Exception as e:
        logger.error('Unexpected exception in SSM (type %s). Error: %s. ', e.__class__, e)

    try:
        sender.close_connection()
    except UnboundLocalError:
        # SSM not set up.
        pass

    logger.debug('SSM has shut down.')
