atlas-chat 0.1.0__py3-none-any.whl

atlas/tests/test_security_header_injection.py

@@ -0,0 +1,191 @@
+"""
+Test header injection vulnerabilities.
+
+This test suite demonstrates the header injection attack vector and
+documents why reverse proxy configuration is critical.
+"""
+import pytest
+from fastapi.testclient import TestClient
+from main import app
+
+client = TestClient(app)
+
+
+def test_direct_access_header_injection_vulnerability():
+    """
+    SECURITY WARNING: This test demonstrates a CRITICAL vulnerability.
+
+    When the app is accessed DIRECTLY (bypassing reverse proxy),
+    attackers can inject X-User-Email headers to impersonate any user.
+
+    This test documents the vulnerability. In production:
+    - Main app MUST be network-isolated (not publicly accessible)
+    - ALL traffic MUST go through reverse proxy
+    - Reverse proxy MUST strip client X-User-Email headers
+
+    This test will PASS because the app is designed to trust headers
+    when behind a properly configured reverse proxy. The test serves
+    as documentation of the security requirement.
+    """
+    # Attacker tries to impersonate admin by injecting header
+    response = client.get(
+        "/api/config",
+        headers={"X-User-Email": "attacker-pretending-to-be-admin@evil.com"}
+    )
+
+    # In direct access mode (no proxy), the app trusts this header
+    # This is why network isolation is CRITICAL
+    assert response.status_code == 200
+
+    # The app will treat this as a legitimate request from the attacker's email
+    # In production, this request should NEVER reach the app (network isolation)
+
+
+def test_websocket_header_injection_vulnerability():
+    """
+    Demonstrates WebSocket header injection vulnerability with direct access.
+
+    This shows why the reverse proxy MUST strip X-User-Email headers
+    before adding the authenticated user's header.
+    """
+    # Attacker connects with injected header
+    with client.websocket_connect(
+        "/ws",
+        headers={"X-User-Email": "attacker@evil.com"}
+    ) as websocket:
+        # Connection succeeds because app trusts the header
+        # This is the EXPECTED behavior when behind a proxy that strips headers
+        # This is VULNERABLE behavior when directly accessible
+
+        # Send a test message
+        websocket.send_json({
+            "type": "chat",
+            "content": "test message"
+        })
+
+        # The WebSocket will use "attacker@evil.com" as the user
+        # This demonstrates why network isolation is critical
+
+
+def test_multiple_headers_first_wins():
+    """
+    Demonstrates the danger of improperly configured reverse proxies.
+
+    If the reverse proxy adds X-User-Email without stripping the client's
+    version first, both headers arrive. Most frameworks (including FastAPI)
+    return the FIRST header, allowing the attacker to win.
+
+    Proper nginx config:
+        proxy_set_header X-User-Email "";              # Strip first!
+        proxy_set_header X-User-Email $authenticated_user;  # Then add
+
+    Vulnerable nginx config:
+        proxy_set_header X-User-Email $authenticated_user;  # Only adds, doesn't strip!
+    """
+    # Simulate what happens when proxy doesn't strip headers
+    # We can't easily test multiple headers with TestClient,
+    # but we document the expected behavior
+
+    # When client sends: X-User-Email: attacker@evil.com
+    # And proxy adds: X-User-Email: realuser@example.com
+    # The app receives BOTH headers
+
+    # FastAPI's request.headers.get() returns the FIRST occurrence
+    # So the attacker's header would win!
+
+    # This test documents the requirement for header stripping
+    assert True, "Documented: Proxy must strip headers first"
+
+
+@pytest.mark.skip(reason="Requires production environment with reverse proxy")
+def test_production_header_stripping():
+    """
+    Test to run in production/staging to verify header stripping works.
+
+    This test should be run manually against the actual deployment to verify
+    that the reverse proxy properly strips client-provided headers.
+
+    Usage:
+        1. Deploy to staging/production with reverse proxy
+        2. Get a valid authentication token/cookie
+        3. Run this test against the deployed URL
+        4. Verify logs show the REAL user, not the injected one
+
+    Expected behavior:
+        - Request includes malicious X-User-Email header
+        - Reverse proxy strips it
+        - Reverse proxy adds real authenticated user header
+        - Backend receives only the real user header
+        - Logs confirm backend saw the real user
+    """
+    import os
+
+    import requests
+
+    deployment_url = os.getenv("PRODUCTION_URL")
+    auth_cookie = os.getenv("VALID_AUTH_COOKIE")
+
+    if not deployment_url or not auth_cookie:
+        pytest.skip("Set PRODUCTION_URL and VALID_AUTH_COOKIE env vars")
+
+    # Try to inject a malicious header
+    response = requests.get(
+        f"{deployment_url}/api/config",
+        headers={"X-User-Email": "attacker@evil.com"},
+        cookies={"session": auth_cookie}
+    )
+
+    assert response.status_code == 200
+
+    # Manual verification required:
+    # Check backend logs to confirm it received the REAL user from auth,
+    # not the injected "attacker@evil.com"
+    print("✓ Request succeeded")
+    print("⚠ MANUAL VERIFICATION REQUIRED:")
+    print("  Check backend logs to confirm user was NOT 'attacker@evil.com'")
+    print("  The backend should have received the real authenticated user")
+
+
+def test_header_injection_documentation():
+    """
+    Documentation test: Lists all security requirements for production deployment.
+
+    This test always passes but serves as executable documentation of the
+    security requirements needed to prevent header injection attacks.
+    """
+    security_requirements = [
+        "Main app MUST be network-isolated (not publicly accessible)",
+        "ALL traffic MUST flow through reverse proxy",
+        "Reverse proxy MUST strip client-provided X-User-Email headers",
+        "Reverse proxy MUST add X-User-Email header AFTER stripping client headers",
+        "Direct access to main app ports MUST be blocked by firewall/VPC",
+        "Nginx config MUST include: proxy_set_header X-User-Email '' before setting it",
+        "Apache config MUST include: RequestHeader unset X-User-Email before setting it",
+        "Network isolation MUST be tested (attempt direct access should fail)",
+        "Header injection test MUST be run in production (test_production_header_stripping)",
+        "Deployment checklist in docs/reverse-proxy-examples.md MUST be completed",
+    ]
+
+    for i, requirement in enumerate(security_requirements, 1):
+        print(f"{i}. {requirement}")
+
+    assert True, "Security requirements documented"
+
+
+# Additional test to verify the current behavior
+def test_x_user_email_header_is_used():
+    """
+    Verifies that X-User-Email header is properly extracted.
+
+    This is the expected behavior when behind a properly configured proxy.
+    """
+    test_user = "alice@example.com"
+
+    response = client.get(
+        "/api/config",
+        headers={"X-User-Email": test_user}
+    )
+
+    assert response.status_code == 200
+    # The middleware should have processed this header
+    # In production, this header comes from the reverse proxy, not the client

