Merge pull request #78 from hjacobs/screen-tokens

#65 allow creating and redeeming screen tokens

app.py

@@ -14,7 +14,9 @@ import os
 import re
 import requests
 import datetime
+import random
 import redis
+import string
 import time
 import tokens
 from queue import Queue
@@ -24,12 +26,41 @@ from flask import Flask, redirect
 from flask_oauthlib.client import OAuth, OAuthRemoteApp
 from urllib.parse import urljoin
 
+ONE_YEAR = 3600 * 24 * 365
+
 logging.basicConfig(level=logging.INFO)
 
 
+def generate_token(n: int):
+    '''Generate a random ASCII token of length n'''
+    # uses os.urandom()
+    rng = random.SystemRandom()
+    return ''.join([rng.choice(string.ascii_letters + string.digits) for i in range(n)])
+
+
+def generate_token_data():
+    '''Generate screen token data for storing'''
+    token = generate_token(10)
+    now = time.time()
+    return {'token': token, 'created': now, 'expires': now + ONE_YEAR}
+
+
+def check_token(token: str, remote_addr: str, data: dict):
+    '''Check whether the given screen token is valid, raises exception if not'''
+    now = time.time()
+    if data and now < data['expires'] and data.get('remote_addr', remote_addr) == remote_addr:
+        data['remote_addr'] = remote_addr
+        return data
+    else:
+        raise ValueError('Invalid token')
+
+
 class MemoryStore:
+    '''Memory-only backend, mostly useful for local debugging'''
+
     def __init__(self):
         self._queues = []
+        self._screen_tokens = {}
 
     def acquire_lock(self):
         # no-op for memory store
@@ -53,9 +84,22 @@ class MemoryStore:
         finally:
             self._queues.remove(queue)
 
+    def create_screen_token(self):
+        data = generate_token_data()
+        token = data['token']
+        self._screen_tokens[token] = data
+        return token
+
+    def redeem_screen_token(self, token: str, remote_addr: str):
+        data = self._screen_tokens.get(token)
+        data = check_token(token, remote_addr, data)
+        self._screen_tokens[token] = data
+
 
 class RedisStore:
-    def __init__(self, url):
+    '''Redis-based backend for deployments with replicas > 1'''
+
+    def __init__(self, url: str):
         logging.info('Connecting to Redis on {}..'.format(url))
         self._redis = redis.StrictRedis.from_url(url)
         self._redlock = Redlock([url])
@@ -77,6 +121,23 @@ class RedisStore:
                 event_type, data = message['data'].decode('utf-8').split(':', 1)
                 yield (event_type, json.loads(data))
 
+    def create_screen_token(self):
+        '''Generate a new screen token and store it in Redis'''
+        data = generate_token_data()
+        token = data['token']
+        self._redis.set('screen-tokens:{}'.format(token), json.dumps(data))
+        return token
+
+    def redeem_screen_token(self, token: str, remote_addr: str):
+        '''Validate the given token and bind it to the IP'''
+        redis_key = 'screen-tokens:{}'.format(token)
+        data = self._redis.get(redis_key)
+        if not data:
+            raise ValueError('Invalid token')
+        data = json.loads(data.decode('utf-8'))
+        data = check_token(token, remote_addr, data)
+        self._redis.set(redis_key, json.dumps(data))
+
 
 CLUSTER_ID_INVALID_CHARS = re.compile('[^a-z0-9:-]')
 
@@ -335,6 +396,27 @@ def get_events():
     return flask.Response(event(cluster_ids), mimetype='text/event-stream')
 
 
+@app.route('/screen-tokens', methods=['GET', 'POST'])
+@authorize
+def screen_tokens():
+    new_token = None
+    if flask.request.method == 'POST':
+        new_token = STORE.create_screen_token()
+    return flask.render_template('screen-tokens.html', new_token=new_token)
+
+
+@app.route('/screen/<token>')
+def redeem_screen_token(token: str):
+    remote_addr = flask.request.headers.get('X-Forwarded-For') or flask.request.remote_addr
+    logging.info('Trying to redeem screen token "{}" for IP {}..'.format(token, remote_addr))
+    try:
+        STORE.redeem_screen_token(token, remote_addr)
+    except:
+        flask.abort(401)
+    flask.session['auth_token'] = (token, '')
+    return redirect(urljoin(APP_URL, '/'))
+
+
 @app.route('/login')
 def login():
     redirect_uri = urljoin(APP_URL, '/login/authorized')
@@ -361,11 +443,6 @@ def authorized():
     return redirect(urljoin(APP_URL, '/'))
 
 
-@auth.tokengetter
-def get_auth_oauth_token():
-    return flask.session.get('auth_token')
-
-
 def update():
     while True:
         lock = STORE.acquire_lock()

templates/screen-tokens.html

@@ -0,0 +1,18 @@
+<!doctype html>
+<html>
+    <head>
+        <meta charset="utf-8">
+        <meta name="viewport" content="width=device-width, initial-scale=1">
+        <title>Screen Tokens</title>
+        <link rel="shortcut icon" href="static/favicon.ico">
+    </head>
+    <body>
+        <h1>Screen Tokens</h1>
+        {% if new_token: %}
+            <p>The new token is: <code>{{ new_token }}</code></p>
+        {% endif %}
+        <form action="" method="post">
+            <button type="submit">Create new token</button>
+        </form>
+    </body>
+</html>