##########################################################################
#
# pgAdmin 4 - PostgreSQL Tools
#
# Copyright (C) 2013 - 2026, The pgAdmin Development Team
# This software is released under the PostgreSQL Licence
#
##########################################################################

"""
Implements the server-side session management.

Credit/Reference: http://flask.pocoo.org/snippets/109/

Modified to support both Python 2.6+ & Python 3.x
"""

import base64
import datetime
import hmac
import hashlib
import logging
import os
import pickle
import secrets
import string
import time
import config
from uuid import uuid4
from threading import Lock
from flask import current_app, request, flash, redirect, has_request_context
from flask_login import login_url

from collections import OrderedDict
from cryptography.fernet import Fernet, InvalidToken
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from itsdangerous import signer

from flask.sessions import SessionInterface, SessionMixin
from werkzeug.datastructures import CallbackDict
from werkzeug.security import safe_join
from werkzeug.exceptions import InternalServerError

from pgadmin.utils.ajax import make_json_response


def _calc_hmac(body, secret):
    return base64.b64encode(
        hmac.new(
            secret.encode(), body.encode(), hashlib.sha256
        ).digest()
    ).decode()


# File-HMAC: protects on-disk session files against tampering before
# deserialization. SHA-256 is hard-coded (not driven by SESSION_DIGEST_METHOD)
# so the header length is a fixed, format-stable constant.
_FILE_HMAC = hashlib.sha256
_HMAC_HEX_LEN = _FILE_HMAC().digest_size * 2  # 64 for SHA-256

logger = logging.getLogger(__name__)


def _compute_file_hmac(secret, body):
    """Compute the hex-encoded HMAC over a session-file body.

    Flask's SECRET_KEY may be str or bytes; both are accepted.
    """
    key = secret.encode() if isinstance(secret, str) else secret
    return hmac.new(key, body, _FILE_HMAC).hexdigest().encode()


# HKDF salt + info for the session-body Fernet key. Both fixed and version-
# tagged so a future format change can derive a fresh key without reusing
# bytes that may already be in flight on disk.
_SESSION_FERNET_SALT = b'pga-session-body-v1'
_SESSION_FERNET_INFO = b'pgadmin session body encryption v1'


def _derive_session_fernet(secret):
    """Derive a Fernet instance from `SECRET_KEY` for session-body
    confidentiality.

    Layer 1 of the on-disk session protection. The HMAC header (added in
    PR 1) protects integrity of the on-disk bytes; this Fernet layer
    additionally protects confidentiality so that a leak of `sessions/`
    files alone (without the SECRET_KEY in `pgadmin4.db`) does not
    expose OAuth tokens, cloud credentials, MFA OTPs, `pass_enc_key`
    (the KEK that decrypts saved server passwords), or any other
    sensitive material the application places in the session.

    Caveat (real, not theoretical): SECRET_KEY currently lives in
    `pgadmin4.db` in the same DATA_DIR. A leak that includes BOTH
    `sessions/` and `pgadmin4.db` recovers the derived Fernet key and
    can decrypt session bodies. Closing this gap requires moving
    SECRET_KEY out of DATA_DIR (e.g., into the OS keychain via
    `USE_OS_SECRET_STORAGE`); tracked as Layer 2 follow-up.

    HKDF with a fixed salt and info string is used so the derivation is
    deterministic across processes (multiple workers must produce the
    same key) and so a future format change can swap in a new
    salt/info without compromising existing files.
    """
    key_material = secret.encode() if isinstance(secret, str) else secret
    derived = HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=_SESSION_FERNET_SALT,
        info=_SESSION_FERNET_INFO,
    ).derive(key_material)
    return Fernet(base64.urlsafe_b64encode(derived))


def _open_session_file(path):
    """Open a session file for writing with mode 0o600 (owner-only).

    Session files contain OAuth access/refresh tokens, AWS/Google/Azure
    cloud credentials, the Kerberos cache path, MFA OTP material,
    and `pass_enc_key` (the symmetric KEK used to decrypt the user's saved
    Postgres server passwords). The default `open(path, 'wb')` uses the
    process umask, which on most systems leaves files world-readable
    (0o644). Force 0o600 so that even if the directory permissions are
    misconfigured (e.g., a volume mount in a container with shared uids),
    individual session files remain owner-only.
    """
    fd = os.open(
        path,
        os.O_WRONLY | os.O_CREAT | os.O_TRUNC,
        0o600,
    )
    return os.fdopen(fd, 'wb')


