tests/test_launch_blocker_security.py

from __future__ import annotations

import asyncio
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timedelta, timezone

import jwt
import pytest
from starlette.requests import Request


def _auth(token: str) -> dict[str, str]:
    return {"Authorization": f"Bearer {token}"}


def _signup(client, email: str = "launch-security@example.test") -> str:
    response = client.post(
        "/auth/signup",
        json={"name": "Launch Security", "email": email, "password": "Pass123!beta"},
    )
    assert response.status_code == 201, response.text
    return response.json()["access_token"]


def _request_with_headers(headers: dict[str, str]) -> Request:
    return Request(
        {
            "type": "http",
            "method": "POST",
            "path": "/ask",
            "headers": [(name.lower().encode(), value.encode()) for name, value in headers.items()],
            "client": ("198.51.100.25", 12345),
            "server": ("testserver", 80),
            "scheme": "http",
            "query_string": b"",
        },
    )


def test_production_auth_disabled_fails_before_init_db(monkeypatch, tmp_path):
    from app.core.config import get_settings
    from app import main

    monkeypatch.setenv("ENVIRONMENT", "production")
    monkeypatch.setenv("AUTH_ENABLED", "false")
    monkeypatch.setenv("AUTH_PROVIDER", "jwt")
    monkeypatch.setenv("JWT_SECRET_KEY", "production-safe-test-secret-32-chars")
    monkeypatch.setenv("FRONTEND_BASE_URL", "https://docdoe.ai")
    monkeypatch.setenv("DATABASE_URL", f"sqlite:///{tmp_path / 'prod.db'}")
    get_settings.cache_clear()

    def fail_if_called() -> None:
        raise AssertionError("init_db should not run when production auth is disabled")

    async def enter_lifespan() -> None:
        async with main.lifespan(main.app):
            pass

    monkeypatch.setattr(main, "init_db", fail_if_called)
    with pytest.raises(RuntimeError, match="AUTH_ENABLED must be true"):
        asyncio.run(enter_lifespan())


def test_production_sqlite_database_url_is_blocked(monkeypatch, tmp_path):
    from app.core.config import get_settings
    from app.main import _startup_safety_checks

    monkeypatch.setenv("ENVIRONMENT", "production")
    monkeypatch.setenv("AUTH_ENABLED", "true")
    monkeypatch.setenv("AUTH_PROVIDER", "jwt")
    monkeypatch.setenv("JWT_SECRET_KEY", "production-safe-test-secret-32-chars")
    monkeypatch.setenv("FRONTEND_BASE_URL", "https://docdoe.ai")
    monkeypatch.setenv("CORS_ORIGINS", "https://docdoe.ai")
    monkeypatch.setenv("DATABASE_URL", f"sqlite:///{tmp_path / 'prod.db'}")
    get_settings.cache_clear()

    with pytest.raises(RuntimeError, match="DATABASE_URL must be PostgreSQL"):
        _startup_safety_checks()


def test_compute_video_routes_are_rate_limited_without_polling_gets() -> None:
    from app.main import _RATE_LIMIT_PREFIXES, _should_rate_limit_request

    assert "/video/study-video-jobs" in _RATE_LIMIT_PREFIXES
    assert "/video/render-jobs" in _RATE_LIMIT_PREFIXES
    assert _should_rate_limit_request("POST", "/video/study-video-jobs/job_1/render-preview")
    assert not _should_rate_limit_request("GET", "/video/study-video-jobs/job_1")


def test_invalid_forged_jwt_not_trusted_for_rate_limit_key(monkeypatch):
    from app.core.config import get_settings
    from app.main import _rate_limit_key

    monkeypatch.setenv("AUTH_PROVIDER", "jwt")
    monkeypatch.setenv("JWT_SECRET_KEY", "correct-rate-limit-secret-32-chars")
    monkeypatch.setenv("JWT_ALGORITHM", "HS256")
    get_settings.cache_clear()

    expires_at = datetime.now(timezone.utc) + timedelta(minutes=10)
    forged = jwt.encode(
        {"sub": "forged-user", "exp": expires_at},
        "wrong-rate-limit-secret-at-least-32",
        algorithm="HS256",
    )
    request = _request_with_headers(
        {
            "Authorization": f"Bearer {forged}",
            "X-Forwarded-For": "203.0.113.44, 10.0.0.1",
        },
    )

    assert _rate_limit_key(request) == "ip:198.51.100.25"


