auth.py

"""
OmniParse AI - Authentication Module

Handles user registration, login, session management, and API key generation.
Uses bcrypt for password hashing and delegates all storage to database.py.

Late-imports database functions inside call-bodies to avoid circular imports
(auth ↔ database).
"""

import secrets

import bcrypt

from validation import validate_email, validate_name, validate_password, sanitize_string


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

def hash_pw(password: str) -> str:
    """Hash *password* with bcrypt (12 rounds) and return the hash as a UTF-8 string."""
    salt = bcrypt.gensalt(rounds=12)
    hashed = bcrypt.hashpw(password.encode("utf-8"), salt)
    return hashed.decode("utf-8")


def verify_pw(password: str, hashed: str) -> bool:
    """Return True if *password* matches the bcrypt *hashed* string."""
    try:
        return bcrypt.checkpw(password.encode("utf-8"), hashed.encode("utf-8"))
    except Exception:
        return False


# ---------------------------------------------------------------------------
# Token / API-key generators
# ---------------------------------------------------------------------------

def gen_token() -> str:
    """Generate a cryptographically random session token (URL-safe, 32 bytes)."""
    return secrets.token_urlsafe(32)


def gen_api_key() -> str:
    """Generate a unique API key with a recognizable prefix."""
    return "op_live_" + secrets.token_hex(20)


# ---------------------------------------------------------------------------
# User creation
# ---------------------------------------------------------------------------

def create_user(email: str, name: str, password: str, plan: str = "free"):
    """Register a new user after validating all inputs.

    Args:
        email: User email address.
        name: Display name.
        password: Plain-text password (will be hashed).
        plan: Subscription plan (default "free").

    Returns:
        (user_dict, "") on success.
        (None, "error message") on failure.
    """
    # --- Input validation ---
    valid, err = validate_email(email)
    if not valid:
        return None, err

    valid, err = validate_name(name)
    if not valid:
        return None, err

    valid, errs = validate_password(password)
    if not valid:
        return None, "; ".join(errs)

    email = sanitize_string(email).lower()
    name = sanitize_string(name)

    if plan not in ("free", "basic", "pro", "enterprise"):
        return None, "Invalid plan."

    # --- Late import to avoid circular dependency ---
    from database import get_user_by_email, create_user as db_create_user

    # --- Check uniqueness ---
    existing = get_user_by_email(email)
    if existing is not None:
        return None, "An account with this email already exists."

    # --- Hash & store ---
    hashed = hash_pw(password)
    user = db_create_user(email=email, name=name, password=hashed, plan=plan)
    return user, ""


# ---------------------------------------------------------------------------
# User lookups
# ---------------------------------------------------------------------------

def get_user_by_email(email: str):
    """Fetch a user by email. Returns user dict or None."""
    from database import get_user_by_email as db_get_user_by_email
    return db_get_user_by_email(email)


def get_user_by_id(uid: str):
    """Fetch a user by ID. Returns user dict or None."""
    from database import get_user_by_id as db_get_user_by_id
    return db_get_user_by_id(uid)


# ---------------------------------------------------------------------------
# User update
# ---------------------------------------------------------------------------

def update_user(uid: str, fields: dict):
    """Update user fields. Returns the updated user dict or None on failure."""
    from database import update_user as db_update_user
    return db_update_user(uid, fields)


# ---------------------------------------------------------------------------
# Session management
# ---------------------------------------------------------------------------

def create_session(user_id: str) -> str:
    """Create a new session for *user_id* and return the session token."""
    from database import create_session as db_create_session
    token = gen_token()
    db_create_session(user_id=user_id, token=token)
    return token


def get_session_user(token: str):
    """Look up the user associated with *token*.

    Returns:
        user_dict if the session is valid and not expired, else None.
    """
    from database import get_session_user as db_get_session_user
    return db_get_session_user(token)


def delete_session(token: str) -> None:
    """Delete a session token (logout)."""
    from database import delete_session as db_delete_session
    db_delete_session(token)


# ---------------------------------------------------------------------------
# Authentication (login)
# ---------------------------------------------------------------------------

def authenticate_user(email: str, password: str, ip: str = ""):
    """Verify credentials (session is created separately by the caller).

    Args:
        email: User email.
        password: Plain-text password.
        ip: Client IP address for rate-limiting.

    Returns:
        (user_dict, "") on success.
        (None, "error message") on failure.
    """
    # --- Input validation ---
    valid, err = validate_email(email)
    if not valid:
        return None, err

    if not isinstance(password, str) or not password:
        return None, "Password is required."

    email = sanitize_string(email).lower()

    # --- Rate limit check ---
    if ip:
        from rate_limiter import check_auth_rate
        if not check_auth_rate(ip):
            return None, "Too many login attempts. Please try again later."

    # --- Look up user ---
    from database import get_user_by_email as db_get_user_by_email
    user = db_get_user_by_email(email)
    if user is None:
        return None, "Invalid email or password."

    # --- Verify password ---
    stored_hash = user.get("password", "")
    if not stored_hash:
        return None, "Invalid email or password."

    if not verify_pw(password, stored_hash):
        return None, "Invalid email or password."

    return user, ""


# ---------------------------------------------------------------------------
# Demo account bootstrap
# ---------------------------------------------------------------------------

def ensure_demo_account():
    """Create a demo account if one does not already exist."""
    from database import get_user_by_email, create_user as db_create_user
    existing = get_user_by_email("demo@omniparse.ai")
    if existing is not None:
        return
    try:
        hashed = hash_pw("demo1234")
        db_create_user(
            email="demo@omniparse.ai",
            name="Demo User",
            password=hashed,
            plan="pro",
        )
    except Exception as e:
        print(f"[INFO] Demo account not created: {e}")