Cortex-Inara/cortex/auth_utils.py

"""
Authentication utilities — password hashing and JWT session tokens.

Passwords are stored as bcrypt hashes in home/{username}/auth.json.
Sessions are JWT cookies signed with JWT_SECRET from settings.

Usage:
  set_password("scott", "mypassword")       # admin setup
  check_credentials("scott", "mypassword")  # login validation
  create_token("scott")                     # returns JWT string
  decode_token(token)                       # returns username or raises
"""

import json
import logging
import secrets
from datetime import datetime, timedelta, timezone
from pathlib import Path

import bcrypt
import jwt

from config import settings

logger = logging.getLogger(__name__)

COOKIE_NAME = "cortex_session"
ALGORITHM = "HS256"


# ---------------------------------------------------------------------------
# Password helpers
# ---------------------------------------------------------------------------

def _auth_path(username: str) -> Path:
    return settings.home_root() / username / "auth.json"


def set_password(username: str, password: str) -> None:
    """Hash and store a password for a user. Creates auth.json if needed."""
    hashed = bcrypt.hashpw(password.encode(), bcrypt.gensalt()).decode()
    _auth_path(username).write_text(json.dumps({"password_hash": hashed}) + "\n")
    logger.info("password set for user: %s", username)


def check_credentials(username: str, password: str) -> bool:
    """Return True if username+password are valid, False otherwise."""
    path = _auth_path(username)
    if not path.exists():
        return False
    try:
        data = json.loads(path.read_text())
        stored = data.get("password_hash", "").encode()
        return bcrypt.checkpw(password.encode(), stored)
    except Exception:
        return False


# ---------------------------------------------------------------------------
# JWT helpers
# ---------------------------------------------------------------------------

def create_token(username: str) -> str:
    """Return a signed JWT encoding the username."""
    expire = datetime.now(timezone.utc) + timedelta(days=settings.jwt_expire_days)
    payload = {"sub": username, "exp": expire}
    return jwt.encode(payload, settings.jwt_secret, algorithm=ALGORITHM)


def decode_token(token: str) -> str:
    """Decode a JWT and return the username. Raises jwt.InvalidTokenError on failure."""
    payload = jwt.decode(token, settings.jwt_secret, algorithms=[ALGORITHM])
    return payload["sub"]


# ---------------------------------------------------------------------------
# Invite tokens — one-time setup links for new users
# ---------------------------------------------------------------------------

def _invite_path(username: str) -> Path:
    return settings.home_root() / username / "invite.json"


def create_invite(username: str, expire_hours: int = 72) -> str:
    """
    Generate a one-time invite token for a user and save it to invite.json.
    Returns the raw token string (embed in a URL).
    """
    token = secrets.token_urlsafe(32)
    expires = (datetime.now(timezone.utc) + timedelta(hours=expire_hours)).isoformat()
    user_dir = settings.home_root() / username
    user_dir.mkdir(parents=True, exist_ok=True)
    _invite_path(username).write_text(
        json.dumps({"token": token, "expires_at": expires, "used": False}) + "\n"
    )
    logger.info("invite created for user: %s (expires %s)", username, expires[:10])
    return token


def validate_invite(token: str) -> str | None:
    """
    Check an invite token across all users.
    Returns the username if valid and unused, None otherwise.
    """
    root = settings.home_root()
    if not root.exists():
        return None
    for user_dir in root.iterdir():
        if not user_dir.is_dir():
            continue
        inv_path = user_dir / "invite.json"
        if not inv_path.exists():
            continue
        try:
            data = json.loads(inv_path.read_text())
        except Exception:
            continue
        if data.get("used"):
            continue
        if data.get("token") != token:
            continue
        expires = datetime.fromisoformat(data["expires_at"])
        if datetime.now(timezone.utc) > expires:
            continue
        return user_dir.name
    return None


def consume_invite(username: str) -> None:
    """Mark the invite token for a user as used."""
    path = _invite_path(username)
    if path.exists():
        try:
            data = json.loads(path.read_text())
            data["used"] = True
            path.write_text(json.dumps(data) + "\n")
        except Exception:
            pass