sess_lock = Lock()
LAST_CHECK_SESSION_FILES = None


class ManagedSession(CallbackDict, SessionMixin):
    def __init__(self, initial=None, sid=None, new=False, randval=None,
                 hmac_digest=None):
        def on_update(self):
            self.modified = True

        CallbackDict.__init__(self, initial, on_update)
        self.sid = sid
        self.new = new
        self.modified = False
        self.randval = randval
        self.last_write = None
        self.force_write = False
        self.hmac_digest = hmac_digest
        self.permanent = True

    def sign(self, secret):
        if not self.hmac_digest:
            population = string.ascii_lowercase + string.digits

            self.randval = ''.join(
                secrets.choice(population) for i in range(20))
            self.hmac_digest = _calc_hmac(
                '%s:%s' % (self.sid, self.randval), secret)


class SessionManager():
    def new_session(self):
        'Create a new session'
        raise NotImplementedError

    def exists(self, sid):
        'Does the given session-id exist?'
        raise NotImplementedError

    def remove(self, sid):
        'Remove the session'
        raise NotImplementedError

    def get(self, sid, digest):
        'Retrieve a managed session by session-id, checking the HMAC digest'
        raise NotImplementedError

    def put(self, session):
        'Store a managed session'
        raise NotImplementedError


class CachingSessionManager(SessionManager):
    def __init__(self, parent, num_to_store, skip_paths=None):
        self.parent = parent
        self.num_to_store = num_to_store
        self._cache = OrderedDict()
        self.skip_paths = [] if skip_paths is None else skip_paths

    def _normalize(self):
        if len(self._cache) > self.num_to_store:
            # Flush 20% of the cache
            with sess_lock:
                while len(self._cache) > (self.num_to_store * 0.8):
                    self._cache.popitem(False)

    def is_session_ready(self, _session):
        if not has_request_context() or _session is None:
            return False

        # Session _id returns the str object
        # or None if it hasn't been set yet.
        try:
            return _session['_id'] is not None
        except (AssertionError, RuntimeError, KeyError, TypeError):
            return False

    def new_session(self):
        session = self.parent.new_session()

        # Do not store the session if skip paths
        for sp in self.skip_paths:
            if request.path.startswith(sp):
                return session

        with sess_lock:
            self._cache[session.sid] = session
        self._normalize()

        return session

    def remove(self, sid):
        with sess_lock:
            self.parent.remove(sid)
            if sid in self._cache:
                del self._cache[sid]

    def exists(self, sid):
        with sess_lock:
            if sid in self._cache:
                return True
            return self.parent.exists(sid)

    def get(self, sid, digest):
        session = None
        with (sess_lock):
            if sid in self._cache:
                session = self._cache[sid]
                if self.is_session_ready(session) and\
                        session.hmac_digest != digest:
                    session = None

                # reset order in Dict
                del self._cache[sid]

            if not self.is_session_ready(session):
                session = self.parent.get(sid, digest)

            # Do not store the session if skip paths
            for sp in self.skip_paths:
                if request.path.startswith(sp):
                    return session

            self._cache[sid] = session
        self._normalize()

        return session

    def put(self, session):
        with sess_lock:
            self.parent.put(session)

            # Do not store the session if skip paths
            for sp in self.skip_paths:
                if request.path.startswith(sp):
                    return

            if session.sid in self._cache:
                try:
                    del self._cache[session.sid]
                except Exception:
                    pass

            self._cache[session.sid] = session
        self._normalize()


