knowledge-adapter-service/tests/test_oss_signer.py

from typing import ClassVar
from urllib.parse import parse_qs, unquote, urlparse

import pytest
from alibabacloud_credentials.models import Config, CredentialModel
from pydantic import SecretStr

import knowledge_adapter.oss_signer as oss_signer_module
from knowledge_adapter.asset_redirect import AssetSignerError
from knowledge_adapter.oss_signer import (
    AlibabaOssV4Signer,
    create_ecs_ram_role_credential_client,
)


class RecordingCredentialClient:
    def __init__(self) -> None:
        self.calls = 0

    def get_credential(self) -> CredentialModel:
        self.calls += 1
        return CredentialModel(
            access_key_id="test-access-key-id",
            access_key_secret="test-access-key-secret",
            security_token="test-security-token",
            type="ecs_ram_role",
        )


async def test_oss_signer_generates_v4_temporary_get_url() -> None:
    credential_client = RecordingCredentialClient()
    signer = AlibabaOssV4Signer(
        credential_client=credential_client,
        region="cn-shenzhen",
        bucket="private-assets",
    )

    signed_url = await signer.sign_get(
        object_key="weknora/10008/assets/image.png",
        expires_seconds=600,
    )

    parsed = urlparse(signed_url)
    query = parse_qs(parsed.query)
    assert parsed.scheme == "https"
    assert parsed.hostname == "private-assets.oss-cn-shenzhen.aliyuncs.com"
    assert unquote(parsed.path) == "/weknora/10008/assets/image.png"
    assert query["x-oss-signature-version"] == ["OSS4-HMAC-SHA256"]
    assert query["x-oss-expires"] == ["600"]
    assert query["x-oss-security-token"] == ["test-security-token"]
    assert credential_client.calls == 1


@pytest.mark.parametrize(
    ("object_key", "expires_seconds"),
    [
        ("../other-tenant/secret.png", 600),
        ("weknora/10008/assets/image.png", 601),
        ("weknora/10008/assets/image.png", 60),
    ],
)
async def test_oss_signer_rejects_unsafe_key_or_ttl_before_credentials(
    object_key: str,
    expires_seconds: int,
) -> None:
    credential_client = RecordingCredentialClient()
    signer = AlibabaOssV4Signer(
        credential_client=credential_client,
        region="cn-shenzhen",
        bucket="private-assets",
    )

    with pytest.raises(AssetSignerError, match=r"^asset signing failed$"):
        await signer.sign_get(
            object_key=object_key,
            expires_seconds=expires_seconds,
        )

    assert credential_client.calls == 0


class UnsafeBucket:
    def sign_url(
        self,
        method: str,
        object_key: str,
        expires_seconds: int,
    ) -> str:
        return "http://127.0.0.1/internal?x-oss-signature-version=OSS4-HMAC-SHA256"


async def test_oss_signer_rejects_untrusted_signed_url() -> None:
    signer = AlibabaOssV4Signer(
        credential_client=RecordingCredentialClient(),
        region="cn-shenzhen",
        bucket="private-assets",
    )
    signer._bucket = UnsafeBucket()

    with pytest.raises(AssetSignerError, match=r"^asset signing failed$"):
        await signer.sign_get(
            object_key="weknora/10008/assets/image.png",
            expires_seconds=600,
        )


class RecordingCredentialsClientFactory:
    configs: ClassVar[list[Config]] = []

    def __init__(self, config: Config) -> None:
        self.configs.append(config)


async def test_ecs_ram_role_client_uses_imdsv2_without_static_keys(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    RecordingCredentialsClientFactory.configs = []
    monkeypatch.setattr(
        oss_signer_module,
        "CredentialsClient",
        RecordingCredentialsClientFactory,
    )

    create_ecs_ram_role_credential_client("knowledge-adapter-role")

    assert len(RecordingCredentialsClientFactory.configs) == 1
    config = RecordingCredentialsClientFactory.configs[0]
    assert config.type == "ecs_ram_role"
    assert config.role_name == "knowledge-adapter-role"
    assert config.enable_imds_v2 is True
    assert config.disable_imds_v1 is True
    assert config.connect_timeout == 1000
    assert config.timeout == 1000
    assert not config.access_key_id
    assert not config.access_key_secret


def test_static_ak_client_uses_secret_values_without_security_token() -> None:
    client = oss_signer_module.create_static_ak_credential_client(
        SecretStr("test-access-key-id"),
        SecretStr("test-access-key-secret"),
    )

    credential = client.get_credential()

    assert credential.access_key_id == "test-access-key-id"
    assert credential.access_key_secret == "test-access-key-secret"
    assert not credential.security_token
    assert credential.type == "access_key"
    assert "test-access-key-secret" not in repr(client)