80 lines
2.5 KiB
Python
80 lines
2.5 KiB
Python
|
import flask
|
||
|
|
from flask import request
|
||
|
|
import functools
|
||
|
|
import helpers
|
||
|
|
import uuid
|
||
|
|
|
||
|
|
def _generate_token():
|
||
|
|
token = str(uuid.uuid4())
|
||
|
|
#print('MAKE SESSION', token)
|
||
|
|
return token
|
||
|
|
|
||
|
|
def _normalize_token(token):
|
||
|
|
if isinstance(token, flask.Request):
|
||
|
|
token = token.cookies.get('etiquette_session', None)
|
||
|
|
|
||
|
|
class SessionManager:
|
||
|
|
def __init__(self):
|
||
|
|
self.sessions = {}
|
||
|
|
|
||
|
|
def add(self, session):
|
||
|
|
self.sessions[session.token] = session
|
||
|
|
|
||
|
|
def get(self, token):
|
||
|
|
token = _normalize_token(token)
|
||
|
|
return self.sessions.get(token, None)
|
||
|
|
|
||
|
|
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:
|
||
|
|
token = _generate_token()
|
||
|
|
request.cookies = dict(request.cookies)
|
||
|
|
request.cookies['etiquette_session'] = token
|
||
|
|
|
||
|
|
response = function(*args, **kwargs)
|
||
|
|
|
||
|
|
if not isinstance(response, flask.Response):
|
||
|
|
response = flask.Response(response)
|
||
|
|
|
||
|
|
# Send the token back to the client
|
||
|
|
# but only if the endpoint didn't manually set the cookie.
|
||
|
|
for (headerkey, value) in response.headers:
|
||
|
|
if headerkey == 'Set-Cookie' and value.startswith('etiquette_session='):
|
||
|
|
break
|
||
|
|
else:
|
||
|
|
response.set_cookie('etiquette_session', value=token, max_age=86400)
|
||
|
|
self.maintain(token)
|
||
|
|
|
||
|
|
return response
|
||
|
|
return wrapped
|
||
|
|
|
||
|
|
def maintain(self, token):
|
||
|
|
session = self.get(token)
|
||
|
|
if session:
|
||
|
|
session.maintain()
|
||
|
|
|
||
|
|
def remove(self, token):
|
||
|
|
token = _normalize_token(token)
|
||
|
|
if token in self.sessions:
|
||
|
|
self.sessions.pop(token)
|
||
|
|
|
||
|
|
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 = int(helpers.now())
|
||
|
|
|
||
|
|
def maintain(self):
|
||
|
|
self.last_activity = int(helpers.now())
|