class FileBackedSessionManager(SessionManager):

    def __init__(self, path, secret, disk_write_delay, skip_paths=None):
        if not secret:
            # File-HMAC integrity collapses without a secret. Use raise, not
            # assert: -O strips assertions in production.
            raise RuntimeError("SECRET_KEY must be non-empty")
        self.path = path
        self.secret = secret
        self.disk_write_delay = disk_write_delay
        if not os.path.exists(self.path):
            os.makedirs(self.path)
        self.skip_paths = [] if skip_paths is None else skip_paths

    def exists(self, sid):
        fname = safe_join(self.path, sid)
        return fname is not None and os.path.exists(fname)

    def remove(self, sid):
        fname = safe_join(self.path, sid)
        if fname is not None and os.path.exists(fname):
            os.unlink(fname)

    def new_session(self):
        sid = str(uuid4())
        fname = safe_join(self.path, sid)

        while fname is not None and os.path.exists(fname):
            sid = str(uuid4())
            fname = safe_join(self.path, sid)

        # Do not store the session if skip paths
        for sp in self.skip_paths:
            if request.path.startswith(sp):
                return ManagedSession(sid=sid)

        if fname is None:
            raise InternalServerError('Failed to create new session')

        # touch the file with mode 0o600 — see _open_session_file rationale.
        with _open_session_file(fname):
            return ManagedSession(sid=sid)

        return ManagedSession(sid=sid)

    def get(self, sid, digest):
        """Retrieve a managed session by session-id, verifying integrity
        and decrypting the body.

        File format:
            +-------------------------------------------------------+
            | _HMAC_HEX_LEN bytes : hex HMAC over the ciphertext   |
            +-------------------------------------------------------+
            | N bytes : Fernet(pickle((randval, digest, data)))    |
            +-------------------------------------------------------+

        The HMAC is verified against the ciphertext BEFORE Fernet decrypt
        and BEFORE pickle.loads, so a malicious file dropped in the
        sessions directory cannot trigger arbitrary code execution via
        __reduce__ during deserialization (encrypt-then-MAC; pickle.loads
        only ever sees Fernet-decrypted plaintext from a server-written
        file).
        """
        fname = safe_join(self.path, sid)
        if fname is None or not os.path.exists(fname):
            return self.new_session()

        try:
            with open(fname, 'rb') as f:
                header = f.read(_HMAC_HEX_LEN)
                body = f.read()
            if len(header) != _HMAC_HEX_LEN or len(body) == 0:
                # 0-byte placeholder from new_session(), or pre-fix legacy
                # file shorter than the header. Silently invalidate; one-time
                # re-login on upgrade is acceptable and expected.
                return self.new_session()
            expected = _compute_file_hmac(self.secret, body)
            if not hmac.compare_digest(header, expected):
                logger.warning(
                    "session file rejected: bad file-HMAC for sid=%s",
                    sid[:8] if sid else '<empty>',
                )
                return self.new_session()
            # Body integrity verified — decrypt before deserializing.
            try:
                plaintext = _derive_session_fernet(
                    self.secret).decrypt(body)
            except InvalidToken:
                # Pre-Layer-1 (HMAC-only) session files: HMAC is still a
                # valid HMAC over the legacy plain-pickle body, but
                # Fernet.decrypt on raw pickle bytes raises InvalidToken.
                # Silently invalidate; users see a one-time re-login on
                # the upgrade that introduced this layer.
                logger.warning(
                    "session file rejected: legacy unencrypted body for "
                    "sid=%s", sid[:8] if sid else '<empty>',
                )
                return self.new_session()
            randval, hmac_digest, data = pickle.loads(plaintext)
        except (pickle.UnpicklingError, EOFError, OSError, MemoryError,
                ValueError, TypeError):
            # Narrow catch: silently invalidate corrupted/unreadable files.
            # Programming errors (AttributeError, NameError, etc.) propagate
            # so they surface in tests rather than being masked.
            return self.new_session()

        if not data:
            return self.new_session()

        # Cookie-binding check: the cookie's digest must match the file's
        # stored digest. Prevents using one session's file with another's
        # cookie even when both are server-written.
        if hmac_digest != digest:
            return self.new_session()

        return ManagedSession(
            data, sid=sid, randval=randval, hmac_digest=hmac_digest
        )

    def put(self, session):
        """Store a managed session"""
        current_time = time.time()
        if not session.hmac_digest:
            session.sign(self.secret)
        elif not session.force_write and session.last_write is not None and \
            (current_time - float(session.last_write)) < \
                self.disk_write_delay:
            return

        session.last_write = current_time
        session.force_write = False

        # Do not store the session if skip paths
        for sp in self.skip_paths:
            if request.path.startswith(sp):
                return

        fname = safe_join(self.path, session.sid)

        if fname is None:
            raise InternalServerError('Failed to update the session')

        plaintext = pickle.dumps(
            (session.randval, session.hmac_digest, dict(session)), -1
        )
        # Encrypt-then-MAC: Fernet wrap first, then HMAC over the
        # ciphertext. pickle.loads on a tampered or unauthenticated body
        # is therefore unreachable on the read path.
        body = _derive_session_fernet(self.secret).encrypt(plaintext)
        header = _compute_file_hmac(self.secret, body)
        with _open_session_file(fname) as f:
            f.write(header)
            f.write(body)


