auth.py

"""Authentication, encryption, and session management.

Encryption model (defense in depth)
==================================

Every "important" secret in the panel is encrypted **twice**::

  stored = USER_FERNET.encrypt(STATIC_FERNET.encrypt(plaintext))

* **Inner** — ``STATIC_FERNET`` uses ``PANEL_ENCRYPTION_KEY`` (32-byte
  Fernet key from env).  Compromising the DB on its own is not
  enough — the panel admin key is still required to read anything.
* **Outer** — ``USER_FERNET`` uses a per-user key generated at
  signup, stored encrypted-by-PANEL-KEY on the ``User`` row, and
  optionally mirrored to Supabase.  Even if ``PANEL_ENCRYPTION_KEY``
  leaks, the attacker still needs every user's outer key to read
  their secrets.

Backwards compatibility
-----------------------

Older rows that were encrypted only with the static key are still
readable.  ``SecretVault.decrypt`` tries the two-tier path first,
and falls back to a single-layer decrypt on the same value if the
outer layer is absent or doesn't decrypt cleanly.

When a user *writes* a secret, the panel always uses the
two-tier form (assuming the user has a key).  Unwrapped single-layer
rows persist on disk and get re-wrapped the next time the user
edits that secret.

Supabase mirror
---------------

If ``SUPABASE_URL`` and ``SUPABASE_SERVICE_KEY`` are configured,
``PerUserKeyManager`` mirrors every user key into a ``user_keys``
table in the user's Supabase project.  This gives a second,
cloud-resident copy in case the panel's storage is wiped — the
user can re-import keys from Supabase on a fresh install.
"""
from __future__ import annotations

import base64
import hashlib
import hmac
import logging
import secrets
import time
from dataclasses import dataclass
from typing import Optional

import bcrypt
from cryptography.fernet import Fernet, InvalidToken
from fastapi import Request, HTTPException, status
from itsdangerous import URLSafeTimedSerializer, BadSignature, SignatureExpired

from config import PANEL_CONFIG

log = logging.getLogger("hosting-panel.auth")


# ---------------------------------------------------------------------------
# Two-tier secret vault
# ---------------------------------------------------------------------------


class SecretVault:
    """Encrypt / decrypt bot secrets with the panel's static key.

    Pass a per-user key to enable the second tier.  Decryption tries
    both tiers and falls back to single-tier if the outer is missing
    or malformed (legacy rows).
    """

    def __init__(self) -> None:
        raw = PANEL_CONFIG.encryption_key
        if not raw:
            # Without this key the panel can't decrypt anything, but we
            # don't want startup to die — let the admin see a clear
            # warning at the encrypt/decrypt boundary.
            self._fernet: Optional[Fernet] = None
            log.warning(
                "PANEL_ENCRYPTION_KEY is not set — encrypted secrets "
                "will fail to decrypt. Set the env var and restart."
            )
            return
        try:
            self._fernet = Fernet(raw.encode() if isinstance(raw, str) else raw)
        except Exception as exc:
            raise RuntimeError(f"Invalid PANEL_ENCRYPTION_KEY: {exc}") from exc

    @property
    def is_configured(self) -> bool:
        return self._fernet is not None

    # -- encrypt -----------------------------------------------------------

    def encrypt(self, plaintext: str, user_key: Optional[str] = None) -> str:
        """Encrypt ``plaintext`` with the static key, then optionally
        wrap the result with ``user_key``.

        With no user key, this is a single-layer Fernet — legacy format.
        """
        if not plaintext:
            return plaintext
        if not self._fernet:
            raise RuntimeError("PANEL_ENCRYPTION_KEY not configured")

        inner = self._fernet.encrypt(plaintext.encode()).decode()
        if not user_key:
            return inner
        user_fernet = self._user_fernet(user_key)
        return user_fernet.encrypt(inner.encode()).decode()

    # -- decrypt -----------------------------------------------------------

    def decrypt(self, ciphertext: Optional[str], user_key: Optional[str] = None) -> str:
        """Decrypt ``ciphertext``.  Tries the layered path first, then
        falls back to the single-tier path.  Returns ``""`` on empty
        input — never raises unless the value is real but undecryptable
        with *both* keys.
        """
        if not ciphertext:
            return ""
        if not self._fernet:
            raise RuntimeError("PANEL_ENCRYPTION_KEY not configured")

        # Try two-tier first (with user key)
        if user_key:
            try:
                outer = self._user_fernet(user_key).decrypt(ciphertext.encode())
                return self._fernet.decrypt(outer).decode()
            except InvalidToken:
                pass  # may be a single-tier legacy row — fall through
            except Exception:
                pass

        # Fallback: single-tier (PANEL key only)
        try:
            return self._fernet.decrypt(ciphertext.encode()).decode()
        except InvalidToken as exc:
            raise ValueError(
                "Cannot decrypt secret — PANEL_ENCRYPTION_KEY has likely "
                "changed since this row was written, or the value is corrupted. "
                "If you rotated the key, this secret must be re-entered."
            ) from exc

    # -- helpers -----------------------------------------------------------

    @staticmethod
    def _user_fernet(user_key: str) -> Fernet:
        """Validate a user-supplied Fernet key.

        Accepts raw bytes (b"..."), urlsafe-base64, or urlsafe-base64
        without padding.  We pad and validate.
        """
        if isinstance(user_key, str):
            kb = user_key.encode()
        else:
            kb = user_key
        try:
            return Fernet(kb)
        except Exception:
            # try padding
            pad = -len(kb) % 4
            try:
                return Fernet(kb + b"=" * pad)
            except Exception as exc:
                raise ValueError(f"Invalid user key: {exc}") from exc

    @staticmethod
    def generate_user_key() -> str:
        """Generate a fresh 32-byte urlsafe-base64 Fernet key."""
        return Fernet.generate_key().decode()


