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

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 sid:
            try:
                con = self._connect()
                row = con.execute(
                    '''SELECT s.session_id, s.account_id, s.tz_offset_seconds,
                              s.apply_changes_immediately,
                              a.email, a.access_level
                       FROM sessions s
                       JOIN accounts a ON a.account_id = s.account_id
                       WHERE s.session_id=?''',
                    (sid,)
                ).fetchone()
                con.close()
                if row:
                    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(row['apply_changes_immediately']),
                    }
                    return SqliteSession(data, sid=sid, new=False)
            except Exception:
                pass
        return SqliteSession(sid=str(uuid.uuid4()), new=True)

    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:
            if 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')
        if not account_id:
            return

        now           = int(time.time())
        tz_offset     = int(session.get('tz_offset_seconds', 0))
        apply_changes = 1 if session.get('apply_changes_immediately') else 0

        try:
            con = self._connect()
            if session.new:
                con.execute(
                    '''INSERT INTO sessions
                       (session_id, account_id, tz_offset_seconds, apply_changes_immediately,
                        session_started_ts, last_seen_ts)
                       VALUES (?,?,?,?,?,?)''',
                    (session.sid, account_id, tz_offset, apply_changes, now, now)
                )
            elif session.modified:
                con.execute(
                    '''UPDATE sessions SET tz_offset_seconds=?, apply_changes_immediately=?,
                       last_seen_ts=? WHERE session_id=?''',
                    (tz_offset, apply_changes, now, session.sid)
                )
            else:
                con.execute(
                    'UPDATE sessions SET last_seen_ts=? WHERE session_id=?',
                    (now, session.sid)
                )
            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),
        )