split multiuser files

multiuser/user.py

@@ -0,0 +1,180 @@
+"""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()