class ManagedSessionInterface(SessionInterface):
    def __init__(self, manager):
        self.manager = manager
        signer.Signer.default_digest_method = \
            eval(config.SESSION_DIGEST_METHOD)

    def open_session(self, app, request):
        cookie_val = request.cookies.get(app.config['SESSION_COOKIE_NAME'])

        if not cookie_val or '!' not in cookie_val:
            return self.manager.new_session()

        sid, digest = cookie_val.split('!', 1)

        if self.manager.exists(sid):
            return self.manager.get(sid, digest)

        return self.manager.new_session()

    def save_session(self, app, session, response):
        domain = self.get_cookie_domain(app)
        if not session:
            self.manager.remove(session.sid)
            if session.modified:
                response.delete_cookie(app.config['SESSION_COOKIE_NAME'],
                                       domain=domain)
            return

        if not session.modified:
            # No need to save an unaltered session
            # TODO: put logic here to test if the cookie is older than N days,
            # if so, update the expiration date
            return

        self.manager.put(session)
        session.modified = False

        cookie_exp = self.get_expiration_time(app, session)
        response.set_cookie(
            app.config['SESSION_COOKIE_NAME'],
            '%s!%s' % (session.sid, session.hmac_digest),
            expires=cookie_exp,
            path=config.SESSION_COOKIE_PATH,
            secure=config.SESSION_COOKIE_SECURE,
            httponly=config.SESSION_COOKIE_HTTPONLY,
            samesite=config.SESSION_COOKIE_SAMESITE,
            domain=domain
        )


def create_session_interface(app, skip_paths=[]):
    return ManagedSessionInterface(
        CachingSessionManager(
            FileBackedSessionManager(
                app.config['SESSION_DB_PATH'],
                app.config['SECRET_KEY'],
                app.config.get('PGADMIN_SESSION_DISK_WRITE_DELAY', 10),
                skip_paths
            ),
            1000,
            skip_paths
        ))


def pga_unauthorised():

    lm = current_app.login_manager
    login_message = None

    if lm.login_message:
        if lm.localize_callback is not None:
            login_message = lm.localize_callback(lm.login_message)
        else:
            login_message = lm.login_message

    if not lm.login_view:
        # Only 401 is not enough to distinguish pgAdmin login is required.
        # There are other cases when we return 401. For eg. wrong password
        # supplied while connecting to server.
        # So send additional 'info' message.
        return make_json_response(
            status=401,
            success=0,
            errormsg=login_message,
            info='PGADMIN_LOGIN_REQUIRED'
        )

    # flash messages are only required if the request was from a
    # security page, otherwise it will be redirected to login page
    # anyway
    if login_message and 'security' in request.endpoint:
        flash(login_message, category=lm.login_message_category)

    return redirect(login_url(lm.login_view, request.url))


def cleanup_session_files():
    """
    This function will iterate through session directory and check the last
    modified time, if it older than (session expiration time + 1) days then
    delete that file.
    """
    iterate_session_files = False

    global LAST_CHECK_SESSION_FILES
    if LAST_CHECK_SESSION_FILES is None or \
        datetime.datetime.now() >= LAST_CHECK_SESSION_FILES + \
            datetime.timedelta(hours=config.CHECK_SESSION_FILES_INTERVAL):
        iterate_session_files = True
        LAST_CHECK_SESSION_FILES = datetime.datetime.now()

    if iterate_session_files:
        for root, dirs, files in os.walk(
                current_app.config['SESSION_DB_PATH']):
            for file_name in files:
                absolute_file_name = os.path.join(root, file_name)
                st = os.stat(absolute_file_name)

                # Get the last modified time of the session file
                last_modified_time = \
                    datetime.datetime.fromtimestamp(st.st_mtime)

                # Calculate session file expiry time.
                file_expiration_time = \
                    last_modified_time + \
                    current_app.permanent_session_lifetime + \
                    datetime.timedelta(days=1)

                if file_expiration_time <= datetime.datetime.now() and \
                        os.path.exists(absolute_file_name):
                    os.unlink(absolute_file_name)
