#!/usr/bin/python3

import multiprocessing
import os
import pwd
import sys
import socket
import time
import optparse

# It looks like HTCondor spawns "daemon core" processes with a signal
# mask set, probably because it doesn't want to get signals between the
# fork() (clone()?) and exec().  These two lines of code unblock all
# signals.  This doesn't cause this program to respond as you'd expect
# gunicorn to, but they may be a problem with this script, instead.  (It
# would have to forward signals along to the 'server' process, at the
# least.)
import signal
signal.pthread_sigmask(signal.SIG_UNBLOCK, range(1, signal.NSIG))

os.environ.setdefault('CONDOR_CONFIG', '/etc/condor-ce/condor_config')

import gunicorn.app.base
import gunicorn.glogging

import htcondor2 as htcondor
import htcondorce.web

ALIVE_HEARTBEAT = 60
def send_heartbeat():
    euid = os.geteuid()
    try:
        htcondor.send_alive(timeout=ALIVE_HEARTBEAT)
        htcondor.log(htcondor.LogLevel.Always, 'Sent heartbeat to HTCondor-CE master')
    except (RuntimeError, ValueError) as exc:
        if 'CONDOR_INHERIT' in exc or 'location ClassAd' in exc:
            htcondor.log(htcondor.LogLevel.FullDebug,
                         'WARNING: Could not find location of HTCondor-CE master daemon to send keepalive')
        else:
            htcondor.log(htcondor.LogLevel.Always,
                         'ERROR: Failed to send keepalive to the HTCondor-CE master daemon with EUID {0}'.format(euid))
    if euid != os.geteuid():
        os.seteuid(euid)


def condor_ids():
    if 'CONDOR_IDS' in htcondor.param:
        info = htcondor.param['CONDOR_IDS'].split('.', 1)
        return int(info[0]), int(info[1])
    info = pwd.getpwnam('condor')
    return info.pw_uid, info.pw_gid


def setup_logging():
    if not hasattr(htcondor, 'log'):
        return
    euid = os.geteuid()
    if os.isatty(1):
        htcondor.enable_debug()
    else:
        htcondor.enable_log()
    if euid != os.geteuid():
        os.seteuid(euid)


def parse_opts(localname=None, pidfile=None):
    parser = optparse.OptionParser()
    # Unimplemented, DaemonCore short-name options
    parser.add_option("-a")
    parser.add_option("-b", action="store_true")
    parser.add_option("-c")
    parser.add_option("-d", action="store_true")
    parser.add_option("-f", action="store_true")
    parser.add_option("-k")
    parser.add_option("-l")
    parser.add_option("-q", action="store_true")
    parser.add_option("-r", type="int")
    parser.add_option("-t", action="store_true")
    parser.add_option("-v", action="store_true")
    # Implemented options
    parser.add_option("--pool", help="HTCondor-CE pool to consider.", dest="pool")
    parser.add_option("-n", "--name", help="HTCondor-CE schedd to consider.", dest="name")
    parser.add_option("--spool", help="Spool directory to use.", dest="spool")
    parser.add_option("-p", "--port", help="Port to use for webapp.", dest="port")

    opts = parser.parse_args()[0]

    if hasattr(htcondor, 'set_subsystem'):
        htcondor.set_subsystem(localname or 'CEVIEW')
    setup_logging()

    return opts


class MyWSGIGateway(gunicorn.app.base.BaseApplication):
    """From the Gunicorn documentation example:
    https://docs.gunicorn.org/en/20.0.4/custom.html
    """

    def __init__(self, app, gunicorn_opts=None):
        self.options = gunicorn_opts or {}
        self.application = app
        super().__init__()

    def load_config(self):
        config = {key: value for key, value in self.options.items()
                  if key in self.cfg.settings and value is not None}
        for key, value in config.items():
            self.cfg.set(key.lower(), value)

    def load(self):
        return self.application


def to_bytes(str_or_bytes, encoding="utf-8", errors="backslashreplace"):
    if isinstance(str_or_bytes, str):
        return str_or_bytes.encode(encoding, errors)
    else:
        return str_or_bytes


