litellm/tests/litellm_utils_tests/test_cyberark.py

"""
Integration test for CyberArk Conjur Secret Manager.
"""

import os
import sys
import pytest
import yaml
from dotenv import load_dotenv

load_dotenv()

sys.path.insert(0, os.path.abspath("../.."))
from unittest.mock import AsyncMock, MagicMock, patch
from litellm._uuid import uuid

# Set up environment variables for testing
os.environ["CYBERARK_API_KEY"] = "test-cyberark-api-key-909"
os.environ["CYBERARK_API_BASE"] = "http://0.0.0.0:8080"
os.environ["CYBERARK_ACCOUNT"] = "default"
os.environ["CYBERARK_USERNAME"] = "admin"

from litellm.secret_managers.cyberark_secret_manager import CyberArkSecretManager


def create_mock_response(status_code: int, text: str = ""):
    """
    Helper function to create a mock HTTP response.
    """
    mock_response = MagicMock()
    mock_response.status_code = status_code
    mock_response.text = text
    mock_response.raise_for_status = MagicMock()

    if status_code >= 400:
        import httpx

        error = httpx.HTTPStatusError(
            message=f"HTTP {status_code}", request=MagicMock(), response=mock_response
        )
        mock_response.raise_for_status.side_effect = error

    return mock_response


@pytest.mark.asyncio
async def test_cyberark_write_secret_rejects_yaml_injection():
    """
    Regression test: async_write_secret must reject a secret_name that is not
    safe to embed in the Conjur policy body, before any HTTP call is made.
    """
    with patch("litellm.proxy.proxy_server.premium_user", True):
        malicious_secret_name = "foo\n- !grant\n  role: !!admin\n  member: attacker"

        mock_sync_client = MagicMock()
        mock_async_client = AsyncMock()

        with (
            patch(
                "litellm.secret_managers.cyberark_secret_manager._get_httpx_client",
                return_value=mock_sync_client,
            ),
            patch(
                "litellm.secret_managers.cyberark_secret_manager.get_async_httpx_client",
                return_value=mock_async_client,
            ),
        ):
            cyberark_manager = CyberArkSecretManager()

            response = await cyberark_manager.async_write_secret(
                secret_name=malicious_secret_name,
                secret_value="sk-1234",
            )

            assert response["status"] == "error"
            assert "Invalid secret_name" in response["message"]
            # The malicious policy YAML must never reach the wire.
            mock_sync_client.client.post.assert_not_called()
            mock_async_client.post.assert_not_called()


@pytest.mark.parametrize(
    "secret_name",
    [
        "foo: bar",
        "foo # bar",
        "plain-alias",
        "team/user@example.com",
    ],
)
def test_cyberark_ensure_variable_exists_escapes_yaml_metacharacters(secret_name):
    """
    Regression test: _ensure_variable_exists must escape secret_name (not just
    denylist-check it) so the policy body always parses back to exactly one
    '!variable' scalar node holding the untouched secret_name.
    """
    with patch("litellm.proxy.proxy_server.premium_user", True):
        captured = {}

        def _capture_post(url, headers=None, content=None):
            captured["content"] = content
            return create_mock_response(status_code=201, text="")

        mock_sync_client = MagicMock()
        mock_sync_client.client.post.side_effect = _capture_post

        with patch(
            "litellm.secret_managers.cyberark_secret_manager._get_httpx_client",
            return_value=mock_sync_client,
        ):
            cyberark_manager = CyberArkSecretManager()
            cyberark_manager._ensure_variable_exists(secret_name)

        policy_yaml = captured["content"]
        parsed = yaml.compose(policy_yaml)
        assert len(parsed.value) == 1
        node = parsed.value[0]
        assert node.tag == "!variable"
        assert node.value == secret_name


@pytest.mark.asyncio
async def test_cyberark_write_and_read_secret():
    """
    Test writing a secret to CyberArk Conjur and reading it back using mocked HTTP requests.
    """
    with patch("litellm.proxy.proxy_server.premium_user", True):
        # Generate unique secret name and value
        secret_name = f"test-secret-{uuid.uuid4()}"
        secret_value = f"test-value-{uuid.uuid4()}"

        # Mock sync httpx client (for auth, ensure variable exists, sync read)
        # The _get_httpx_client returns an HTTPHandler with a .client property
        mock_sync_client = MagicMock()
        # Auth response - note: the actual client is accessed via .client property
        mock_sync_client.client.post.return_value = create_mock_response(
            status_code=200, text="mock-token"
        )
        # Sync read response
        mock_sync_client.client.get.return_value = create_mock_response(
            status_code=200, text=secret_value
        )

        # Mock async httpx client (for async write)
        mock_async_client = AsyncMock()
        mock_async_client.post.return_value = create_mock_response(
            status_code=201, text=""
        )

        with (
            patch(
                "litellm.secret_managers.cyberark_secret_manager._get_httpx_client",
                return_value=mock_sync_client,
            ),
            patch(
                "litellm.secret_managers.cyberark_secret_manager.get_async_httpx_client",
                return_value=mock_async_client,
            ),
        ):
            # Create CyberArk secret manager instance
            cyberark_manager = CyberArkSecretManager()

            # Write the secret
            write_response = await cyberark_manager.async_write_secret(
                secret_name=secret_name,
                secret_value=secret_value,
            )
            print("write_response=", write_response)

            # Validate write was successful
            assert write_response["status"] == "success"

            # Read the secret back
            read_value = cyberark_manager.sync_read_secret(secret_name=secret_name)
            print("READ VALUE=", read_value)

            # Validate the secret exists and has the correct value
            assert read_value is not None
            assert read_value == secret_value