atlas/tests/test_security_headers_and_filename.py

@@ -0,0 +1,63 @@
+from main import app
+from starlette.testclient import TestClient
+
+
+def test_security_headers_present_by_default():
+    client = TestClient(app)
+    r = client.get("/api/files/healthz", headers={"X-User-Email": "test@test.com"})
+    assert r.status_code == 200
+    # HSTS intentionally omitted
+    assert r.headers.get("X-Content-Type-Options") == "nosniff"
+    assert r.headers.get("X-Frame-Options") in ("SAMEORIGIN", "DENY")
+    assert r.headers.get("Referrer-Policy") is not None
+    # CSP may be present per default value
+    assert "Content-Security-Policy" in r.headers
+
+
+def test_download_filename_sanitized(monkeypatch):
+    # Insert a file into mock S3 listing by calling upload
+    from atlas.infrastructure.app_factory import app_factory
+    app_factory.get_file_manager()
+
+    # Prepare malicious filename
+    bad_name = 'evil\r\nInjected.txt'
+    content = "SGVsbG8="  # base64(Hello)
+
+    async def upload_stub(user_email, filename, content_base64, content_type, tags, source_type):
+        return {
+            "key": "k_mal",
+            "filename": filename,
+            "size": 5,
+            "content_type": "text/plain",
+            "last_modified": "now",
+            "etag": "etag",
+            "tags": tags or {},
+            "user_email": user_email,
+        }
+
+    async def get_stub(user_email, key):
+        return {
+            "key": key,
+            "filename": bad_name,
+            "size": 5,
+            "content_base64": content,
+            "content_type": "text/plain",
+            "last_modified": "now",
+            "etag": "etag",
+            "tags": {},
+        }
+
+    # Patch storage client
+    s3 = app_factory.get_file_storage()
+    monkeypatch.setattr(s3, "upload_file", upload_stub)
+    monkeypatch.setattr(s3, "get_file", get_stub)
+
+    client = TestClient(app)
+
+    # Trigger download endpoint directly (no need to actually upload first)
+    r = client.get("/api/files/download/k_mal", headers={"X-User-Email": "test@test.com"})
+    assert r.status_code == 200
+    cd = r.headers.get("Content-Disposition", "")
+    assert "\r" not in cd and "\n" not in cd
+    assert cd.startswith("attachment;") or cd.startswith("inline;")
+    assert r.headers.get("X-Content-Type-Options") == "nosniff"

