From 8a1f2c65f9b83f46bd1f87bec2c1886e32a128c9 Mon Sep 17 00:00:00 2001 From: Ethan Dalool Date: Mon, 15 Jan 2018 18:39:40 -0800 Subject: [PATCH] Tighten the types for _normalize_token; include werkzeug wrappers. --- .../etiquette_flask/etiquette_flask/sessions.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/frontends/etiquette_flask/etiquette_flask/sessions.py b/frontends/etiquette_flask/etiquette_flask/sessions.py index 46e79ae..0ae5c3c 100644 --- a/frontends/etiquette_flask/etiquette_flask/sessions.py +++ b/frontends/etiquette_flask/etiquette_flask/sessions.py @@ -13,8 +13,16 @@ def _generate_token(length=32): return token def _normalize_token(token): - if isinstance(token, flask.Request): - token = token.cookies.get('etiquette_session', None) + if isinstance(token, (flask.Request, werkzeug.wrappers.Request, werkzeug.local.LocalProxy)): + request = token + token = request.cookies.get('etiquette_session', None) + if token is None: + message = 'Cannot normalize token for request with no etiquette_session header.' + raise TypeError(message, request) + elif isinstance(token, str): + pass + else: + raise TypeError('Unsupported token normalization', type(token)) return token