"""script to monitor and cull idle single-user servers

Adapted to SWAN from https://github.com/jupyterhub/jupyterhub/blob/1.1.0/examples/cull-idle/cull_idle_servers.py

Caveats:

last_activity is not updated with high frequency,
so cull timeout should be greater than the sum of:

- single-user websocket ping interval (default: 30s)
- JupyterHub.last_activity_interval (default: 5 minutes)

You can run this as a service managed by JupyterHub with this in your config::


    c.JupyterHub.services = [
        {
            'name': 'cull-idle',
            'admin': True,
            'command': [sys.executable, 'cull_idle_servers.py', '--timeout=3600'],
        }
    ]

Or run it manually by generating an API token and storing it in `JUPYTERHUB_API_TOKEN`:

    export JUPYTERHUB_API_TOKEN=$(jupyterhub token)
    python3 cull_idle_servers.py [--timeout=900] [--url=http://127.0.0.1:8081/hub/api]

This script uses the same ``--timeout`` and ``--max-age`` values for
culling users and users' servers.  If you want a different value for
users and servers, you should add this script to the services list
twice, just with different ``name``s, different values, and one with
the ``--cull-users`` option.
"""
import json
import os
from datetime import datetime, timezone
from functools import partial

try:
    from urllib.parse import quote, urlencode
except ImportError:
    from urllib import quote, urlencode

from subprocess import call

import dateutil.parser
from tornado.gen import coroutine, multi
from tornado.httpclient import AsyncHTTPClient, HTTPClientError, HTTPRequest
from tornado.ioloop import IOLoop, PeriodicCallback
from tornado.locks import Semaphore
from tornado.log import app_log
from tornado.options import define, options, parse_command_line

# Global variables
# The users list is shared between the functions that check for idle
# sessions and the one that validates wether a user is blocked
users = []

def check_ticket(username):
    app_log.info("Checking ticket for user %s", username)
    call(['sudo', '--preserve-env=SWAN_DEV', "%s/check_ticket.sh" % options.hooks_dir, username])
# End SWAN code

def parse_date(date_string):
    """Parse a timestamp

    If it doesn't have a timezone, assume utc

    Returned datetime object will always be timezone-aware
    """
    dt = dateutil.parser.parse(date_string)
    if not dt.tzinfo:
        # assume naïve timestamps are UTC
        dt = dt.replace(tzinfo=timezone.utc)
    return dt


def format_td(td):
    """
    Nicely format a timedelta object

    as HH:MM:SS
    """
    if td is None:
        return "unknown"
    if isinstance(td, str):
        return td
    seconds = int(td.total_seconds())
    h = seconds // 3600
    seconds = seconds % 3600
    m = seconds // 60
    seconds = seconds % 60
    return f"{h:02}:{m:02}:{seconds:02}"

@coroutine
def delete_server(user_name, server_name, url, auth_header, fetch):
    if server_name:
        # culling a named server
        delete_url = url + "/users/%s/servers/%s" % (
            quote(user_name),
            quote(server_name),
        )
    else:
        delete_url = url + '/users/%s/server' % quote(user_name)

    req = HTTPRequest(url=delete_url, method='DELETE', headers=auth_header)
    resp = yield fetch(req)
    return resp

@coroutine
def delete_user(user_name, url, auth_header, fetch):
    req = HTTPRequest(
        url=url + '/users/%s' % user_name, method='DELETE', headers=auth_header
    )
    resp = yield fetch(req)
    return resp

