372 lines
13 KiB
Python
372 lines
13 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
from cryptography.fernet import Fernet
|
||
|
|
|
||
|
|
from yuxi.channels.auth.backoff import BackoffConfig, ExponentialBackoff
|
||
|
|
from yuxi.channels.auth.health_monitor import AuthHealth, AuthHealthMonitor
|
||
|
|
from yuxi.channels.auth.secret_manager import SecretManager, SecretSource
|
||
|
|
from yuxi.channels.auth.security_policy import (
|
||
|
|
DmPolicy,
|
||
|
|
GroupPolicy,
|
||
|
|
SecurityPolicy,
|
||
|
|
SecurityPolicyEngine,
|
||
|
|
)
|
||
|
|
from yuxi.channels.auth.sensitive_filter import SensitiveDataFilter, redact_sensitive
|
||
|
|
from yuxi.channels.auth.ssrf_guard import (
|
||
|
|
is_private_url,
|
||
|
|
sanitize_transport_kwargs,
|
||
|
|
validate_url,
|
||
|
|
)
|
||
|
|
from yuxi.channels.auth.token_manager import TokenState, UnifiedTokenManager
|
||
|
|
from yuxi.channels.auth.token_providers import (
|
||
|
|
OAuth2TokenProvider,
|
||
|
|
StaticTokenProvider,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class TestSecretManager:
|
||
|
|
def test_encrypt_decrypt_with_key(self):
|
||
|
|
key = Fernet.generate_key().decode()
|
||
|
|
sm = SecretManager(fernet_key=key)
|
||
|
|
original = "my-super-secret-token"
|
||
|
|
encrypted = sm.encrypt_secret(original)
|
||
|
|
assert encrypted != original
|
||
|
|
decrypted = sm.decrypt_secret(encrypted)
|
||
|
|
assert decrypted == original
|
||
|
|
|
||
|
|
def test_encrypt_without_key(self):
|
||
|
|
sm = SecretManager()
|
||
|
|
original = "my-super-secret-token"
|
||
|
|
encrypted = sm.encrypt_secret(original)
|
||
|
|
assert encrypted == original
|
||
|
|
|
||
|
|
def test_redact_short_value(self):
|
||
|
|
result = SecretManager.redact("abc")
|
||
|
|
assert result == "***"
|
||
|
|
|
||
|
|
def test_redact_long_value(self):
|
||
|
|
result = SecretManager.redact("abcdefghijklmnopqrstuvwxyz")
|
||
|
|
assert result.startswith("abcd")
|
||
|
|
assert result.endswith("wxyz")
|
||
|
|
assert "****" in result
|
||
|
|
|
||
|
|
def test_redact_none(self):
|
||
|
|
result = SecretManager.redact(None)
|
||
|
|
assert result == "***"
|
||
|
|
|
||
|
|
def test_is_sensitive_key(self):
|
||
|
|
assert SecretManager.is_sensitive_key("bot_token")
|
||
|
|
assert SecretManager.is_sensitive_key("api_secret")
|
||
|
|
assert SecretManager.is_sensitive_key("password")
|
||
|
|
assert not SecretManager.is_sensitive_key("username")
|
||
|
|
assert not SecretManager.is_sensitive_key("channel_name")
|
||
|
|
|
||
|
|
def test_redact_config(self):
|
||
|
|
config = {
|
||
|
|
"bot_token": "secret123456789",
|
||
|
|
"username": "mybot",
|
||
|
|
"api_secret": "shhh",
|
||
|
|
"nested": {"password": "nested_secret", "name": "test"},
|
||
|
|
}
|
||
|
|
redacted = SecretManager.redact_config(config)
|
||
|
|
assert redacted["bot_token"] != "secret123456789"
|
||
|
|
assert "***" in redacted["bot_token"]
|
||
|
|
assert redacted["username"] == "mybot"
|
||
|
|
assert redacted["api_secret"] != "shhh"
|
||
|
|
assert redacted["nested"]["password"] != "nested_secret"
|
||
|
|
assert redacted["nested"]["name"] == "test"
|
||
|
|
|
||
|
|
def test_is_encryption_available(self):
|
||
|
|
key = Fernet.generate_key().decode()
|
||
|
|
sm = SecretManager(fernet_key=key)
|
||
|
|
assert sm.is_encryption_available
|
||
|
|
|
||
|
|
sm2 = SecretManager()
|
||
|
|
assert not sm2.is_encryption_available
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_resolve_secret_from_config(self):
|
||
|
|
sm = SecretManager()
|
||
|
|
result = await sm.resolve_secret(
|
||
|
|
"bot_token",
|
||
|
|
sources=[SecretSource.CONFIG],
|
||
|
|
config={"bot_token": "token123"},
|
||
|
|
)
|
||
|
|
assert result == "token123"
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_resolve_secret_with_fallback(self):
|
||
|
|
sm = SecretManager()
|
||
|
|
result = await sm.resolve_secret(
|
||
|
|
"missing_key",
|
||
|
|
sources=[SecretSource.CONFIG],
|
||
|
|
config={"bot_token": "token123"},
|
||
|
|
)
|
||
|
|
assert result is None
|
||
|
|
|
||
|
|
|
||
|
|
class TestBackoff:
|
||
|
|
def test_reset(self):
|
||
|
|
backoff = ExponentialBackoff()
|
||
|
|
backoff.mark_attempt()
|
||
|
|
backoff.mark_attempt()
|
||
|
|
assert backoff.attempt == 2
|
||
|
|
backoff.reset()
|
||
|
|
assert backoff.attempt == 0
|
||
|
|
|
||
|
|
def test_next_delay_increases(self):
|
||
|
|
config = BackoffConfig(base_seconds=1.0, multiplier=2.0, max_seconds=60.0)
|
||
|
|
backoff = ExponentialBackoff(config)
|
||
|
|
delays = []
|
||
|
|
for _ in range(5):
|
||
|
|
delays.append(backoff.next_delay)
|
||
|
|
backoff.mark_attempt()
|
||
|
|
assert delays[0] == 0
|
||
|
|
assert delays[1] > 0
|
||
|
|
assert delays[1] < delays[3]
|
||
|
|
|
||
|
|
def test_max_delay(self):
|
||
|
|
config = BackoffConfig(base_seconds=1.0, multiplier=2.0, max_seconds=5.0)
|
||
|
|
backoff = ExponentialBackoff(config)
|
||
|
|
for _ in range(20):
|
||
|
|
backoff.mark_attempt()
|
||
|
|
delay = backoff.next_delay
|
||
|
|
assert delay <= 5.0 * (1 + config.jitter_pct)
|
||
|
|
|
||
|
|
|
||
|
|
class TestStaticTokenProvider:
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_get_token(self):
|
||
|
|
provider = StaticTokenProvider("static-token", scopes=["read", "write"])
|
||
|
|
state = await provider.get_token()
|
||
|
|
assert state.token == "static-token"
|
||
|
|
assert state.scopes == ["read", "write"]
|
||
|
|
assert state.expires_at is None
|
||
|
|
assert not state.is_expired
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_validate_token(self):
|
||
|
|
provider = StaticTokenProvider("valid-token")
|
||
|
|
state = await provider.get_token()
|
||
|
|
assert await provider.validate_token(state)
|
||
|
|
|
||
|
|
|
||
|
|
class TestTokenManager:
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_register_and_get_token(self):
|
||
|
|
tm = UnifiedTokenManager()
|
||
|
|
provider = StaticTokenProvider("test-token")
|
||
|
|
tm.register_provider("test_channel", provider)
|
||
|
|
|
||
|
|
state = await tm.get_token("test_channel")
|
||
|
|
assert state.token == "test-token"
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_unregistered_provider_raises(self):
|
||
|
|
tm = UnifiedTokenManager()
|
||
|
|
with pytest.raises(ValueError, match="No token provider"):
|
||
|
|
await tm.get_token("nonexistent")
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_invalidate_token(self):
|
||
|
|
tm = UnifiedTokenManager()
|
||
|
|
provider = StaticTokenProvider("test-token")
|
||
|
|
tm.register_provider("test_channel", provider)
|
||
|
|
|
||
|
|
await tm.get_token("test_channel")
|
||
|
|
await tm.invalidate_token("test_channel")
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_unregister_provider(self):
|
||
|
|
tm = UnifiedTokenManager()
|
||
|
|
provider = StaticTokenProvider("test-token")
|
||
|
|
tm.register_provider("test_channel", provider)
|
||
|
|
tm.unregister_provider("test_channel")
|
||
|
|
|
||
|
|
with pytest.raises(ValueError, match="No token provider"):
|
||
|
|
await tm.get_token("test_channel")
|
||
|
|
|
||
|
|
|
||
|
|
class TestAuthHealthMonitor:
|
||
|
|
def test_update_health(self):
|
||
|
|
monitor = AuthHealthMonitor()
|
||
|
|
health = AuthHealth(status="healthy", message="OK")
|
||
|
|
monitor.update_health("ch1", health)
|
||
|
|
|
||
|
|
retrieved = monitor.get_health("ch1")
|
||
|
|
assert retrieved.status == "healthy"
|
||
|
|
assert retrieved.is_healthy
|
||
|
|
|
||
|
|
def test_update_unhealthy_triggers_callback(self):
|
||
|
|
monitor = AuthHealthMonitor()
|
||
|
|
monitor.update_health("ch1", AuthHealth(status="healthy"))
|
||
|
|
|
||
|
|
notified = []
|
||
|
|
monitor.on_auth_failure(lambda cid, h: notified.append((cid, h.status)))
|
||
|
|
|
||
|
|
monitor.update_health("ch1", AuthHealth(status="unhealthy"))
|
||
|
|
|
||
|
|
def test_get_all_health(self):
|
||
|
|
monitor = AuthHealthMonitor()
|
||
|
|
monitor.update_health("ch1", AuthHealth(status="healthy"))
|
||
|
|
monitor.update_health("ch2", AuthHealth(status="degraded", message="slow"))
|
||
|
|
|
||
|
|
all_health = monitor.get_all_health()
|
||
|
|
assert len(all_health) == 2
|
||
|
|
assert all_health["ch1"].is_healthy
|
||
|
|
assert not all_health["ch2"].is_healthy
|
||
|
|
|
||
|
|
|
||
|
|
class TestSecurityPolicy:
|
||
|
|
def test_open_dm(self):
|
||
|
|
policy = SecurityPolicy(dm_policy=DmPolicy.OPEN)
|
||
|
|
assert policy.check_dm("any_user")
|
||
|
|
|
||
|
|
def test_disabled_dm(self):
|
||
|
|
policy = SecurityPolicy(dm_policy=DmPolicy.DISABLED)
|
||
|
|
assert not policy.check_dm("any_user")
|
||
|
|
|
||
|
|
def test_allowlist_dm(self):
|
||
|
|
policy = SecurityPolicy(dm_policy=DmPolicy.ALLOWLIST, allowlist={"user1", "user2"})
|
||
|
|
assert policy.check_dm("user1")
|
||
|
|
assert not policy.check_dm("user3")
|
||
|
|
|
||
|
|
def test_open_group(self):
|
||
|
|
policy = SecurityPolicy(group_policy=GroupPolicy.OPEN)
|
||
|
|
assert policy.check_group()
|
||
|
|
|
||
|
|
def test_disabled_group(self):
|
||
|
|
policy = SecurityPolicy(group_policy=GroupPolicy.DISABLED)
|
||
|
|
assert not policy.check_group()
|
||
|
|
|
||
|
|
def test_from_config(self):
|
||
|
|
config = {
|
||
|
|
"dm_policy": "allowlist",
|
||
|
|
"group_policy": "allowlist",
|
||
|
|
"allowlist": ["alice", "bob"],
|
||
|
|
}
|
||
|
|
policy = SecurityPolicy.from_config(config)
|
||
|
|
assert policy.dm_policy == DmPolicy.ALLOWLIST
|
||
|
|
assert policy.group_policy == GroupPolicy.ALLOWLIST
|
||
|
|
assert policy.is_allowlisted("alice")
|
||
|
|
assert not policy.is_allowlisted("charlie")
|
||
|
|
|
||
|
|
def test_add_remove_allowlist(self):
|
||
|
|
policy = SecurityPolicy(dm_policy=DmPolicy.ALLOWLIST)
|
||
|
|
policy.add_to_allowlist("new_user")
|
||
|
|
assert policy.is_allowlisted("new_user")
|
||
|
|
policy.remove_from_allowlist("new_user")
|
||
|
|
assert not policy.is_allowlisted("new_user")
|
||
|
|
|
||
|
|
def test_blocklist(self):
|
||
|
|
policy = SecurityPolicy(blocklist={"spammer"})
|
||
|
|
assert policy.is_blocklisted("spammer")
|
||
|
|
assert not policy.is_blocklisted("good_user")
|
||
|
|
|
||
|
|
|
||
|
|
class TestSecurityPolicyEngine:
|
||
|
|
def test_register_and_resolve(self):
|
||
|
|
engine = SecurityPolicyEngine()
|
||
|
|
policy = SecurityPolicy(dm_policy=DmPolicy.ALLOWLIST, allowlist={"user1"})
|
||
|
|
engine.register_policy("ch1", policy)
|
||
|
|
|
||
|
|
assert engine.resolve_dm_policy("ch1") == DmPolicy.ALLOWLIST
|
||
|
|
assert engine.check_allowlist("ch1", "user1")
|
||
|
|
assert not engine.check_allowlist("ch1", "user2")
|
||
|
|
|
||
|
|
def test_default_policy(self):
|
||
|
|
engine = SecurityPolicyEngine()
|
||
|
|
assert engine.resolve_dm_policy("unknown") == DmPolicy.OPEN
|
||
|
|
|
||
|
|
|
||
|
|
class TestSensitiveDataFilter:
|
||
|
|
def test_redact_dict(self):
|
||
|
|
data = {
|
||
|
|
"bot_token": "secret123",
|
||
|
|
"channel_name": "mychannel",
|
||
|
|
"api_secret": "shhh",
|
||
|
|
}
|
||
|
|
redacted = SensitiveDataFilter.redact_dict(data)
|
||
|
|
assert redacted["bot_token"] != "secret123"
|
||
|
|
assert redacted["channel_name"] == "mychannel"
|
||
|
|
assert redacted["api_secret"] != "shhh"
|
||
|
|
|
||
|
|
def test_redact_nested_dict(self):
|
||
|
|
data = {"channels": {"telegram": {"bot_token": "tg_secret", "enabled": True}}}
|
||
|
|
redacted = SensitiveDataFilter.redact_dict(data)
|
||
|
|
assert redacted["channels"]["telegram"]["bot_token"] != "tg_secret"
|
||
|
|
assert redacted["channels"]["telegram"]["enabled"] is True
|
||
|
|
|
||
|
|
def test_is_sensitive(self):
|
||
|
|
assert SensitiveDataFilter._is_sensitive("bot_token")
|
||
|
|
assert SensitiveDataFilter._is_sensitive("APP_SECRET")
|
||
|
|
assert not SensitiveDataFilter._is_sensitive("channel_id")
|
||
|
|
|
||
|
|
def test_redact_for_logging(self):
|
||
|
|
msg = "Starting bot with token=abc123def456ghi789 secret=shhh123"
|
||
|
|
redacted = SensitiveDataFilter.redact_for_logging(msg)
|
||
|
|
assert "abc123" not in redacted
|
||
|
|
assert "shhh123" not in redacted
|
||
|
|
assert "***" in redacted
|
||
|
|
|
||
|
|
def test_redact_sensitive_helper(self):
|
||
|
|
assert redact_sensitive("mysecret", key="bot_token") != "mysecret"
|
||
|
|
assert redact_sensitive("myname", key="channel_name") == "myname"
|
||
|
|
|
||
|
|
|
||
|
|
class TestSSRFGuard:
|
||
|
|
def test_localhost_is_private(self):
|
||
|
|
assert is_private_url("http://localhost:8080/api")
|
||
|
|
assert is_private_url("http://127.0.0.1:8080/api")
|
||
|
|
|
||
|
|
def test_private_ip_is_private(self):
|
||
|
|
assert is_private_url("http://192.168.1.1/api")
|
||
|
|
assert is_private_url("http://10.0.0.1/api")
|
||
|
|
|
||
|
|
def test_public_ip_is_not_private(self):
|
||
|
|
assert not is_private_url("http://8.8.8.8/api")
|
||
|
|
assert not is_private_url("http://1.1.1.1/api")
|
||
|
|
|
||
|
|
def test_sanitize_transport_kwargs(self):
|
||
|
|
kwargs = {"agent": "custom", "cert": "fake.crt", "headers": {"X-Key": "val"}}
|
||
|
|
cleaned = sanitize_transport_kwargs(kwargs)
|
||
|
|
assert "agent" not in cleaned
|
||
|
|
assert "cert" not in cleaned
|
||
|
|
assert "headers" in cleaned
|
||
|
|
|
||
|
|
def test_validate_url_with_private(self):
|
||
|
|
with pytest.raises(ValueError, match="private"):
|
||
|
|
validate_url("http://192.168.1.1/api")
|
||
|
|
|
||
|
|
def test_validate_url_with_public(self):
|
||
|
|
result = validate_url("http://8.8.8.8/api")
|
||
|
|
assert result == "http://8.8.8.8/api"
|
||
|
|
|
||
|
|
|
||
|
|
class TestTokenState:
|
||
|
|
def test_not_expired_when_no_expiry(self):
|
||
|
|
state = TokenState(token="test")
|
||
|
|
assert not state.is_expired
|
||
|
|
|
||
|
|
def test_expired_when_past_expiry(self):
|
||
|
|
import time
|
||
|
|
|
||
|
|
state = TokenState(token="test", expires_at=time.monotonic() - 3600)
|
||
|
|
assert state.is_expired
|
||
|
|
|
||
|
|
def test_not_expired_before_expiry(self):
|
||
|
|
import time
|
||
|
|
|
||
|
|
state = TokenState(token="test", expires_at=time.monotonic() + 3600)
|
||
|
|
assert not state.is_expired
|
||
|
|
|
||
|
|
def test_remaining_seconds(self):
|
||
|
|
import time
|
||
|
|
|
||
|
|
state = TokenState(token="test", expires_at=time.monotonic() + 600)
|
||
|
|
remaining = state.remaining_seconds
|
||
|
|
assert remaining is not None
|
||
|
|
assert 0 < remaining < 600
|