"""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"])