@coroutine
def cull_idle(
    url, api_token, inactive_limit, cull_users=False, disable_hooks=False, max_age=0, concurrency=10
):
    """Shutdown idle single-user servers

    If cull_users, inactive *users* will be deleted as well.
    """
    auth_header = {'Authorization': 'token %s' % api_token}
    req = HTTPRequest(url=url + '/users', headers=auth_header)
    now = datetime.now(timezone.utc)
    client = AsyncHTTPClient()

    if concurrency:
        semaphore = Semaphore(concurrency)

        @coroutine
        def fetch(req):
            """client.fetch wrapped in a semaphore to limit concurrency"""
            yield semaphore.acquire()
            try:
                return (yield client.fetch(req))
            finally:
                yield semaphore.release()

    else:
        fetch = client.fetch

    resp = yield fetch(req)

    global users
    users = json.loads(resp.body.decode('utf8', 'replace'))
    futures = []

    @coroutine
    def handle_server(user, server_name, server, max_age, inactive_limit):
        """Handle (maybe) culling a single server

        "server" is the entire server model from the API.

        Returns True if server is now stopped (user removable),
        False otherwise.
        """
        log_name = user['name']
        if server_name:
            log_name = '%s/%s' % (user['name'], server_name)
        if server.get('pending'):
            app_log.warning(
                "Not culling server %s with pending %s", log_name, server['pending']
            )
            return False

        # jupyterhub < 0.9 defined 'server.url' once the server was ready
        # as an *implicit* signal that the server was ready.
        # 0.9 adds a dedicated, explicit 'ready' field.
        # By current (0.9) definitions, servers that have no pending
        # events and are not ready shouldn't be in the model,
        # but let's check just to be safe.

        if not server.get('ready', bool(server['url'])):
            app_log.warning(
                "Not culling not-ready not-pending server %s: %s", log_name, server
            )
            return False

        if server.get('started'):
            age = now - parse_date(server['started'])
        else:
            # started may be undefined on jupyterhub < 0.9
            age = None

        # check last activity
        # last_activity can be None in 0.9
        if server['last_activity']:
            inactive = now - parse_date(server['last_activity'])
        else:
            # no activity yet, use start date
            # last_activity may be None with jupyterhub 0.9,
            # which introduces the 'started' field which is never None
            # for running servers
            inactive = age

        # CUSTOM CULLING TEST CODE HERE
        # Add in additional server tests here.  Return False to mean "don't
        # cull", True means "cull immediately", or, for example, update some
        # other variables like inactive_limit.
        #
        # Here, server['state'] is the result of the get_state method
        # on the spawner.  This does *not* contain the below by
        # default, you may have to modify your spawner to make this
        # work.  The `user` variable is the user model from the API.
        #
        # if server['state']['profile_name'] == 'unlimited'
        #     return False
        # inactive_limit = server['state']['culltime']

        should_cull = (
            inactive is not None and inactive.total_seconds() >= inactive_limit
        )
        if should_cull:
            app_log.info(
                "Culling server %s (inactive for %s)", log_name, format_td(inactive)
            )

        if max_age and not should_cull:
            # only check started if max_age is specified
            # so that we can still be compatible with jupyterhub 0.8
            # which doesn't define the 'started' field
            if age is not None and age.total_seconds() >= max_age:
                app_log.info(
                    "Culling server %s (age: %s, inactive for %s)",
                    log_name,
                    format_td(age),
                    format_td(inactive),
                )
                should_cull = True

        if not should_cull:
            app_log.debug(
                "Not culling server %s (age: %s, inactive for %s)",
                log_name,
                format_td(age),
                format_td(inactive),
            )
            return False

        resp = yield delete_server(user['name'], server_name, url, auth_header, fetch)
        if resp.code == 202:
            app_log.warning("Server %s is slow to stop", log_name)
            # return False to prevent culling user with pending shutdowns
            return False
        return True

    @coroutine
    def handle_user(user):
        """Handle one user.

        Create a list of their servers, and async exec them.  Wait for
        that to be done, and if all servers are stopped, possibly cull
        the user.
        """
        # shutdown servers first.
        # Hub doesn't allow deleting users with running servers.
        # jupyterhub 0.9 always provides a 'servers' model.
        # 0.8 only does this when named servers are enabled.
        if 'servers' in user:
            servers = user['servers']
        else:
            # jupyterhub < 0.9 without named servers enabled.
            # create servers dict with one entry for the default server
            # from the user model.
            # only if the server is running.
            servers = {}
            if user['server']:
                servers[''] = {
                    'last_activity': user['last_activity'],
                    'pending': user['pending'],
                    'url': user['server'],
                }
        server_futures = [
            handle_server(user, server_name, server, max_age, inactive_limit)
            for server_name, server in servers.items()
        ]
        results = yield multi(server_futures)

        # some servers are still running, cannot cull users
        still_alive = len(results) - sum(results)

        # we need to know if we want to renew the EOS ticket or not
        # therefore this part of the code needs to check if there are sessions alive
        if not cull_users:
            return not still_alive

        if still_alive:
            app_log.debug(
                "Not culling user %s with %i servers still alive",
                user['name'],
                still_alive,
            )
            return False

        should_cull = False
        if user.get('created'):
            age = now - parse_date(user['created'])
        else:
            # created may be undefined on jupyterhub < 0.9
            age = None

        # check last activity
        # last_activity can be None in 0.9
        if user['last_activity']:
            inactive = now - parse_date(user['last_activity'])
        else:
            # no activity yet, use start date
            # last_activity may be None with jupyterhub 0.9,
            # which introduces the 'created' field which is never None
            inactive = age

        should_cull = (
            inactive is not None and inactive.total_seconds() >= inactive_limit
        )
        if should_cull:
            app_log.info("Culling user %s (inactive for %s)", user['name'], inactive)

        if max_age and not should_cull:
            # only check created if max_age is specified
            # so that we can still be compatible with jupyterhub 0.8
            # which doesn't define the 'started' field
            if age is not None and age.total_seconds() >= max_age:
                app_log.info(
                    "Culling user %s (age: %s, inactive for %s)",
                    user['name'],
                    format_td(age),
                    format_td(inactive),
                )
                should_cull = True

        if not should_cull:
            app_log.debug(
                "Not culling user %s (created: %s, last active: %s)",
                user['name'],
                format_td(age),
                format_td(inactive),
            )
            return False

        yield delete_user(user['name'], url, auth_header, fetch)
        return True

    for user in users:
        futures.append((user['name'], handle_user(user)))

    for (name, f) in futures:
        try:
            result = yield f
        except Exception:
            app_log.exception("Error processing %s", name)
        else:
            if result:
                app_log.debug("Finished culling %s", name)
            else:
                if not disable_hooks: check_ticket(name)