def test_rate_limit_ignores_client_supplied_forwarded_for() -> None:
    from app.main import _rate_limit_key

    request = _request_with_headers({"X-Forwarded-For": "203.0.113.99"})

    assert _rate_limit_key(request) == "ip:198.51.100.25"


def test_valid_jwt_is_trusted_for_rate_limit_key(monkeypatch):
    from app.core.config import get_settings
    from app.main import _rate_limit_key

    secret = "correct-rate-limit-secret-32-chars"
    monkeypatch.setenv("AUTH_PROVIDER", "jwt")
    monkeypatch.setenv("JWT_SECRET_KEY", secret)
    monkeypatch.setenv("JWT_ALGORITHM", "HS256")
    get_settings.cache_clear()

    token = jwt.encode(
        {"sub": "real-user", "exp": datetime.now(timezone.utc) + timedelta(minutes=10)},
        secret,
        algorithm="HS256",
    )
    request = _request_with_headers(
        {
            "Authorization": f"Bearer {token}",
            "X-Forwarded-For": "203.0.113.45",
        },
    )

    assert _rate_limit_key(request) == "user:real-user"


def test_billing_paid_plan_direct_upgrade_blocked_in_production(auth_client, monkeypatch):
    from app.core.config import get_settings

    token = _signup(auth_client, "billing-prod@example.test")
    monkeypatch.setenv("ENVIRONMENT", "production")
    monkeypatch.setenv("AUTH_ENABLED", "true")
    monkeypatch.setenv("AUTH_PROVIDER", "jwt")
    monkeypatch.setenv("JWT_SECRET_KEY", "test-only-secret-for-auth-tests-32chars!")
    monkeypatch.setenv("FRONTEND_BASE_URL", "https://docdoe.ai")
    get_settings.cache_clear()

    paid = auth_client.post(
        "/billing/select-plan",
        headers=_auth(token),
        json={"plan": "popular_299"},
    )
    assert paid.status_code == 402
    assert "checkout" in paid.json()["detail"].lower()

    current = auth_client.get("/billing/me", headers=_auth(token))
    assert current.status_code == 200
    assert current.json()["selected_plan"] == "free_trial"

    free = auth_client.post(
        "/billing/select-plan",
        headers=_auth(token),
        json={"plan": "free_trial"},
    )
    assert free.status_code == 200
    assert free.json()["selected_plan"] == "free_trial"


def test_usage_recording_uses_atomic_database_increments(client):
    """Atomic conditional UPDATE prevents over-limit and races simultaneously.

    Production hardening switched ``record_generation``/``record_video_plan``
    from unconditional increments to ``UPDATE ... WHERE used + N <= limit`` so
    we can't exceed the quota. To exercise concurrency we widen the limits so
    every request fits inside the cap.
    """
    from app.core.database import SessionLocal
    from app.models.user import User
    from app.models.user_plan import UserPlan
    from app.services.usage_service import (
        get_or_create_user_plan,
        record_generation,
        record_video_plan,
    )
    from sqlalchemy import select

    user_id = "usr_atomic_launch"
    with SessionLocal() as db:
        db.add(User(id=user_id, name="Atomic User", email="atomic@example.test"))
        db.commit()
        plan = get_or_create_user_plan(db, user_id)
        plan.monthly_generation_used = 0
        plan.monthly_video_used = 0
        # Widen limits so the test isolates atomicity, not quota enforcement.
        plan.monthly_generation_limit = 100
        plan.monthly_video_limit = 50
        db.add(plan)
        db.commit()

    def increment_generation() -> None:
        with SessionLocal() as db:
            record_generation(db, user_id)

    def increment_video() -> None:
        with SessionLocal() as db:
            record_video_plan(db, user_id)

    with ThreadPoolExecutor(max_workers=4) as executor:
        list(executor.map(lambda _: increment_generation(), range(20)))
        list(executor.map(lambda _: increment_video(), range(12)))

    with SessionLocal() as db:
        plan = db.scalar(select(UserPlan).where(UserPlan.user_id == user_id))
        assert plan is not None
        assert plan.monthly_generation_used == 20
        assert plan.monthly_video_used == 12