atlas/tests/test_shared_session_repository.py

@@ -0,0 +1,101 @@
+"""
+Test that sessions are shared across ChatService instances.
+
+This test verifies the fix for the file upload registration issue where
+files attached in one WebSocket connection were not visible in chat messages
+because each ChatService instance had its own session repository.
+"""
+import uuid
+
+import pytest
+
+from atlas.domain.messages.models import Message, MessageRole
+from atlas.infrastructure.app_factory import app_factory
+
+
+@pytest.mark.asyncio
+async def test_sessions_shared_across_chat_service_instances():
+    """
+    Test that sessions are shared across different ChatService instances.
+
+    This simulates the scenario where:
+    1. A file is attached via one WebSocket connection (ChatService instance 1)
+    2. A chat message is sent via the same or different connection (ChatService instance 2)
+    3. The file should be visible in the session retrieved by instance 2
+    """
+    # Create two ChatService instances (simulating two WebSocket connections)
+    chat_service_1 = app_factory.create_chat_service(connection=None)
+    chat_service_2 = app_factory.create_chat_service(connection=None)
+
+    # Verify they share the same session repository
+    assert chat_service_1.session_repository is chat_service_2.session_repository, \
+        "ChatService instances should share the same session repository"
+
+    user_email = "test@example.com"
+    session_id = uuid.uuid4()
+
+    # Step 1: Create a session and attach a file using ChatService 1
+    session = await chat_service_1.create_session(session_id, user_email)
+    filename = "test-document.pdf"
+    session.context["files"] = {
+        filename: {
+            "key": "s3://bucket/test-document.pdf",
+            "content_type": "application/pdf",
+            "size": 1024,
+            "source": "user",
+        }
+    }
+
+    # Add a message to the session history
+    session.history.add_message(Message(
+        role=MessageRole.USER,
+        content="what files can you see?"
+    ))
+
+    # Step 2: Retrieve the session using ChatService 2 (simulating a different connection)
+    session_from_cs2 = await chat_service_2.session_repository.get(session_id)
+
+    # Verify the session exists and contains the file
+    assert session_from_cs2 is not None, "Session should be accessible from ChatService 2"
+    assert session_from_cs2 is session, "Should be the same session object"
+    assert filename in session_from_cs2.context.get("files", {}), \
+        "File attached in ChatService 1 should be visible in ChatService 2"
+
+    # Verify the session history is also shared
+    messages = session_from_cs2.history.get_messages_for_llm()
+    assert len(messages) == 1, "Session history should be shared"
+    assert messages[0]["content"] == "what files can you see?"
+
+
+@pytest.mark.asyncio
+async def test_session_repository_shared_across_app_factory_calls():
+    """
+    Test that the session repository is truly shared at the AppFactory level.
+
+    This verifies that multiple calls to create_chat_service return
+    ChatService instances that all share the same underlying session storage.
+    """
+    # Create three ChatService instances
+    cs1 = app_factory.create_chat_service()
+    cs2 = app_factory.create_chat_service()
+    cs3 = app_factory.create_chat_service()
+
+    # All should share the same session repository instance
+    assert cs1.session_repository is cs2.session_repository
+    assert cs2.session_repository is cs3.session_repository
+    assert cs1.session_repository is app_factory.session_repository
+
+    # Create a session via cs1
+    session_id = uuid.uuid4()
+    await cs1.create_session(session_id, "user@example.com")
+
+    # Verify it's accessible from all instances
+    assert await cs1.session_repository.get(session_id) is not None
+    assert await cs2.session_repository.get(session_id) is not None
+    assert await cs3.session_repository.get(session_id) is not None
+
+    # Verify they all return the same session object
+    s1 = await cs1.session_repository.get(session_id)
+    s2 = await cs2.session_repository.get(session_id)
+    s3 = await cs3.session_repository.get(session_id)
+    assert s1 is s2 is s3