@coroutine
def check_blocked_users(url, api_token, client_id, client_secret, auth_url, audience, authz_api_url):
    """Detect blocked users.

    Detect blocked users and cull their servers.
    """
    # Step 1: Get Token for Authorization Service APIs (no Token Exchange permissions - too powerful)
    token_req = HTTPRequest(
        url=auth_url,
        method="POST",
        headers={'Content-Type': 'application/x-www-form-urlencoded'},
        body="grant_type=client_credentials&client_id=%s&client_secret=%s&audience=%s"
            % (quote(client_id), quote(client_secret), quote(audience)),
    )

    client = AsyncHTTPClient()

    try:
        token_resp = yield client.fetch(token_req)
        token_data = json.loads(token_resp.body.decode("utf-8"))
        access_token = token_data["access_token"]
    except Exception:
        app_log.exception("Failed to get access token for blocked user check")
        return

    # Step 2: Check each user - Get their identity from the Authorization Service API using the obtained token as Bearer
    for user in users:
        app_log.info("Checking if user %s is blocked", user['name'])
        query = urlencode([("field", "uniqueIdentifier"), ("field", "blocked"), ("field", "disabled")])
        id_url = f"{authz_api_url}/{quote(user['name'])}/accounts?{query}"
        id_req = HTTPRequest(
            url=id_url,
            headers={
                "Authorization": "Bearer %s" % access_token,
                "Accept": "*/*",
            }
        )

        try:
            id_resp = yield client.fetch(id_req)
            data = json.loads(id_resp.body.decode("utf-8")).get("data", [])
            if not data:
                app_log.warning(f"No identity data returned for user {user['name']}")
                continue
            is_blocked = data[0].get("blocked", False)
            is_disabled = data[0].get("disabled", False)
        except HTTPClientError as e:
            if e.code == 404:
                app_log.info(f"User {user['name']} not found for blocked user check, skipping.")
                continue
            else:
                app_log.warning(f"Failed to check identity for {user['name']}: {e}")
                continue

        if is_blocked or is_disabled:
            app_log.warning("User %s is blocked. Terminating their sessions.", user['name'])
            auth_header = {'Authorization': 'token %s' % api_token}

            # collect user's server
            if 'servers' in user:
                servers = user['servers']
            else:
                servers = {}
                if user.get('server'):
                    servers[''] = {
                        'last_activity': user.get('last_activity'),
                        'pending': user.get('pending'),
                        'url': user['server'],
                    }

            delete_futures = []
            for server_name in servers:
                delete_futures.append(
                    delete_server(user['name'], server_name, url, auth_header, client.fetch)
                )
            # wait for all delete requests to complete
            results = yield multi(delete_futures)

            for i, server_name in enumerate(servers):
                try:
                    resp = results[i]
                    if resp.code in (204, 202):
                        app_log.info("Deleted server '%s' for user %s", server_name, user['name'])
                    else:
                        app_log.warning("Unexpected response deleting server %s for user %s: %s",
                                        server_name, user['name'], resp.code)
                except Exception:
                    app_log.exception("Failed to delete server %s for user %s", server_name, user['name'])

            yield delete_user(user['name'], url, auth_header, client.fetch)

