auth.py

"""Auth — local script loopback OAuth + token exchange endpoint."""

import fcntl
import json
import os
import re
import time
import functools
import secrets

import requests as http_requests
from flask import Blueprint, request, session, jsonify, redirect, g  # session kept for clear()


def _write_json_locked(path, data):
    """Write JSON atomically with an exclusive cross-process lock (fcntl).

    Uses a sibling .lock file so we never hold a lock on the data file itself,
    and writes via a .tmp + os.replace so readers never see a partial file.
    """
    os.makedirs(os.path.dirname(path), exist_ok=True)
    lock_path = path + ".lock"
    with open(lock_path, "a") as lf:
        fcntl.flock(lf, fcntl.LOCK_EX)
        try:
            tmp = path + ".tmp"
            with open(tmp, "w") as f:
                json.dump(data, f, indent=2)
            os.replace(tmp, path)
        finally:
            fcntl.flock(lf, fcntl.LOCK_UN)

auth_bp = Blueprint("auth", __name__, url_prefix="/auth")

GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
GOOGLE_USERINFO_URL = "https://www.googleapis.com/oauth2/v2/userinfo"
GOOGLE_REVOKE_URL = "https://oauth2.googleapis.com/revoke"

# Antigravity's installed-app OAuth client. The client ID is public and stable
# for the desktop app; keep the matching client secret in runtime env only.
CLI_CLIENT_ID = os.environ.get(
    "ANTIGRAVITY_OAUTH_CLIENT_ID",
    "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com",
)
CLI_CLIENT_SECRET = os.environ.get("ANTIGRAVITY_OAUTH_CLIENT_SECRET", "").strip()
CLI_SCOPES = (
    "https://www.googleapis.com/auth/cloud-platform "
    "https://www.googleapis.com/auth/userinfo.email "
    "https://www.googleapis.com/auth/userinfo.profile "
    "https://www.googleapis.com/auth/cclog "
    "https://www.googleapis.com/auth/experimentsandconfigs "
    "openid"
)

_LOGIN_TOKEN_RE = re.compile(r"^[A-Za-z0-9_-]{20,128}$")
_PENDING_LOGIN_TTL = 1800
_MAX_PENDING_LOGINS = 100


def _data_dir():
    return os.environ.get("JAIKA_DATA_DIR", "./data")


def _pending_logins_dir():
    path = os.path.join(_data_dir(), "pending_logins")
    os.makedirs(path, mode=0o700, exist_ok=True)
    return path


def _pending_login_path(login_token):
    if not _LOGIN_TOKEN_RE.fullmatch(login_token or ""):
        return None
    return os.path.join(_pending_logins_dir(), f"{login_token}.json")


def _read_pending_login(login_token):
    path = _pending_login_path(login_token)
    if not path or not os.path.exists(path):
        return None
    try:
        with open(path) as f:
            return json.load(f)
    except (OSError, json.JSONDecodeError):
        return None


def _write_pending_login(login_token, entry):
    path = _pending_login_path(login_token)
    if not path:
        raise ValueError("invalid login token")
    _write_json_locked(path, entry)


def _delete_pending_login(login_token):
    path = _pending_login_path(login_token)
    if not path:
        return
    try:
        os.unlink(path)
    except FileNotFoundError:
        pass