class WSGILogging(object):

    def __init__(self, app):
        self.app = app

    FORMAT = '{host} "{request}" {status} {size} "{referer}" "{agent}"'

    def __call__(self, environ, start_response):
        status_codes = []
        content_lengths = []
        def custom_start_response(status, response_headers, exc_info=None):
            status_codes.append(int(status.partition(' ')[0]))
            for name, value in response_headers:
                if name.lower() == 'content-length':
                    content_lengths.append(int(value))
                    break
            return start_response(status, response_headers, exc_info)
        retval = self.app(environ, custom_start_response)
        if isinstance(retval, list):
            retval = list(to_bytes(x) for x in retval)
        content_length = content_lengths[0] if content_lengths else len(b"".join(retval))

        msg = { \
            'host': environ.get('REMOTE_ADDR', ''),
            'request': "%s %s %s" % ( \
              environ.get('REQUEST_METHOD', ''),
              environ.get('PATH_INFO', ''),
              environ.get('SERVER_PROTOCOL', '')
            ),
            'size': content_length,
            'status': status_codes[0],
            'referer': environ.get('HTTP_REFERER', ''),
            'agent': environ.get('HTTP_USER_AGENT', ''),
        }
        msg = self.FORMAT.format(**msg)
        htcondor.log(htcondor.LogLevel.Always, msg)
        return retval


def main():
    dc_long_opts = {}
    for long_opt in ['-local-name', '-pidfile']:
        try:
            long_opt_index = sys.argv.index(long_opt)
            dc_long_opts[long_opt.replace('-', '')] = sys.argv.pop(long_opt_index + 1)
            sys.argv.pop(long_opt_index)
        except ValueError:
            pass

    opts = parse_opts(**dc_long_opts)

    spooldir = htcondor.param.get("HTCONDORCE_VIEW_SPOOL")
    if opts.spool:
        spooldir = opts.spool
    if not spooldir:
        if not os.path.exists("tmp"):
            os.mkdir("tmp")
        spooldir = "tmp"

    if opts.port:
        port = int(opts.port)
    else:
        port = int(htcondor.param.get('HTCONDORCE_VIEW_PORT', 8080))

    wsgi_opts = htcondorce.web.GUNICORN_CONFIG.copy()
    wsgi_opts['raw_env'] = [f"htcondorce.spool={spooldir}"]

    if os.path.exists('templates'):
        wsgi_opts['raw_env'].append("htcondorce.templates=templates")
    if opts.pool:
        wsgi_opts['raw_env'].append(f"htcondorce.pool={opts.pool}")
    if opts.name:
        wsgi_opts['raw_env'].append(f"htcondorce.name={opts.name}")

    app = WSGILogging(htcondorce.web.application)

    # Do the IPv6 / v4 dance:
    #   - If bindv6only is enabled, then we must connect to v4 and v6 separately.
    #   - Otherwise, only connect via v6
    bindonly = '/proc/sys/net/ipv6/bindv6only'
    addrs = socket.getaddrinfo(socket.getfqdn(), 9618, socket.AF_UNSPEC, socket.SOCK_STREAM, 0, 0)
    families = [i[0] for i in addrs]
    bind_ipv6_only = False
    wsgi_opts['bind'] = []

    if socket.AF_INET6 in families:
        bind_ipv6_only = os.path.exists(bindonly) and open(bindonly).read() == "1"
        wsgi_opts['bind'].append(f"[::]:{port}")

    if (socket.AF_INET in families) and \
       ((bind_ipv6_only) or (socket.AF_INET6 not in families)):
        wsgi_opts['bind'].append(f"0.0.0.0:{port}")

    if os.geteuid() == 0:
        condor_uid, condor_gid = condor_ids()
        wsgi_opts['user'] = condor_uid
        wsgi_opts['group'] = condor_gid

    gateway = MyWSGIGateway(app, wsgi_opts)
    server = multiprocessing.Process(target=gateway.run)
    server.start()

    if not os.isatty(1):
        while True:
            os.seteuid(condor_uid)
            os.setegid(condor_gid)
            send_heartbeat()
            time.sleep(max(ALIVE_HEARTBEAT/3-1, 1))


if __name__ == '__main__':
    main()