def main():
    define(
        'url',
        default=os.environ.get('JUPYTERHUB_API_URL'),
        help="The JupyterHub API URL",
    )
    define('timeout', default=600, help="The idle timeout (in seconds)")
    define(
        'cull_every',
        default=0,
        help="The interval (in seconds) for checking for idle servers to cull",
    )
    define(
        'max_age',
        default=0,
        help="The maximum age (in seconds) of servers that should be culled even if they are active",
    )
    define(
        'cull_users',
        default=False,
        help="""Cull users in addition to servers.
                This is for use in temporary-user cases such as tmpnb.""",
    )
    define(
        'concurrency',
        default=10,
        help="""Limit the number of concurrent requests made to the Hub.

                Deleting a lot of users at the same time can slow down the Hub,
                so limit the number of API requests we have outstanding at any given time.
                """,
    )
    define('hooks_dir', default="/srv/jupyterhub/culler", help="Path to the directory for the krb tickets script (check_ticket.sh)")
    define('disable_hooks', default=False, help="The user's home is a temporary scratch directory and we should not check krb tickets")
    define('auth_url', default='', help="URL to fetch CERN access token")
    define('auth_client_id', default=os.environ.get('AUTH_CLIENT_ID'), help="Client ID for blocked user check")
    define('auth_client_secret', default=os.environ.get('AUTH_CLIENT_SECRET'), help="Client secret for blocked user check")
    define('audience', default='', help="Audience for CERN access token")
    define('auth_check_interval', default=0, help="The interval (in seconds) for checking blocked users")
    define('authz_api_url', default='', help="URL to fetch user identity from authorization service")


    parse_command_line()
    if not options.cull_every:
        options.cull_every = options.timeout // 2
    api_token = os.environ['JUPYTERHUB_API_TOKEN']

    try:
        AsyncHTTPClient.configure("tornado.curl_httpclient.CurlAsyncHTTPClient")
    except ImportError as e:
        app_log.warning(
            "Could not load pycurl: %s\n"
            "pycurl is recommended if you have a large number of users.",
            e,
        )

    loop = IOLoop.current()
    cull = partial(
        cull_idle,
        url=options.url,
        api_token=api_token,
        inactive_limit=options.timeout,
        cull_users=options.cull_users,
        disable_hooks=options.disable_hooks,
        max_age=options.max_age,
        concurrency=options.concurrency,
    )
    # schedule first cull immediately
    # because PeriodicCallback doesn't start until the end of the first interval
    loop.add_callback(cull)
    # schedule periodic cull
    pc = PeriodicCallback(cull, 1e3 * options.cull_every)
    pc.start()

    blocked_check = partial(
        check_blocked_users,
        url=options.url,
        api_token=api_token,
        auth_url = options.auth_url,
        client_id=options.auth_client_id,
        client_secret=options.auth_client_secret,
        audience = options.audience,
        authz_api_url=options.authz_api_url,
    )
    # schedule first blocked user check
    loop.add_callback(blocked_check)
    # schedule periodic blocked user check
    pc_blocked = PeriodicCallback(blocked_check, 1e3 * options.auth_check_interval)
    pc_blocked.start()

    try:
        loop.start()
    except KeyboardInterrupt:
        pass