VAULT = SecretVault()


# ---------------------------------------------------------------------------
# Per-user key management
# ---------------------------------------------------------------------------


@dataclass
class UserKey:
    """Plaintext per-user Fernet key + metadata.

    ``created_at`` is set locally so the panel can warn the user about
    rotation age; ``last_used`` is updated on every decrypt.
    """
    key: str
    created_at: float
    source: str = "local"        # "local" | "supabase"
    key_id: Optional[str] = None # opaque id (Supabase row primary key)


class PerUserKeyManager:
    """Persistent store + Supabase mirror for per-user Fernet keys.

    The panel keeps the canonical copy on ``User.user_key_enc`` (a
    Fernet-wrapped Fernet key).  When Supabase is configured, we also
    mirror the plaintext key to a ``user_keys`` table there so a
    fresh panel can import keys from a cloud-resident copy.
    """

    def __init__(self) -> None:
        self._supabase = None
        if PANEL_CONFIG.supabase_url and PANEL_CONFIG.supabase_service_key:
            try:
                from supabase import create_client
                self._supabase = create_client(
                    PANEL_CONFIG.supabase_url,
                    PANEL_CONFIG.supabase_service_key,
                )
                log.info("Supabase mirror enabled: %s", PANEL_CONFIG.supabase_url)
            except Exception as exc:
                log.warning("Supabase client init failed — mirror disabled: %s", exc)

    @property
    def mirror_enabled(self) -> bool:
        return self._supabase is not None

    # -- read / write ------------------------------------------------------

    def get_user_key(self, user) -> Optional[str]:
        """Return the plaintext per-user Fernet key for ``user``.

        ``user`` is a ``models.User`` instance.  Returns
        ``None`` when the user doesn't have a key yet (e.g. pre-layer
        accounts).
        """
        if user.user_key_enc:
            try:
                return VAULT.decrypt(user.user_key_enc)
            except Exception as exc:
                log.warning("Failed to decrypt user_key for %s: %s", user.username, exc)
        return None

    async def ensure_user_key(self, user, db_session) -> str:
        """Make sure the user has a key; create one if not.

        Returns the plaintext key.  Persists to the local User row
        *and* mirrors to Supabase if enabled.  Works on both sync and
        async ``db_session`` — detects which one we got.
        """
        existing = self.get_user_key(user)
        if existing:
            return existing

        new_key = VAULT.generate_user_key()
        user.user_key_enc = VAULT.encrypt(new_key)
        await _commit_session(db_session)

        if self._supabase:
            try:
                self._supabase.table("user_keys").insert({
                    "user_id": user.id,
                    "username": user.username,
                    "user_key": new_key,
                    "created_at": time.time(),
                }).execute()
            except Exception as exc:
                log.warning("Supabase mirror insert failed for user %s: %s",
                            user.username, exc)
        return new_key

    async def rotate_user_key(self, user, db_session) -> str:
        """Generate a new key, re-wrap all of the user's secrets with it,
        persist both copies.  Returns the new plaintext key."""
        import inspect
        from sqlalchemy import select
        from models import BotInstance

        old_key = self.get_user_key(user)
        new_key = VAULT.generate_user_key()

        # Re-wrap all of the user's bot secrets.
        cls_name = type(db_session).__name__
        is_async = "Async" in cls_name

        if is_async:
            stmt = select(BotInstance).where(BotInstance.owner_id == user.id)
            result = await db_session.execute(stmt)
            bots = list(result.scalars().all())
        else:
            stmt = select(BotInstance).where(BotInstance.owner_id == user.id)
            result = db_session.execute(stmt)
            bots = list(result.scalars().all())

        old_inner_decrypts: list[tuple] = []
        for bot in bots:
            for field in ("bot_token_enc", "hf_token_enc",
                          "nvidia_api_key_enc", "llm_proxy_token_enc"):
                cipher = getattr(bot, field)
                if not cipher:
                    continue
                try:
                    plain = VAULT.decrypt(cipher, user_key=old_key)
                except Exception:
                    try:
                        plain = VAULT.decrypt(cipher)
                    except Exception:
                        log.warning("Cannot decrypt %s on bot %s during rotation; skipping",
                                    field, getattr(bot, "id", "?"))
                        plain = None
                if plain is not None:
                    old_inner_decrypts.append((bot, field, plain))

        # Persist the new user key on the User row FIRST so the next
        # encrypt() uses it.
        user.user_key_enc = VAULT.encrypt(new_key)
        await _commit_session(db_session)

        # Re-encrypt every secret with the new key (and PANEL key).
        for bot, field, plain in old_inner_decrypts:
            setattr(bot, field, VAULT.encrypt(plain, user_key=new_key))
        await _commit_session(db_session)

        if self._supabase:
            try:
                self._supabase.table("user_keys").upsert({
                    "user_id": user.id,
                    "username": user.username,
                    "user_key": new_key,
                    "created_at": time.time(),
                }).execute()
            except Exception as exc:
                log.warning("Supabase mirror upsert failed: %s", exc)

        log.info("Rotated user key for %s — re-wrapped %d secret(s)",
                 user.username, len(old_inner_decrypts))
        return new_key

    def export_to_supabase(self, user) -> bool:
        """Force-mirror the user's key to Supabase.  Returns ``True``
        on success."""
        if not self._supabase:
            return False
        key = self.get_user_key(user)
        if not key:
            return False
        try:
            self._supabase.table("user_keys").upsert({
                "user_id": user.id,
                "username": user.username,
                "user_key": key,
                "created_at": time.time(),
            }).execute()
            return True
        except Exception as exc:
            log.warning("Supabase mirror failed for user %s: %s", user.username, exc)
            return False

    def import_from_supabase(self, user) -> bool:
        """Pull a Supabase mirror back into the local User row."""
        if not self._supabase:
            return False
        try:
            res = (self._supabase.table("user_keys")
                   .select("user_key")
                   .eq("user_id", user.id)
                   .execute())
        except Exception as exc:
            log.warning("Supabase read failed: %s", exc)
            return False
        rows = res.data or []
        if not rows:
            return False
        key = rows[0]["user_key"]
        user.user_key_enc = VAULT.encrypt(key)
        return True