atlas/tests/test_system_prompt_loading.py

@@ -0,0 +1,181 @@
+import tempfile
+import uuid
+from pathlib import Path
+
+import pytest
+
+from atlas.application.chat.preprocessors.message_builder import MessageBuilder
+from atlas.domain.messages.models import Message, MessageRole
+from atlas.domain.sessions.models import Session
+from atlas.modules.config import ConfigManager
+from atlas.modules.prompts.prompt_provider import PromptProvider
+
+
+@pytest.mark.asyncio
+async def test_prompt_provider_loads_system_prompt(tmp_path):
+    """Test that PromptProvider correctly loads and formats system_prompt.md"""
+    # Create a temporary system prompt file
+    prompts_dir = tmp_path / "prompts"
+    prompts_dir.mkdir()
+    system_prompt_file = prompts_dir / "system_prompt.md"
+    system_prompt_content = "You are a helpful assistant for user {user_email}."
+    system_prompt_file.write_text(system_prompt_content)
+
+    # Create a config manager with custom prompt base path
+    config_manager = ConfigManager()
+    config_manager.app_settings.prompt_base_path = str(prompts_dir)
+    config_manager.app_settings.system_prompt_filename = "system_prompt.md"
+
+    # Create prompt provider
+    prompt_provider = PromptProvider(config_manager)
+
+    # Test loading system prompt
+    result = prompt_provider.get_system_prompt(user_email="test@example.com")
+
+    assert result is not None
+    assert "test@example.com" in result
+    assert "helpful assistant" in result
+
+
+@pytest.mark.asyncio
+async def test_prompt_provider_handles_missing_system_prompt():
+    """Test that PromptProvider returns None when system_prompt.md is missing"""
+    # Create a config manager pointing to non-existent directory
+    config_manager = ConfigManager()
+    config_manager.app_settings.prompt_base_path = "/nonexistent/path"
+    config_manager.app_settings.system_prompt_filename = "system_prompt.md"
+
+    # Create prompt provider
+    prompt_provider = PromptProvider(config_manager)
+
+    # Test loading system prompt
+    result = prompt_provider.get_system_prompt(user_email="test@example.com")
+
+    assert result is None
+
+
+@pytest.mark.asyncio
+async def test_message_builder_includes_system_prompt(tmp_path):
+    """Test that MessageBuilder includes system prompt in messages"""
+    # Create a temporary system prompt file
+    prompts_dir = tmp_path / "prompts"
+    prompts_dir.mkdir()
+    system_prompt_file = prompts_dir / "system_prompt.md"
+    system_prompt_content = "You are a helpful assistant for user {user_email}."
+    system_prompt_file.write_text(system_prompt_content)
+
+    # Create a config manager with custom prompt base path
+    config_manager = ConfigManager()
+    config_manager.app_settings.prompt_base_path = str(prompts_dir)
+    config_manager.app_settings.system_prompt_filename = "system_prompt.md"
+
+    # Create prompt provider and message builder
+    prompt_provider = PromptProvider(config_manager)
+    message_builder = MessageBuilder(prompt_provider=prompt_provider)
+
+    # Create a session with some history
+    session = Session(user_email="test@example.com")
+    session.history.add_message(Message(role=MessageRole.USER, content="Hello"))
+
+    # Build messages
+    messages = await message_builder.build_messages(
+        session=session,
+        include_files_manifest=False,
+        include_system_prompt=True,
+    )
+
+    # Verify system prompt is first message
+    assert len(messages) >= 2  # system prompt + user message
+    assert messages[0]["role"] == "system"
+    assert "helpful assistant" in messages[0]["content"]
+    assert "test@example.com" in messages[0]["content"]
+
+    # Verify user message is second
+    assert messages[1]["role"] == "user"
+    assert messages[1]["content"] == "Hello"
+
+
+@pytest.mark.asyncio
+async def test_message_builder_without_system_prompt(tmp_path):
+    """Test that MessageBuilder works without system prompt when disabled"""
+    # Create prompt provider without system prompt file
+    config_manager = ConfigManager()
+    config_manager.app_settings.prompt_base_path = "/nonexistent"
+    prompt_provider = PromptProvider(config_manager)
+    message_builder = MessageBuilder(prompt_provider=prompt_provider)
+
+    # Create a session with some history
+    session = Session(user_email="test@example.com")
+    session.history.add_message(Message(role=MessageRole.USER, content="Hello"))
+
+    # Build messages with system prompt disabled
+    messages = await message_builder.build_messages(
+        session=session,
+        include_files_manifest=False,
+        include_system_prompt=False,
+    )
+
+    # Verify no system prompt
+    assert len(messages) == 1
+    assert messages[0]["role"] == "user"
+    assert messages[0]["content"] == "Hello"
+
+
+@pytest.mark.asyncio
+async def test_system_prompt_sent_to_llm():
+    """Test that system prompt is sent to LLM in chat flow"""
+    # Create a temporary directory for prompts
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        prompts_dir = Path(tmp_dir) / "prompts"
+        prompts_dir.mkdir()
+        system_prompt_file = prompts_dir / "system_prompt.md"
+        system_prompt_content = "You are a helpful AI assistant for user {user_email}."
+        system_prompt_file.write_text(system_prompt_content)
+
+        # Create config manager
+        config_manager = ConfigManager()
+        config_manager.app_settings.prompt_base_path = str(prompts_dir)
+        config_manager.app_settings.system_prompt_filename = "system_prompt.md"
+
+        # Capture messages sent to LLM
+        captured = {}
+
+        class DummyLLM:
+            async def call_plain(self, model_name, messages, temperature=0.7):
+                captured["messages"] = messages
+                return "Hello! I'm here to help."
+
+        # Create chat service
+        from atlas.application.chat.service import ChatService
+
+        chat_service = ChatService(
+            llm=DummyLLM(),
+            tool_manager=None,
+            connection=None,
+            config_manager=config_manager,
+            file_manager=None,
+        )
+
+        # Create session and send message
+        session_id = uuid.uuid4()
+        await chat_service.handle_chat_message(
+            session_id=session_id,
+            content="Hello",
+            model="test-model",
+            user_email="tester@example.com",
+            selected_tools=None,
+            selected_prompts=None,
+            selected_data_sources=None,
+            only_rag=False,
+            tool_choice_required=False,
+            agent_mode=False,
+            temperature=0.7,
+        )
+
+        # Verify system prompt was sent to LLM
+        msgs = captured.get("messages")
+        assert msgs, "LLM was not called or messages not captured"
+        assert len(msgs) >= 2  # system prompt + user message
+        assert msgs[0]["role"] == "system", f"Expected first message to be system, got: {msgs[0]}"
+        assert "helpful AI assistant" in msgs[0]["content"]
+        assert "tester@example.com" in msgs[0]["content"]