linuxrouter/docker/routlin-dash/app/session_interface.py

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']),
                    '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({'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),
        )