PER_USER_KEYS = PerUserKeyManager()


# ---------------------------------------------------------------------------
# Password hashing
# ---------------------------------------------------------------------------

def hash_password(password: str) -> str:
    """Hash a password with bcrypt."""
    return bcrypt.hashpw(password.encode(), bcrypt.gensalt(rounds=12)).decode()


def verify_password(password: str, hashed: str) -> bool:
    """Verify a password against a bcrypt hash."""
    try:
        return bcrypt.checkpw(password.encode(), hashed.encode())
    except Exception:
        return False


# ---------------------------------------------------------------------------
# Session management (signed cookies)
# ---------------------------------------------------------------------------

@dataclass
class SessionUser:
    user_id: int
    username: str
    role: str
    expires_at: float


_serializer = URLSafeTimedSerializer(PANEL_CONFIG.secret_key, salt="panel-session")


def create_session(user_id: int, username: str, role: str) -> str:
    """Create a signed session token."""
    expires_at = time.time() + PANEL_CONFIG.session_max_age_sec
    data = {"uid": user_id, "u": username, "r": role, "exp": expires_at}
    return _serializer.dumps(data)


def parse_session(token: str) -> Optional[SessionUser]:
    """Parse and validate a session token."""
    try:
        data = _serializer.loads(token, max_age=PANEL_CONFIG.session_max_age_sec)
    except (BadSignature, SignatureExpired):
        return None
    expires_at = data.get("exp", 0)
    if time.time() > expires_at:
        return None
    return SessionUser(
        user_id=data["uid"],
        username=data["u"],
        role=data["r"],
        expires_at=expires_at,
    )