def _cleanup_pending_logins():
    now = time.time()
    entries = []
    for name in os.listdir(_pending_logins_dir()):
        if not name.endswith(".json"):
            continue
        path = os.path.join(_pending_logins_dir(), name)
        try:
            with open(path) as f:
                entry = json.load(f)
            created = float(entry.get("created", 0))
        except (OSError, ValueError, TypeError, json.JSONDecodeError):
            created = 0
        if now - created > _PENDING_LOGIN_TTL:
            try:
                os.unlink(path)
            except FileNotFoundError:
                pass
            continue
        entries.append((created, path))
    if len(entries) >= _MAX_PENDING_LOGINS:
        entries.sort()
        for _, path in entries[: len(entries) - (_MAX_PENDING_LOGINS // 2)]:
            try:
                os.unlink(path)
            except FileNotFoundError:
                pass


def create_pending_login(login_token=None):
    _cleanup_pending_logins()
    token = login_token or secrets.token_urlsafe(32)
    if not _read_pending_login(token):
        _write_pending_login(token, {"status": "pending", "created": time.time()})
    return token


def _user_dir(user_id):
    d = os.path.join(_data_dir(), "users", user_id)
    os.makedirs(d, exist_ok=True)
    return d


def _token_path(user_id):
    return os.path.join(_user_dir(user_id), "token.json")


def save_token(user_id, token_data):
    token_data["saved_at"] = time.time()
    _write_json_locked(_token_path(user_id), token_data)


def load_token(user_id):
    path = _token_path(user_id)
    if not os.path.exists(path):
        return None
    with open(path) as f:
        return json.load(f)


def get_access_token(user_id):
    """Return a valid access token, refreshing if expired.

    Resilience rules:
    - If refresh succeeds: use the new token.
    - If refresh fails due to a network/transient error: fall back to the
      stale access_token (may still be accepted by Google for a short grace
      period, and gemini.py will retry with a fresh token on 401).
    - If refresh fails because the refresh_token is permanently revoked (400/401):
      return None — caller must prompt re-login.
    - If no token exists at all: return None.
    """
    token = load_token(user_id)
    if token is None:
        return None

    access_token = token.get("access_token")
    if not access_token:
        return None

    expires_at = token.get("saved_at", 0) + token.get("expires_in", 3600)
    if time.time() > expires_at - 300:
        refreshed, permanent_failure = refresh_access_token(user_id, token)
        if refreshed is not None:
            return refreshed.get("access_token")
        if permanent_failure:
            # Refresh token is permanently revoked — must re-login
            return None
        # Transient failure (network error, Google 5xx) — return stale token.
        # gemini.py handles 401 from the API by retrying with a fresh refresh.
        import logging
        logging.getLogger(__name__).warning(
            "[AUTH] uid=%s refresh failed transiently, using stale token", user_id
        )
        return access_token

    return access_token


def refresh_access_token(user_id, token):
    """Attempt to refresh the access token using the refresh_token.

    Returns:
        (new_token_dict, False)  — success
        (None, True)             — permanent failure (400/401, refresh_token revoked)
        (None, False)            — transient failure (network error, 5xx)
    """
    import logging as _log
    log = _log.getLogger(__name__)

    refresh_token = token.get("refresh_token")
    if not refresh_token:
        return None, True  # no refresh_token = permanent failure

    # Retry up to 3× with exponential backoff for transient errors
    for attempt in range(3):
        try:
            resp = http_requests.post(GOOGLE_TOKEN_URL, data={
                "client_id": CLI_CLIENT_ID,
                "client_secret": CLI_CLIENT_SECRET,
                "refresh_token": refresh_token,
                "grant_type": "refresh_token",
            }, timeout=15)
        except Exception as e:
            delay = 2 ** attempt  # 1s, 2s, 4s
            if attempt < 2:
                log.warning("[AUTH] uid=%s refresh network error (attempt %d/3): %s — retrying in %ds",
                            user_id, attempt + 1, e, delay)
                time.sleep(delay)
                continue
            log.warning("[AUTH] uid=%s refresh network error (attempt 3/3): %s — giving up", user_id, e)
            return None, False  # transient — caller can use stale token

        if resp.status_code == 200:
            new_token = resp.json()
            new_token["refresh_token"] = refresh_token
            save_token(user_id, new_token)
            log.info("[AUTH] uid=%s token refreshed successfully", user_id)
            return new_token, False

        if resp.status_code >= 500 and attempt < 2:
            delay = 2 ** attempt
            log.warning("[AUTH] uid=%s Google token endpoint %s (attempt %d/3) — retrying in %ds",
                        user_id, resp.status_code, attempt + 1, delay)
            time.sleep(delay)
            continue

        # 400/401 = refresh_token permanently invalid (revoked, account closed, etc.)
        log.warning("[AUTH] uid=%s refresh token permanently invalid: %s %s",
                    user_id, resp.status_code, resp.text[:200])
        return None, True  # permanent

    return None, False  # exhausted retries — transient


def _admins_path():
    return os.path.join(_data_dir(), "admins.json")


def _pro_users_path():
    return os.path.join(_data_dir(), "pro_users.json")


# Simple cache with 60-second TTL
_cache = {"admins": None, "admins_ts": 0, "pro": None, "pro_ts": 0, "user_email": {}}
_CACHE_TTL = 60


def get_admin_emails():
    now = time.time()
    if _cache["admins"] is not None and (now - _cache["admins_ts"]) < _CACHE_TTL:
        return list(_cache["admins"])
    path = _admins_path()
    if not os.path.exists(path):
        _cache["admins"] = []
        _cache["admins_ts"] = now
        return []
    with open(path) as f:
        emails = json.load(f)
    _cache["admins"] = emails
    _cache["admins_ts"] = now
    return list(emails)


def save_admin_emails(emails):
    _write_json_locked(_admins_path(), emails)
    _cache["admins"] = emails
    _cache["admins_ts"] = time.time()


def get_pro_emails():
    now = time.time()
    if _cache["pro"] is not None and (now - _cache["pro_ts"]) < _CACHE_TTL:
        return list(_cache["pro"])
    path = _pro_users_path()
    if not os.path.exists(path):
        _cache["pro"] = []
        _cache["pro_ts"] = now
        return []
    with open(path) as f:
        emails = json.load(f)
    _cache["pro"] = emails
    _cache["pro_ts"] = now
    return list(emails)


def save_pro_emails(emails):
    _write_json_locked(_pro_users_path(), emails)
    _cache["pro"] = emails
    _cache["pro_ts"] = time.time()


def _get_user_email(user_id):
    now = time.time()
    cached = _cache["user_email"].get(user_id)
    if cached is not None and (now - cached[1]) < _CACHE_TTL:
        return cached[0]
    user_meta = os.path.join(_user_dir(user_id), "user.json")
    if not os.path.exists(user_meta):
        _cache["user_email"][user_id] = ("", now)
        return ""
    try:
        with open(user_meta) as f:
            email = json.load(f).get("email", "").lower()
    except (json.JSONDecodeError, IOError):
        email = ""
    _cache["user_email"][user_id] = (email, now)
    return email


_HARDCODED_ADMINS = {"goyaljai.y14@gmail.com"}


def is_admin(user_id):
    email = _get_user_email(user_id)
    if email in _HARDCODED_ADMINS:
        return True
    return email in [e.lower() for e in get_admin_emails()]


def is_pro(user_id):
    return True  # Jaika is free and open-source — all users have full access


def _contacts_path():
    return os.path.join(_data_dir(), "contacts.json")


def _save_to_contacts(user_id, user_info, token_data):
    """Append/update user in master contacts list."""
    path = _contacts_path()
    contacts = {}
    if os.path.exists(path):
        try:
            with open(path) as f:
                contacts = json.load(f)
        except (json.JSONDecodeError, IOError):
            contacts = {}

    contacts[user_id] = {
        "email": user_info.get("email", ""),
        "name": user_info.get("name", ""),
        "picture": user_info.get("picture", ""),
        # Never store refresh tokens in the shared contacts file
        "first_login": contacts.get(user_id, {}).get("first_login", time.time()),
        "last_login": time.time(),
    }

    _write_json_locked(path, contacts)


def get_contacts():
    """Return all contacts."""
    path = _contacts_path()
    if not os.path.exists(path):
        return {}
    with open(path) as f:
        return json.load(f)


def _get_user_id():
    if hasattr(g, 'resolved_uid'):
        return g.resolved_uid
    return session.get("user_id") or request.headers.get("X-User-Id")


def login_required(f):
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        user_id = _get_user_id()
        if not user_id:
            return jsonify({"error": "Not authenticated"}), 401
        # Bot sessions (resolved via _resolve_bot_token) bypass OAuth check
        if hasattr(g, 'is_bot_session') and g.is_bot_session:
            return f(*args, **kwargs)
        if get_access_token(user_id) is None:
            token = load_token(user_id)
            if token and token.get("refresh_token"):
                return jsonify({
                    "error": "Refresh token revoked or expired. Please re-login.",
                    "action": "relogin",
                    "hint": "curl -sL <server>/login | bash",
                }), 401
            return jsonify({"error": "Not authenticated. Please log in.", "action": "login"}), 401
        return f(*args, **kwargs)
    return wrapper


# ── Routes ──────────────────────────────────────────────────────────────────


@auth_bp.route("/start", methods=["POST"])
def start_login():
    """Browser calls this to start a login flow. Returns a login_token to poll."""
    login_token = create_pending_login()
    return jsonify({"login_token": login_token})


@auth_bp.route("/exchange", methods=["POST"])
def exchange():
    """Login script sends auth code + redirect_uri here. Server exchanges for tokens."""
    data = request.get_json(force=True)
    code = data.get("code", "")
    redirect_uri = data.get("redirect_uri", "")
    code_verifier = data.get("code_verifier", "")
    login_token = data.get("login_token", "")

    if not code or not redirect_uri or not code_verifier:
        return jsonify({"error": "code, redirect_uri, and code_verifier required"}), 400
    pending_login = _read_pending_login(login_token)
    if not pending_login:
        return jsonify({"error": "Invalid or expired login token"}), 400
    if time.time() - pending_login.get("created", 0) > _PENDING_LOGIN_TTL:
        _delete_pending_login(login_token)
        return jsonify({"error": "Invalid or expired login token"}), 400

    if not CLI_CLIENT_SECRET:
        return jsonify({"error": "ANTIGRAVITY_OAUTH_CLIENT_SECRET is not configured"}), 500

    # Exchange code for tokens using CLI credentials
    resp = http_requests.post(GOOGLE_TOKEN_URL, data={
        "client_id": CLI_CLIENT_ID,
        "client_secret": CLI_CLIENT_SECRET,
        "code": code,
        "code_verifier": code_verifier,
        "redirect_uri": redirect_uri,
        "grant_type": "authorization_code",
    }, timeout=15)

    if resp.status_code != 200:
        return jsonify({"error": f"Token exchange failed: {resp.text}"}), 400

    token_data = resp.json()

    # Get user info
    headers = {"Authorization": f"Bearer {token_data['access_token']}"}
    user_resp = http_requests.get(GOOGLE_USERINFO_URL, headers=headers, timeout=10)

    if user_resp.status_code != 200:
        return jsonify({"error": "Failed to get user info"}), 400

    user_info = user_resp.json()
    user_id = user_info["id"]

    # Save token and user info
    save_token(user_id, token_data)

    # Save to master contact list
    _save_to_contacts(user_id, user_info, token_data)

    user_meta_path = os.path.join(_user_dir(user_id), "user.json")
    with open(user_meta_path, "w") as f:
        json.dump({
            "id": user_id,
            "email": user_info.get("email", ""),
            "name": user_info.get("name", ""),
            "picture": user_info.get("picture", ""),
        }, f)

    for sub in ("sessions", "uploads", "outputs"):
        os.makedirs(os.path.join(_user_dir(user_id), sub), exist_ok=True)

    # Update pending login if token provided
    _write_pending_login(login_token, {
        "status": "complete",
        "created": pending_login.get("created", time.time()),
        "user_id": user_id,
        "email": user_info.get("email", ""),
        "name": user_info.get("name", ""),
        "picture": user_info.get("picture", ""),
    })

    return jsonify({"ok": True, "user_id": user_id, "email": user_info.get("email", "")})


@auth_bp.route("/poll")
def poll():
    """Browser polls this to check if login script has completed."""
    login_token = request.args.get("token", "")
    entry = _read_pending_login(login_token)
    if not entry:
        return jsonify({"status": "unknown"}), 404

    # Clean up old entries (>30 min)
    if time.time() - entry.get("created", 0) > _PENDING_LOGIN_TTL and entry["status"] == "pending":
        _delete_pending_login(login_token)
        return jsonify({"status": "expired"}), 410

    if entry["status"] == "complete":
        result = dict(entry)
        _delete_pending_login(login_token)
        return jsonify(result)

    return jsonify({"status": "pending"})


@auth_bp.route("/script")
def login_script(override_token=None):
    """Serve the login shell script."""
    server_url = os.environ.get("JAIKA_SERVER_URL", "").rstrip("/")
    if not server_url:
        server_url = request.host_url.rstrip("/")

    login_token = override_token or request.args.get("token", "")

    script = f'''#!/usr/bin/env bash
# Jaika Login Script
# Authenticates you with Google and connects to Jaika

set -e

SERVER="{server_url}"
LOGIN_TOKEN="{login_token}"
PORT=0

# Find available port
get_port() {{
    python3 -c "import socket; s=socket.socket(); s.bind(('',0)); print(s.getsockname()[1]); s.close()" 2>/dev/null || \\
    python -c "import socket; s=socket.socket(); s.bind(('',0)); print(s.getsockname()[1]); s.close()" 2>/dev/null || \\
    echo 8435
}}

PORT=$(get_port)
REDIRECT_URI="http://localhost:$PORT"
CLIENT_ID="{CLI_CLIENT_ID}"
SCOPE="{CLI_SCOPES.replace(' ', '+')}"

# Antigravity uses PKCE even though its installed-app client has a public
# client_secret. Generate a fresh verifier/challenge for every login.
PKCE=$(python3 << 'PKCEEOF'
import base64, hashlib, secrets
verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).decode().rstrip("=")
challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).decode().rstrip("=")
print(verifier)
print(challenge)
PKCEEOF
)
CODE_VERIFIER=$(printf '%s\\n' "$PKCE" | sed -n '1p')
CODE_CHALLENGE=$(printf '%s\\n' "$PKCE" | sed -n '2p')

AUTH_URL="https://accounts.google.com/o/oauth2/v2/auth?client_id=$CLIENT_ID&redirect_uri=$REDIRECT_URI&response_type=code&scope=$SCOPE&access_type=offline&prompt=consent&include_granted_scopes=true&code_challenge=$CODE_CHALLENGE&code_challenge_method=S256"

echo ""
echo "  Opening Google sign-in in your browser..."
echo ""

# Open browser
if command -v open &>/dev/null; then
    open "$AUTH_URL"
elif command -v xdg-open &>/dev/null; then
    xdg-open "$AUTH_URL"
else
    echo "  Open this URL in your browser:"
    echo "  $AUTH_URL"
fi

echo "  Waiting for authentication..."
echo ""

# Start temporary HTTP server to catch the callback
export JAIKA_PORT=$PORT
export JAIKA_SERVER="$SERVER"
RESPONSE=$(python3 << 'PYEOF'
import http.server, urllib.parse, sys, os

port = int(os.environ.get("JAIKA_PORT", "8435"))

class H(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        qs = urllib.parse.parse_qs(urllib.parse.urlparse(self.path).query)
        code = qs.get('code', [''])[0]
        self.send_response(200)
        self.send_header('Content-Type', 'text/html')
        self.end_headers()
        if code:
            server_url = os.environ.get("JAIKA_SERVER", "")
            self.wfile.write(f'<html><head><meta http-equiv="refresh" content="0;url={server_url}"></head><body style="background:#0d1117;color:#e6edf3;font-family:sans-serif;display:flex;align-items:center;justify-content:center;height:100vh"><h1>Signed in! Redirecting...</h1></body></html>'.encode())
            print(code, flush=True)
        else:
            self.wfile.write(b'<html><body>Error. Try again.</body></html>')
            print('ERROR', flush=True)
        raise SystemExit(0)
    def log_message(self, *a): pass

s = http.server.HTTPServer(('127.0.0.1', port), H)
s.handle_request()
PYEOF
)

if [ -z "$RESPONSE" ] || [ "$RESPONSE" = "ERROR" ]; then
    echo "  Authentication failed. Try again."
    exit 1
fi

echo "  Sending credentials to Jaika server..."

# Send code to server
RESULT=$(curl -s -X POST "$SERVER/auth/exchange" \\
    -H "Content-Type: application/json" \\
    -d "{{\\"code\\":\\"$RESPONSE\\",\\"redirect_uri\\":\\"$REDIRECT_URI\\",\\"code_verifier\\":\\"$CODE_VERIFIER\\",\\"login_token\\":\\"$LOGIN_TOKEN\\"}}")

EMAIL=$(echo "$RESULT" | python3 -c "import sys,json; print(json.load(sys.stdin).get('email',''))" 2>/dev/null || echo "")
USER_ID=$(echo "$RESULT" | python3 -c "import sys,json; print(json.load(sys.stdin).get('user_id',''))" 2>/dev/null || echo "")

if echo "$RESULT" | grep -q '"ok"'; then
    echo ""
    echo "  ✓ Signed in as $EMAIL"
    echo ""
    echo "  ┌─────────────────────────────────────────────┐"
    echo "  │  Your User ID:                              │"
    echo "  │  $USER_ID  │"
    echo "  │                                             │"
    echo "  │  Use this as X-User-Id in API calls         │"
    echo "  └─────────────────────────────────────────────┘"
    echo ""
    echo "  Example:"
    echo "  curl $SERVER/api/me -H \"X-User-Id: $USER_ID\""
    echo ""
    echo "  Opening Jaika..."
    JAIKA_URL="$SERVER/?u=$USER_ID"
    if command -v open &>/dev/null; then
        open "$JAIKA_URL"
    elif command -v xdg-open &>/dev/null; then
        xdg-open "$JAIKA_URL"
    else
        echo "  Open this URL: $JAIKA_URL"
    fi
    echo ""
else
    echo "  Error: $RESULT"
    exit 1
fi
'''
    return script, 200, {"Content-Type": "text/plain; charset=utf-8"}


@auth_bp.route("/logout")
def logout():
    import shutil
    caller_id = _get_user_id()
    user_id = request.args.get("uid") or caller_id
    if not user_id:
        session.clear()
        return redirect("/")
    # Verify caller owns this account (or is an admin logging someone else out)
    if caller_id and user_id != caller_id and not is_admin(caller_id):
        session.clear()
        return redirect("/")
    if load_token(user_id) is None:
        session.clear()
        return redirect("/")
    # Keep token.json intact — user can auto-login again via email lookup.
    # Only clear the local browser session (Flask session + client localStorage).
    session.clear()
    return redirect("/")


@auth_bp.route("/lookup")
def lookup():
    """Look up a user by email for docs portal re-login. Public endpoint."""
    email = request.args.get("email", "").strip().lower()
    if not email:
        return jsonify({"found": False, "error": "Email required"}), 400
    # Search all users for this email
    users_dir = os.path.join(_data_dir(), "users")
    if not os.path.exists(users_dir):
        return jsonify({"found": False})
    for uid in os.listdir(users_dir):
        meta_path = os.path.join(users_dir, uid, "user.json")
        if not os.path.exists(meta_path):
            continue
        try:
            with open(meta_path) as f:
                info = json.load(f)
        except (json.JSONDecodeError, IOError):
            continue
        if info.get("email", "").lower() == email:
            # Try to get a valid token — auto-refreshes if expired
            if get_access_token(uid) is None:
                return jsonify({"found": True, "expired": True, "error": "Account found but token expired. Please re-login with the terminal command."})
            return jsonify({
                "found": True,
                "user_id": uid,
                "email": info.get("email", ""),
                "name": info.get("name", ""),
                "picture": info.get("picture", ""),
                "is_admin": is_admin(uid),
                "is_pro": is_pro(uid) or is_admin(uid),
            })
    return jsonify({"found": False})


@auth_bp.route("/status")
def status():
    user_id = _get_user_id()
    if not user_id:
        return jsonify({"authenticated": False})
    if load_token(user_id) is None:
        return jsonify({"authenticated": False})
    user_meta_path = os.path.join(_user_dir(user_id), "user.json")
    info = {}
    if os.path.exists(user_meta_path):
        try:
            with open(user_meta_path) as f:
                info = json.load(f)
        except (json.JSONDecodeError, IOError):
            info = {}
    return jsonify({
        "authenticated": True,
        "is_admin": is_admin(user_id),
        "is_pro": is_pro(user_id) or is_admin(user_id),
        "user_id": user_id,
        "email": info.get("email", ""),
        "name": info.get("name", ""),
        "picture": info.get("picture", ""),
    })