@pytest.mark.asyncio
async def test_cyberark_rotate_secret():
    """
    Test key rotation in CyberArk Conjur using mocked HTTP requests.

    This test simulates what happens when a virtual key is rotated:
    1. Write initial secret with alias (like sk-1234)
    2. Rotate to new value (like sk-12359)
    3. Verify reading the secret returns the NEW value
    """
    with patch("litellm.proxy.proxy_server.premium_user", True):
        # Simulate initial virtual key creation
        secret_alias = f"test-rotation-key-{uuid.uuid4()}"
        initial_key_value = f"sk-initial-{uuid.uuid4()}"
        rotated_key_value = f"sk-rotated-{uuid.uuid4()}"

        print(f"\n=== Testing Key Rotation ===")
        print(f"Alias: {secret_alias}")
        print(f"Initial value: {initial_key_value}")
        print(f"Rotated value: {rotated_key_value}")

        # Store the current value to simulate actual storage behavior
        current_value = {"value": initial_key_value}

        # Mock sync httpx client (for auth, ensure variable exists, sync reads)
        # The _get_httpx_client returns an HTTPHandler with a .client property
        mock_sync_client = MagicMock()
        # Auth response - note: the actual client is accessed via .client property
        mock_sync_client.client.post.return_value = create_mock_response(
            status_code=200, text="mock-token"
        )

        # Sync reads return the current value from our simulated storage
        def get_mock_sync_read_response(*args, **kwargs):
            return create_mock_response(status_code=200, text=current_value["value"])

        mock_sync_client.client.get.side_effect = get_mock_sync_read_response

        # Mock async httpx client (for async writes and reads)
        mock_async_client = AsyncMock()

        # Async writes update the current value
        async def mock_async_post(*args, **kwargs):
            content = kwargs.get("content", "")
            if content:
                current_value["value"] = content
            return create_mock_response(status_code=201, text="")

        mock_async_client.post.side_effect = mock_async_post

        # Async reads also return the current value
        async def get_mock_async_read_response(*args, **kwargs):
            return create_mock_response(status_code=200, text=current_value["value"])

        mock_async_client.get.side_effect = get_mock_async_read_response

        with (
            patch(
                "litellm.secret_managers.cyberark_secret_manager._get_httpx_client",
                return_value=mock_sync_client,
            ),
            patch(
                "litellm.secret_managers.cyberark_secret_manager.get_async_httpx_client",
                return_value=mock_async_client,
            ),
        ):
            # Create CyberArk secret manager instance
            cyberark_manager = CyberArkSecretManager()

            # Step 1: Write initial secret (simulates key creation)
            write_response = await cyberark_manager.async_write_secret(
                secret_name=secret_alias,
                secret_value=initial_key_value,
            )
            print(f"\n1. Initial write response: {write_response}")
            assert write_response["status"] == "success"

            # Verify initial value was written
            initial_read = cyberark_manager.sync_read_secret(secret_name=secret_alias)
            print(f"2. Initial read value: {initial_read}")
            assert initial_read == initial_key_value

            # Step 2: Rotate the secret (simulates key rotation)
            # In key rotation, we keep the same secret_name but update the value
            rotation_response = await cyberark_manager.async_rotate_secret(
                current_secret_name=secret_alias,
                new_secret_name=secret_alias,
                new_secret_value=rotated_key_value,
            )
            print(f"3. Rotation response: {rotation_response}")
            assert rotation_response["status"] == "success"

            # Clear cache to force a fresh read
            cyberark_manager.cache.flush_cache()

            # Step 3: Verify the secret now returns the NEW value
            rotated_read = cyberark_manager.sync_read_secret(secret_name=secret_alias)
            print(f"4. After rotation, read value: {rotated_read}")

            # This is the key assertion: after rotation, reading should return the NEW value
            assert rotated_read is not None
            assert rotated_read == rotated_key_value
            assert rotated_read != initial_key_value

            print(
                f"\n✅ Rotation successful: {initial_key_value} → {rotated_key_value}"
            )


