app/security.py

"""Auth, RBAC, and API-key management (the three must-haves).

1. Superadmin login  -> a login session token with expiry.
2. User session expiry -> login_sessions.expires_at (LOGIN_SESSION_TTL_HOURS).
3. Provider API keys  -> api_keys.expires_at (lifetime or N days, default 2 weeks).

Roles:
- 'user'       : end users. scopes = chat + websearch + image (NO agentic tools).
- 'superadmin' : full access (chat, websearch, image, agents, commands, files, mcp).
"""
from __future__ import annotations

import hashlib
import hmac
import re
import secrets
import time
from datetime import datetime, timedelta, timezone

from . import config, db

# ---------- brute-force protection (in-memory, per-instance) ----------
_fails: dict[tuple, list[float]] = {}
MAX_FAILS = 5         # failures allowed per window before lockout
FAIL_WINDOW = 300.0   # 5-minute sliding window


def too_many_fails(ip: str, bucket: str) -> bool:
    now = time.time()
    k = (ip, bucket)
    arr = [t for t in _fails.get(k, []) if now - t < FAIL_WINDOW]
    _fails[k] = arr
    return len(arr) >= MAX_FAILS


def record_fail(ip: str, bucket: str) -> None:
    _fails.setdefault((ip, bucket), []).append(time.time())


def clear_fails(ip: str, bucket: str) -> None:
    _fails.pop((ip, bucket), None)


# ---------- input hardening ----------
_SAFE_SID = re.compile(r"[^a-zA-Z0-9_.\-]")


def safe_session_id(sid: str) -> str:
    """Strip anything that could traverse the filesystem; keep it a flat token."""
    sid = _SAFE_SID.sub("_", str(sid))[:128]
    sid = sid.replace("..", "_")
    return sid or "default"


def within(child, parent) -> bool:
    """True only if `child` resolves to a path inside `parent` (containment guarantee)."""
    from pathlib import Path
    try:
        Path(child).resolve().relative_to(Path(parent).resolve())
        return True
    except (ValueError, OSError):
        return False


# ---------- prompt-extraction / jailbreak guard (deterministic, pre-model) ----------
_JAILBREAK = re.compile(
    r"(?is)("
    r"ignore\s+(?:the\s+|all\s+|any\s+|your\s+|previous\s+|prior\s+|above\s+|earlier\s+)*"
    r"(?:instruction|rule|prompt|guideline|directive)s?"
    r"|(?:reveal|show|print|repeat|give\s+me|tell\s+me|share|output|display|leak|expose|"
    r"what(?:'s| is| are)|state|recite)\b[^.\n]{0,40}\b(?:your|the|its|first|initial|original|hidden|secret)\b"
    r"[^.\n]{0,25}\b(?:system\s+)?(?:prompt|instruction|rule|configuration|guideline|directive)s?"
    r"|repeat\s+(?:everything|all|the\s+text|what(?:'s| is)?)\b[^.\n]{0,20}\babove"
    r"|(?:developer|debug|god|admin|dev|sudo|dan)\s*mode"
    r"|\bjailbreak\b|\bsystem\s*prompt\b"
    r"|pretend\s+(?:you\s+are|to\s+be)\s+(?:a\s+)?(?:different|another|new)\s+(?:ai|system|model|assistant)"
    r")"
)

EXTRACTION_REFUSAL = (
    "I can't share my internal configuration or instructions — but I'm happy to help "
    "with a question, search, or task."
)


def is_extraction_attempt(text: str) -> bool:
    return bool(_JAILBREAK.search(text or ""))

# RBAC: what each role is allowed to do.
ROLE_SCOPES = {
    "user": ["chat", "websearch", "image"],
    "superadmin": ["chat", "websearch", "image", "agents", "commands", "files", "mcp"],
}


# ---------- password hashing (stdlib pbkdf2, no extra deps) ----------
def hash_password(password: str) -> str:
    salt = secrets.token_hex(16)
    dk = hashlib.pbkdf2_hmac("sha256", password.encode(), salt.encode(), 200_000)
    return f"pbkdf2$200000${salt}${dk.hex()}"


def verify_password(password: str, stored: str) -> bool:
    try:
        algo, iters, salt, want = stored.split("$")
        dk = hashlib.pbkdf2_hmac("sha256", password.encode(), salt.encode(), int(iters))
        return hmac.compare_digest(dk.hex(), want)
    except Exception:
        return False


# ---------- superadmin bootstrap ----------
def ensure_superadmin() -> None:
    """Create/refresh the superadmin from env on startup (no-op without DB or password)."""
    if not db.available() or not config.SUPERADMIN_PASSWORD:
        return
    db.init_schema()
    db.upsert_user(config.SUPERADMIN_USER, hash_password(config.SUPERADMIN_PASSWORD), "superadmin")


# ---------- login sessions (with expiry) ----------
def login(username: str, password: str) -> dict | None:
    user = db.get_user(username)
    if not user or not verify_password(password, user["password_hash"]):
        return None
    token = "sess_" + secrets.token_urlsafe(32)
    expires = datetime.now(timezone.utc) + timedelta(hours=config.LOGIN_SESSION_TTL_HOURS)
    db.create_login(token, user["id"], user["role"], expires)
    return {"token": token, "role": user["role"], "expires_at": expires.isoformat(),
            "user_id": user["id"]}


def check_login(token: str) -> dict | None:
    """Return the login session if valid and not expired, else None."""
    if not token:
        return None
    return db.get_login(token)


# ---------- provider API keys (with expiry + scopes) ----------
def _hash_key(raw: str) -> str:
    return hashlib.sha256(raw.encode()).hexdigest()


def generate_api_key(role: str, label: str = "", scopes: list[str] | None = None,
                     ttl_days: int | None = None, created_by: int | None = None,
                     max_sessions: int | None = None, agents_md: str | None = None) -> dict:
    """Create a key. Returns the FULL key once (store it; only the hash is persisted).
    ttl_days: None -> default; 0 -> lifetime. max_sessions: None/0 -> unlimited.
    agents_md: optional per-user AGENTS.md instructions applied to this key only."""
    role = role if role in ROLE_SCOPES else "user"
    scopes = scopes or ROLE_SCOPES[role]
    raw = "ant_" + secrets.token_urlsafe(32)
    prefix = raw[:12] + "…"
    if ttl_days is None:
        ttl_days = config.DEFAULT_API_KEY_TTL_DAYS
    expires_at = None if ttl_days == 0 else datetime.now(timezone.utc) + timedelta(days=ttl_days)
    if max_sessions in (None, 0):
        max_sessions = None  # unlimited
    agents_md = (agents_md or "").strip() or None
    rec = db.create_api_key(_hash_key(raw), prefix, label, role, scopes, created_by,
                            expires_at, max_sessions, agents_md)
    rec["api_key"] = raw  # shown once
    rec["expires_at"] = expires_at.isoformat() if expires_at else "lifetime"
    rec["max_sessions"] = max_sessions or "unlimited"
    return rec


def validate_api_key(raw: str) -> dict | None:
    """Return the key record if valid (exists, not revoked, not expired), else None."""
    if not raw:
        return None
    rec = db.get_api_key_by_hash(_hash_key(raw))
    if not rec or rec["revoked"]:
        return None
    if rec["expires_at"] and rec["expires_at"] <= datetime.now(timezone.utc):
        return None
    return rec


def scopes_for(rec: dict) -> list[str]:
    return rec.get("scopes") or ROLE_SCOPES.get(rec.get("role", "user"), ROLE_SCOPES["user"])