jupyterhub/multiuser/user.py

"""Classes for managing users."""

import json
import os
import signal
import sys
import time
import uuid
from subprocess import Popen

import requests

from tornado.log import app_log

from IPython.utils.traitlets import Any, Unicode, Integer, Dict
from IPython.config import LoggingConfigurable

from .utils import random_port, wait_for_server

class UserSession(LoggingConfigurable):
    env_prefix = Unicode('IPY_')
    process = Any()
    port = Integer()
    user = Unicode()
    cookie_secret = Unicode()
    cookie_name = Unicode()
    def _cookie_name_default(self):
        return 'ipy-multiuser-%s' % self.user
    
    multiuser_prefix = Unicode()
    multiuser_api_url = Unicode()
    
    url_prefix = Unicode()
    def _url_prefix_default(self):
        return '/user/%s/' % self.user
    
    api_token = Unicode()
    def _api_token_default(self):
        return str(uuid.uuid4())
    
    cookie_token = Unicode()
    def _cookie_token_default(self):
        return str(uuid.uuid4())
    
    def _env_key(self, d, key, value):
        d['%s%s' % (self.env_prefix, key)] = value
    
    env = Dict()
    def _env_default(self):
        env = os.environ.copy()
        self._env_key(env, 'COOKIE_SECRET', self.cookie_secret)
        self._env_key(env, 'API_TOKEN', self.api_token)
        return env
    
    @property
    def auth_data(self):
        return dict(
            user=self.user,
        )
    
    def start(self):
        assert self.process is None or self.process.poll() is not None
        cmd = [sys.executable, '-m', 'multiuser.singleuser',
            '--user=%s' % self.user, '--port=%i' % self.port,
            '--cookie-name=%s' % self.cookie_name,
            '--multiuser-prefix=%s' % self.multiuser_prefix,
            '--multiuser-api-url=%s' % self.multiuser_api_url,
            '--base-url=%s' % self.url_prefix,
            ]
        app_log.info("Spawning: %s" % cmd)
        self.process = Popen(cmd, env=self.env,
            # don't forward signals:
            preexec_fn=os.setpgrp,
        )
    
    def running(self):
        if self.process is None:
            return False
        if self.process.poll() is not None:
            self.process = None
            return False
        return True
    
    def request_stop(self):
        if self.running():
            self.process.send_signal(signal.SIGINT)
            time.sleep(0.1)
        if self.running():
            self.process.send_signal(signal.SIGINT)
        
    def stop(self):
        for i in range(100):
            if self.running():
                time.sleep(0.1)
            else:
                break
        if self.running():
            self.process.terminate()


class UserManager(LoggingConfigurable):
    
    users = Dict()
    routes_t = Unicode('http://{ip}:{port}/api/routes{uri}')
    single_user_t = Unicode('http://{ip}:{port}')
    
    single_user_ip = Unicode('localhost')
    proxy_ip = Unicode('localhost')
    proxy_port = Integer(8001)
    
    proxy_auth_token = Unicode()
    def _proxy_auth_token_default(self):
        return str(uuid.uuid4())
    
    def get_session(self, user, **kwargs):
        if user not in self.users:
            kwargs['user'] = user
            self.users[user] = UserSession(**kwargs)
        return self.users[user]
            
    def spawn(self, user):
        session = self.get_session(user)
        if session.running():
            app_log.warn("User session %s already running", user)
            return
        session.port = port = random_port()
        session.start()
        
        r = requests.post(
            self.routes_t.format(
                ip=self.proxy_ip,
                port=self.proxy_port,
                uri=session.url_prefix,
            ),
            data=json.dumps(dict(
                target=self.single_user_t.format(
                    ip=self.single_user_ip,
                    port=port
                ),
                user=user,
            )),
            headers={'Authorization': "token %s" % self.proxy_auth_token},
        )
        wait_for_server(self.single_user_ip, port)
        r.raise_for_status()
    
    def user_for_api_token(self, token):
        """Get the user session object for a given API token"""
        for session in self.users.values():
            if session.api_token == token:
                return session
    
    def user_for_cookie_token(self, token):
        """Get the user session object for a given cookie token"""
        for session in self.users.values():
            if session.cookie_token == token:
                return session
    
    def shutdown(self, user):
        assert user in self.users
        session = self.users.pop(user)
        session.stop()
        r = requests.delete(self.routes_url,
            data=json.dumps(user=user, port=session.port),
        )
        r.raise_for_status()
    
    def cleanup(self):
        sessions = list(self.users.values())
        self.users = {}
        for session in sessions:
            self.log.info("Cleaning up %s's server" % session.user)
            session.request_stop()
        for i in range(100):
            if any([ session.running() for session in sessions ]):
                time.sleep(0.1)
            else:
                break
        for session in sessions:
            session.stop()