import json import sqlite3 import time import uuid from flask.sessions import SessionInterface, SessionMixin from werkzeug.datastructures import CallbackDict _LEVEL_INT_TO_STR = {0: 'nothing', 1: 'viewer', 2: 'administrator', 3: 'manager'} class SqliteSession(CallbackDict, SessionMixin): def __init__(self, initial=None, sid=None, new=False): def on_update(self): self.modified = True CallbackDict.__init__(self, initial or {}, on_update) self.sid = sid self.new = new self.modified = False class SqliteSessionInterface(SessionInterface): def __init__(self, db_path): self.db_path = db_path def _connect(self): con = sqlite3.connect(self.db_path, timeout=5) con.execute('PRAGMA journal_mode=WAL') con.row_factory = sqlite3.Row return con def open_session(self, app, request): name = app.config.get('SESSION_COOKIE_NAME', 'session') sid = request.cookies.get(name) if not sid: return SqliteSession(sid=str(uuid.uuid4()), new=True) try: con = self._connect() row = con.execute( '''SELECT s.session_id, s.account_id, s.tz_offset_seconds, s.preferences_json, s.flashes_json, a.email, a.access_level FROM sessions s JOIN accounts a ON a.account_id = s.account_id WHERE s.session_id=?''', (sid,) ).fetchone() if row: prefs = json.loads(row['preferences_json'] or '{}') flashes = json.loads(row['flashes_json'] or '[]') data = { 'account_id': str(row['account_id']), 'email_address': str(row['email']), 'access_level': _LEVEL_INT_TO_STR.get(row['access_level'], 'viewer'), 'tz_offset_seconds': int(row['tz_offset_seconds']), 'timezone': str(prefs.get('timezone', '')), 'apply_changes_immediately': bool(prefs.get('apply_changes_immediately', False)), '_flashes': flashes, } con.close() return SqliteSession(data, sid=sid, new=False) client = con.execute( 'SELECT flashes_json FROM clients WHERE cookie_unique_token=?', (sid,) ).fetchone() con.close() flashes = json.loads(client['flashes_json'] or '[]') if client else [] data = {'_flashes': flashes} if flashes else {} return SqliteSession(data, sid=sid, new=False) except Exception: pass return SqliteSession(sid=sid, new=False) def save_session(self, app, session, response): name = app.config.get('SESSION_COOKIE_NAME', 'session') domain = self.get_cookie_domain(app) path = self.get_cookie_path(app) if not session and session.modified and not session.new: try: con = self._connect() con.execute('DELETE FROM sessions WHERE session_id=?', (session.sid,)) con.commit() con.close() except Exception: pass response.delete_cookie(name, domain=domain, path=path) return account_id = session.get('account_id') flashes_json = json.dumps(session.get('_flashes', [])) now = int(time.time()) try: con = self._connect() if account_id: prefs = json.dumps({ 'timezone': session.get('timezone', ''), 'apply_changes_immediately': bool(session.get('apply_changes_immediately', False)), }) tz_offset = int(session.get('tz_offset_seconds', 0)) con.execute('INSERT OR IGNORE INTO clients (cookie_unique_token) VALUES (?)', (session.sid,)) con.execute( '''INSERT INTO sessions (session_id, account_id, tz_offset_seconds, preferences_json, flashes_json, session_started_ts, last_seen_ts) VALUES (?,?,?,?,?,?,?) ON CONFLICT(session_id) DO UPDATE SET account_id=excluded.account_id, tz_offset_seconds=excluded.tz_offset_seconds, preferences_json=excluded.preferences_json, flashes_json=excluded.flashes_json, last_seen_ts=excluded.last_seen_ts''', (session.sid, account_id, tz_offset, prefs, flashes_json, now, now) ) else: con.execute( '''INSERT INTO clients (cookie_unique_token, flashes_json) VALUES (?,?) ON CONFLICT(cookie_unique_token) DO UPDATE SET flashes_json=excluded.flashes_json''', (session.sid, flashes_json) ) con.commit() con.close() except Exception: pass response.set_cookie( name, session.sid, expires=self.get_expiration_time(app, session), httponly=self.get_cookie_httponly(app), domain=domain, path=path, secure=self.get_cookie_secure(app), samesite=self.get_cookie_samesite(app), )