import flask; from flask import request import functools import werkzeug.wrappers import werkzeug.datastructures from voussoirkit import cacheclass from voussoirkit import passwordy import etiquette SESSION_MAX_AGE = 86400 REQUEST_TYPES = (flask.Request, werkzeug.wrappers.Request, werkzeug.local.LocalProxy) RESPONSE_TYPES = (flask.Response, werkzeug.wrappers.Response) def _generate_token(length=32): return passwordy.random_hex(length=length) def _normalize_token(token): if isinstance(token, REQUEST_TYPES): request = token token = request.cookies.get('etiquette_session', None) if token is None: # During normal usage, this does not occur because give_token is # applied *before* the request handler even sees the request. # Just a precaution. 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 class SessionManager: def __init__(self, maxlen=None): self.sessions = cacheclass.Cache(maxlen=maxlen) def add(self, session): self.sessions[session.token] = session def get(self, request): token = _normalize_token(request) session = self.sessions[token] invalid = ( request.remote_addr != session.ip_address or session.expired() ) if invalid: self.remove(token) raise KeyError(token) return session def give_token(self, function): ''' This decorator ensures that the user has an `etiquette_session` cookie before reaching the request handler. If the user does not have the cookie, they are given one. If they do, its lifespan is reset. ''' @functools.wraps(function) def wrapped(*args, **kwargs): # Inject new token so the function doesn't know the difference token = request.cookies.get('etiquette_session', None) if not token or token not in self.sessions: token = _generate_token() # cookies is currently an ImmutableMultiDict, but in order to # trick the wrapped function I'm gonna have to mutate it. # It is important to use a werkzeug MultiDict and not a plain # Python dict, because werkzeug puts cookies into lists like # {name: [value]} and then cookies.get pulls the first item out # of that list. A plain dict wouldn't have this .get behavior. request.cookies = werkzeug.datastructures.MultiDict(request.cookies) request.cookies['etiquette_session'] = token try: session = self.get(request) except KeyError: session = Session(request, user=None) self.add(session) else: session.maintain() response = function(*args, **kwargs) if not isinstance(response, RESPONSE_TYPES): response = flask.Response(response) # Send the token back to the client # but only if the endpoint didn't manually set the cookie. function_cookies = response.headers.get_all('Set-Cookie') if not any('etiquette_session=' in cookie for cookie in function_cookies): response.set_cookie( 'etiquette_session', value=session.token, max_age=SESSION_MAX_AGE, httponly=True, ) return response return wrapped def remove(self, token): token = _normalize_token(token) try: self.sessions.remove(token) except KeyError: pass class Session: def __init__(self, request, user): self.token = _normalize_token(request) self.user = user self.ip_address = request.remote_addr self.user_agent = request.headers.get('User-Agent', '') self.last_activity = etiquette.helpers.now() def __repr__(self): if self.user: return f'Session {self.token} for user {self.user}' else: return f'Session {self.token} for anonymous' def expired(self): now = etiquette.helpers.now() age = now - self.last_activity return age > SESSION_MAX_AGE def maintain(self): self.last_activity = etiquette.helpers.now()