@pytest.mark.asyncio
async def test_cyberark_rotate_secret_with_new_alias():
    """
    Test key rotation with a new alias using mocked HTTP requests.

    This simulates rotating a key and changing its alias at the same time:
    1. Write secret with alias-v1
    2. Rotate to alias-v2 with new value
    3. Verify alias-v2 has the new value
    4. Verify alias-v1 still exists with old value (CyberArk doesn't delete)
    """
    with patch("litellm.proxy.proxy_server.premium_user", True):
        # Simulate key rotation with alias change
        base_alias = f"test-alias-change-{uuid.uuid4()}"
        old_alias = f"{base_alias}-v1"
        new_alias = f"{base_alias}-v2"
        old_value = f"sk-old-{uuid.uuid4()}"
        new_value = f"sk-new-{uuid.uuid4()}"

        print(f"\n=== Testing Key Rotation with Alias Change ===")
        print(f"Old alias: {old_alias} = {old_value}")
        print(f"New alias: {new_alias} = {new_value}")

        # Store secrets in a dict to simulate actual storage
        secrets_store = {}

        # Mock sync httpx client (for auth, ensure variable exists, sync reads)
        # The _get_httpx_client returns an HTTPHandler with a .client property
        mock_sync_client = MagicMock()
        # Auth response - note: the actual client is accessed via .client property
        mock_sync_client.client.post.return_value = create_mock_response(
            status_code=200, text="mock-token"
        )

        # Mock sync reads to return from our store
        def get_mock_sync_read(*args, **kwargs):
            url = args[0] if args else kwargs.get("url", "")
            # Extract secret name from URL
            for secret_name, secret_val in secrets_store.items():
                if secret_name in url:
                    return create_mock_response(status_code=200, text=secret_val)
            return create_mock_response(status_code=404, text="Not found")

        mock_sync_client.client.get.side_effect = get_mock_sync_read

        # Mock async httpx client (for async writes and reads)
        mock_async_client = AsyncMock()

        # Mock async write to update our store
        async def mock_async_post(*args, **kwargs):
            url = args[0] if args else kwargs.get("url", "")
            content = kwargs.get("content", "")

            # Extract secret name from URL and store the value
            if old_alias in url:
                secrets_store[old_alias] = content
            elif new_alias in url:
                secrets_store[new_alias] = content

            return create_mock_response(status_code=201, text="")

        mock_async_client.post.side_effect = mock_async_post

        # Mock async reads to return from our store
        async def get_mock_async_read(*args, **kwargs):
            url = args[0] if args else kwargs.get("url", "")
            # Extract secret name from URL
            for secret_name, secret_val in secrets_store.items():
                if secret_name in url:
                    return create_mock_response(status_code=200, text=secret_val)
            return create_mock_response(status_code=404, text="Not found")

        mock_async_client.get.side_effect = get_mock_async_read

        with (
            patch(
                "litellm.secret_managers.cyberark_secret_manager._get_httpx_client",
                return_value=mock_sync_client,
            ),
            patch(
                "litellm.secret_managers.cyberark_secret_manager.get_async_httpx_client",
                return_value=mock_async_client,
            ),
        ):
            # Create CyberArk secret manager instance
            cyberark_manager = CyberArkSecretManager()

            # Step 1: Create initial secret with old alias
            write_response = await cyberark_manager.async_write_secret(
                secret_name=old_alias,
                secret_value=old_value,
            )
            print(f"\n1. Initial write: {write_response}")
            assert write_response["status"] == "success"

            # Step 2: Rotate to new alias with new value
            rotation_response = await cyberark_manager.async_rotate_secret(
                current_secret_name=old_alias,
                new_secret_name=new_alias,
                new_secret_value=new_value,
            )
            print(f"2. Rotation response: {rotation_response}")
            assert rotation_response["status"] == "success"

            # Clear cache to force fresh reads
            cyberark_manager.cache.flush_cache()

            # Step 3: Verify new alias has new value
            new_read = cyberark_manager.sync_read_secret(secret_name=new_alias)
            print(f"3. Read new alias: {new_read}")
            assert new_read == new_value

            # Step 4: Verify old alias still exists (CyberArk doesn't delete via API)
            old_read = cyberark_manager.sync_read_secret(secret_name=old_alias)
            print(f"4. Read old alias (should still exist): {old_read}")
            assert old_read == old_value

            print(f"\n✅ Alias rotation successful: {old_alias} → {new_alias}")
            print(f"   Note: Old alias still exists in CyberArk (expected behavior)")