def get_session_user(request: Request) -> Optional[SessionUser]:
    """Extract session from cookie."""
    token = request.cookies.get("session")
    if not token:
        return None
    return parse_session(token)


def require_auth(request: Request) -> SessionUser:
    """Dependency: raise 302 to /login if not authenticated."""
    user = get_session_user(request)
    if not user:
        raise HTTPException(status_code=status.HTTP_302_FOUND, headers={"Location": "/login"})
    return user


def require_role(*roles: str):
    """Dependency factory: require specific role(s)."""
    def _check(request: Request) -> SessionUser:
        user = require_auth(request)
        if user.role not in roles:
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail="Insufficient permissions.",
            )
        return user
    return _check


# ---------------------------------------------------------------------------
# CSRF protection
# ---------------------------------------------------------------------------

def generate_csrf_token() -> str:
    """Generate a CSRF token."""
    return secrets.token_urlsafe(32)


def validate_csrf(request: Request, token: str) -> bool:
    """Validate CSRF token from form against cookie."""
    cookie_token = request.cookies.get("csrf_token")
    if not cookie_token or not token:
        return False
    return hmac.compare_digest(cookie_token, token)


# ---------------------------------------------------------------------------
# Invite codes
# ---------------------------------------------------------------------------

def generate_invite_code() -> str:
    """Generate a random invite code."""
    return secrets.token_urlsafe(16)


# ---------------------------------------------------------------------------
# Client info for audit logging
# ---------------------------------------------------------------------------

def get_client_info(request: Request) -> dict[str, Optional[str]]:
    """Extract client IP and UA from request."""
    forwarded = request.headers.get("x-forwarded-for")
    ip = forwarded.split(",")[0].strip() if forwarded else request.client.host if request.client else None
    return {
        "ip_address": ip,
        "user_agent": request.headers.get("user-agent"),
    }


# ---------------------------------------------------------------------------
# Sync/async session adapter
# ---------------------------------------------------------------------------
# Some helpers (the per-user key rotation, secret re-wrapping, Supabase
# mirror sync) need to work whether the caller is in a sync context
# (CLI tools, ``run_admin.py``) or an async request handler.
# ---------------------------------------------------------------------------

async def _commit_session(db_session) -> None:
    """Commit a SQLAlchemy session whether it's sync or async."""
    import inspect
    commit = db_session.commit()
    if inspect.isawaitable(commit):
        await commit
    else:
        # Already committed in the sync path.
        pass


async def _exec(db_session, query_fn):
    """Run a synchronous query in both sync and async sessions.

    Returns the raw result.  For async sessions, this uses
    ``run_sync`` to bridge.
    """
    import inspect
    cls = db_session if not inspect.isclass(db_session) else None
    name = type(db_session).__name__
    if "Async" in name or "AsyncSession" in name:
        from sqlalchemy import select
        # Sync query into async session.
        return await db_session.run_sync(query_fn)
    return query_fn(db_session)