From 8d71edae5e017c4ccf4fd8299e98650210290c34 Mon Sep 17 00:00:00 2001
From: Gergely Polonkai <gergely@polonkai.eu>
Date: Wed, 18 Jul 2018 08:26:39 +0200
Subject: [PATCH] Save sessions in the cache

---
 calsocial/__init__.py |   4 +-
 calsocial/cache.py    | 131 ++++++++++++++++++++++++++++++++++++++++++
 calsocial/security.py |   4 +-
 3 files changed, 137 insertions(+), 2 deletions(-)

diff --git a/calsocial/__init__.py b/calsocial/__init__.py
index 59480ac..7705892 100644
--- a/calsocial/__init__.py
+++ b/calsocial/__init__.py
@@ -26,7 +26,7 @@ from flask_security import SQLAlchemyUserDatastore, current_user, login_required
 from sqlalchemy.orm.exc import NoResultFound, MultipleResultsFound
 
 from calsocial.account import AccountBlueprint
-from calsocial.cache import cache
+from calsocial.cache import CachedSessionInterface, cache
 from calsocial.utils import RoutedMixin
 
 
@@ -68,6 +68,8 @@ class CalendarSocialApp(Flask, RoutedMixin):
 
         Flask.__init__(self, name)
 
+        self.session_interface = CachedSessionInterface()
+
         self._timezone = None
 
         config_name = os.environ.get('ENV', config or 'dev')
diff --git a/calsocial/cache.py b/calsocial/cache.py
index e4e923c..5628662 100644
--- a/calsocial/cache.py
+++ b/calsocial/cache.py
@@ -17,6 +17,137 @@
 """Caching functionality for Calendar.social
 """
 
+from datetime import timedelta
+import pickle
+from uuid import uuid4
+
+from flask import current_app, has_request_context, request, session
+from flask.sessions import SessionInterface, SessionMixin
 from flask_caching import Cache
+from werkzeug.datastructures import CallbackDict
 
 cache = Cache()  # pylint: disable=invalid-name
+
+
+class CachedSession(CallbackDict, SessionMixin):  # pylint: disable=too-many-ancestors
+    """Object for session data saved in the cache
+    """
+
+    def __init__(self, initial=None, sid=None, new=False):
+        self.__modifying = False
+
+        def on_update(self):
+            """Function to call when session data is updated
+            """
+
+            if self.__modifying:
+                return
+
+            self.__modifying = True
+
+            if has_request_context():
+                self['ip'] = request.remote_addr
+
+            self.modified = True
+
+            self.__modifying = False
+
+        CallbackDict.__init__(self, initial, on_update)
+        self.sid = sid
+        self.new = new
+        self.modified = False
+
+    @property
+    def user(self):
+        from calsocial.models import User
+
+        if 'user_id' not in self:
+            return None
+
+        return User.query.get(self['user_id'])
+
+
+class CachedSessionInterface(SessionInterface):
+    """A session interface that loads/saves session data from the cache
+    """
+
+    serializer = pickle
+    session_class = CachedSession
+    global_cache = cache
+
+    def __init__(self, prefix='session:'):
+        self.cache = cache
+        self.prefix = prefix
+
+    @staticmethod
+    def generate_sid():
+        """Generade a new session ID
+        """
+
+        return str(uuid4())
+
+    @staticmethod
+    def get_cache_expiration_time(app, session):
+        """Get the expiration time of the cache entry
+        """
+
+        if session.permanent:
+            return app.permanent_session_lifetime
+
+        return timedelta(days=1)
+
+    def open_session(self, app, request):
+        sid = request.cookies.get(app.session_cookie_name)
+
+        if not sid:
+            sid = self.generate_sid()
+
+            return self.session_class(sid=sid, new=True)
+
+        session = self.load_session(sid)
+
+        if session is None:
+            return self.session_class(sid=sid, new=True)
+
+        return session
+
+    def load_session(self, sid):
+        """Load a specific session from the cache
+        """
+
+        val = self.cache.get(self.prefix + sid)
+
+        if val is None:
+            return None
+
+        data = self.serializer.loads(val)
+
+        return self.session_class(data, sid=sid)
+
+    def save_session(self, app, session, response):
+        domain = self.get_cookie_domain(app)
+
+        if not session:
+            self.cache.delete(self.prefix + session.sid)
+
+            if session.modified:
+                response.delete_cookie(app.session_cookie_name, domain=domain)
+
+            return
+
+        cache_exp = self.get_cache_expiration_time(app, session)
+        cookie_exp = self.get_expiration_time(app, session)
+        val = self.serializer.dumps(dict(session))
+        self.cache.set(self.prefix + session.sid, val, int(cache_exp.total_seconds()))
+
+        response.set_cookie(app.session_cookie_name,
+                            session.sid,
+                            expires=cookie_exp,
+                            httponly=True,
+                            domain=domain)
+
+    def delete_session(self, sid):
+        if has_request_context() and session.sid == sid:
+            raise ValueError('Will not delete the current session')
+
+        cache.delete(self.prefix + sid)
diff --git a/calsocial/security.py b/calsocial/security.py
index 0eab762..99ab50b 100644
--- a/calsocial/security.py
+++ b/calsocial/security.py
@@ -17,7 +17,7 @@
 """Security related things for Calendar.social
 """
 
-from flask import current_app
+from flask import current_app, session
 from flask_login.signals import user_logged_in, user_logged_out
 from flask_security import Security, AnonymousUser as BaseAnonymousUser
 
@@ -45,6 +45,8 @@ def login_handler(app, user):  # pylint: disable=unused-argument
 
     AuditLog.log(user, AuditLog.TYPE_LOGIN_SUCCESS)
 
+    user.active_sessions += [session.sid]
+
 
 @user_logged_out.connect
 def logout_handler(app, user):  # pylint: disable=unused-argument