def test_real_user_api_does_not_return_bundled_demo_data(auth_client):
    token = _signup(auth_client, "no-demo-data@example.test")
    headers = _auth(token)

    sources = auth_client.get("/sources", headers=headers)
    assert sources.status_code == 200
    assert sources.json()["sources"] == []

    dashboard = auth_client.get("/dashboard/student", headers=headers)
    assert dashboard.status_code == 200
    assert dashboard.json()["materials"]["total"] == 0
    assert dashboard.json()["recent_results"] == []

    papers = auth_client.get("/previous-papers", headers=headers)
    assert papers.status_code == 200
    assert papers.json() == []

    pyq = auth_client.post("/pyq/analyze", headers=headers, json={"subject": "Physics"})
    assert pyq.status_code == 200
    pyq_body = pyq.json()
    assert pyq_body["available"] is False
    assert pyq_body["predicted_questions"] == []

    combined = f"{sources.text}\n{dashboard.text}\n{papers.text}\n{pyq.text}"
    assert "Electromagnetic Induction" not in combined
    assert "JEE Main Physics PYQ Set" not in combined


def test_empty_previous_paper_account_never_runtime_seeds(
    auth_client,
    monkeypatch,
):
    """Production behavior must match tests: an empty user owns zero papers."""
    monkeypatch.delenv("PYTEST_CURRENT_TEST", raising=False)
    token = _signup(auth_client, "empty-pyq-account@example.test")

    response = auth_client.get("/previous-papers", headers=_auth(token))

    assert response.status_code == 200
    assert response.json() == []


def test_weak_topic_boost_does_not_dominate_irrelevant_chunks(client):
    from app.core.database import SessionLocal
    from app.models.document import Document
    from app.models.document_chunk import DocumentChunk
    from app.models.user import User
    from app.services.retrieval import retrieve_relevant_chunks
    from app.services.weak_topic_service import record_weak_topic

    user_id = "usr_weak_boost"
    document_id = "doc_weak_boost"
    relevant_chunk_id = "chunk_relevant_photosynthesis"

    with SessionLocal() as db:
        db.add(User(id=user_id, name="Weak Topic User", email="weak@example.test"))
        db.add(
            Document(
                id=document_id,
                user_id=user_id,
                title="Biology Notes",
                file_name="biology.txt",
                file_type="text/plain",
                file_path="/tmp/biology.txt",
                subject="Biology",
                status="ready",
                extracted_text="Photosynthesis converts carbon dioxide and water into glucose.",
                chunk_count=2,
            ),
        )
        db.add_all(
            [
                DocumentChunk(
                    id=relevant_chunk_id,
                    document_id=document_id,
                    chunk_index=0,
                    chunk_text=(
                        "Photosynthesis uses light energy to convert carbon dioxide "
                        "and water into glucose and oxygen."
                    ),
                    token_estimate=20,
                    heading="Photosynthesis",
                ),
                DocumentChunk(
                    id="chunk_irrelevant_weak_topic",
                    document_id=document_id,
                    chunk_index=1,
                    chunk_text="Quantum tunneling is a weak area but it is unrelated to plant nutrition.",
                    token_estimate=16,
                    heading="Unrelated physics note",
                ),
            ],
        )
        db.commit()
        record_weak_topic(db, user_id, "quantum tunneling", subject="Biology")

        results = retrieve_relevant_chunks(
            db,
            document_id=document_id,
            query="photosynthesis glucose oxygen",
            limit=2,
            user_id=user_id,
        )

    assert results
    assert results[0].chunk.id == relevant_chunk_id