ForcePilot/backend/test/unit/channels/test_nostr_guard_ext.py

from __future__ import annotations

import time

import pytest

from yuxi.channels.adapters.nostr.guard import (
    AuditRecord,
    GuardPolicy,
    NostrGuard,
    RateLimiter,
    SeenTracker,
)


class TestGuardPolicy:
    def test_default_values(self):
        policy = GuardPolicy()
        assert policy.allowed_kinds == {1, 4, 5, 7, 1059}
        assert policy.max_ciphertext_bytes == 50_000
        assert policy.max_plaintext_bytes == 10_000
        assert policy.max_future_skew_sec == 30
        assert policy.rate_limit_window_ms == 10_000
        assert policy.rate_limit_max_per_sender_per_window == 20
        assert policy.rate_limit_max_global_per_window == 200
        assert policy.dm_policy == "pairing"
        assert policy.allow_from == []

    def test_custom_values(self):
        policy = GuardPolicy(
            allowed_kinds={1, 4},
            max_ciphertext_bytes=100,
            max_plaintext_bytes=50,
            max_future_skew_sec=60,
            rate_limit_window_ms=5000,
            rate_limit_max_per_sender_per_window=5,
            rate_limit_max_global_per_window=50,
            dm_policy="open",
            allow_from=["pub_a"],
        )
        assert policy.allowed_kinds == {1, 4}
        assert policy.dm_policy == "open"
        assert policy.allow_from == ["pub_a"]


class TestSeenTrackerExtended:
    def test_mark_seen_twice(self):
        tracker = SeenTracker(max_size=100)
        tracker.mark_seen("event_001")
        tracker.mark_seen("event_001")
        assert len(tracker) == 1

    def test_is_seen_updates_position(self):
        tracker = SeenTracker(max_size=100, ttl_sec=10)
        tracker.mark_seen("event_a")
        tracker.mark_seen("event_b")
        tracker.mark_seen("event_c")
        tracker.is_seen("event_a")
        assert len(tracker) == 3

    def test_ttl_expiry_removes_on_check(self):
        tracker = SeenTracker(max_size=100, ttl_sec=0)
        tracker.mark_seen("event_001")
        time.sleep(0.01)
        assert tracker.is_seen("event_001") is False
        assert len(tracker) == 0

    def test_periodic_eviction(self):
        tracker = SeenTracker(max_size=5, ttl_sec=0)
        for i in range(5):
            tracker.mark_seen(f"event_{i}")
        time.sleep(0.01)
        tracker.mark_seen("event_5")
        assert len(tracker) <= 5

    def test_len_property(self):
        tracker = SeenTracker(max_size=100)
        tracker.mark_seen("e1")
        tracker.mark_seen("e2")
        tracker.mark_seen("e3")
        assert len(tracker) == 3


class TestRateLimiterExtended:
    def test_window_expiry(self):
        limiter = RateLimiter(window_ms=1, max_per_sender=3, max_global=10)
        limiter.check_and_record("pubkey_a")
        limiter.check_and_record("pubkey_a")
        assert limiter.check_and_record("pubkey_a") is True
        time.sleep(0.002)
        assert limiter.check_and_record("pubkey_a") is True

    def test_clear_resets(self):
        limiter = RateLimiter(max_per_sender=3, max_global=10)
        limiter.check_and_record("pubkey_a")
        limiter.check_and_record("pubkey_a")
        limiter.clear()
        assert limiter.check_and_record("pubkey_a") is True

    def test_separate_senders_dont_interfere(self):
        limiter = RateLimiter(max_per_sender=2, max_global=10)
        assert limiter.check_and_record("sender_a") is True
        assert limiter.check_and_record("sender_a") is True
        assert limiter.check_and_record("sender_a") is False
        assert limiter.check_and_record("sender_b") is True


class TestNostrGuardExtended:
    @pytest.fixture
    def guard(self):
        return NostrGuard(own_pubkey="a" * 64)

    def test_reject_missing_event_id(self, guard):
        event = {
            "kind": 1,
            "content": "no id",
            "pubkey": "b" * 64,
            "created_at": int(time.time()),
        }
        reason = guard.check(event)
        assert reason == "missing event id"

    def test_reject_empty_event_id(self, guard):
        event = {
            "id": "",
            "kind": 1,
            "content": "empty id",
            "pubkey": "b" * 64,
            "created_at": int(time.time()),
        }
        reason = guard.check(event)
        assert reason == "missing event id"

    def test_reject_stale_event(self, guard):
        event = {
            "id": "stale_event",
            "kind": 1,
            "content": "old",
            "pubkey": "b" * 64,
            "created_at": 1000,
        }
        reason = guard.check(event, since_ts=999999)
        assert reason is not None
        assert "stale" in reason

    def test_reject_dm_disabled(self, guard):
        guard.policy.dm_policy = "disabled"
        event = {
            "id": "disabled_dm",
            "kind": 1,
            "content": "hello",
            "pubkey": "b" * 64,
            "created_at": int(time.time()),
        }
        reason = guard.check(event)
        assert "disabled" in reason

    def test_reject_whitelist_empty_allow_from(self, guard):
        guard.policy.dm_policy = "allowlist"
        guard.policy.allow_from = []
        event = {
            "id": "whitelist_reject",
            "kind": 1,
            "content": "not allowed",
            "pubkey": "b" * 64,
            "created_at": int(time.time()),
        }
        reason = guard.check(event)
        assert "whitelist" in reason or "白名单" in reason

    def test_reject_not_in_allow_from(self):
        from yuxi.channels.adapters.nostr.guard import NostrGuard, GuardPolicy

        policy = GuardPolicy(allow_from=["c" * 64])
        guard = NostrGuard(own_pubkey="a" * 64, policy=policy)
        event = {
            "id": "not_allowed",
            "kind": 1,
            "content": "hello",
            "pubkey": "b" * 64,
            "created_at": int(time.time()),
        }
        reason = guard.check(event)
        assert "allow_from" in reason

    def test_pass_in_allow_from(self, guard):
        guard.policy.allow_from = ["b" * 64]
        event = {
            "id": "allowed_event",
            "kind": 1,
            "content": "hello",
            "pubkey": "b" * 64,
            "created_at": int(time.time()),
        }
        reason = guard.check(event)
        assert reason is None

    def test_check_plaintext_size_pass(self, guard):
        reason = guard.check_plaintext_size("short text")
        assert reason is None

    def test_check_plaintext_size_reject(self, guard):
        guard.policy.max_plaintext_bytes = 10
        reason = guard.check_plaintext_size("x" * 100)
        assert reason is not None
        assert "plaintext too large" in reason

    def test_done_processing_removes_from_inflight(self, guard):
        event = {
            "id": "inflight_test",
            "kind": 1,
            "content": "hello",
            "pubkey": "b" * 64,
            "created_at": int(time.time()),
        }
        guard.check(event)
        assert len(guard.inflight) == 1
        guard.done_processing("inflight_test")
        assert len(guard.inflight) == 0

    def test_inflight_blocks_concurrent(self, guard):
        event = {
            "id": "concurrent_event",
            "kind": 1,
            "content": "hello",
            "pubkey": "b" * 64,
            "created_at": int(time.time()),
        }
        assert guard.check(event) is None
        reason = guard.check(event)
        assert reason in ("event in-flight (concurrent processing)", "duplicate event")

    def test_own_pubkey_property(self, guard):
        assert guard.own_pubkey == "a" * 64

    def test_inflight_is_frozenset(self, guard):
        assert isinstance(guard.inflight, frozenset)

    def test_audit_log_records_rejections(self, guard):
        event = {
            "id": "audit_event",
            "kind": 999,
            "content": "bad",
            "pubkey": "b" * 64,
            "created_at": int(time.time()),
        }
        guard.check(event)
        audit = guard.get_audit_log()
        assert len(audit) >= 1
        record = audit[0]
        assert "disallowed" in record["reason"]

    def test_audit_log_max_size(self, guard):
        for i in range(1000):
            event = {
                "id": f"audit_{i}",
                "kind": 999,
                "content": "bad",
                "pubkey": "b" * 64,
                "created_at": int(time.time()),
            }
            guard.check(event)
        audit = guard.get_audit_log()
        assert 500 <= len(audit) <= 1000

    def test_clear_audit_log(self, guard):
        event = {
            "id": "clear_test",
            "kind": 999,
            "content": "bad",
            "pubkey": "b" * 64,
            "created_at": int(time.time()),
        }
        guard.check(event)
        assert len(guard.get_audit_log()) > 0
        guard.clear_audit_log()
        assert len(guard.get_audit_log()) == 0

    def test_rate_limit_rejection(self):
        from yuxi.channels.adapters.nostr.guard import (
            NostrGuard,
            GuardPolicy,
            RateLimiter,
        )

        policy = GuardPolicy(rate_limit_max_per_sender_per_window=2, rate_limit_max_global_per_window=100)
        guard = NostrGuard(own_pubkey="a" * 64, policy=policy)
        for i in range(3):
            event = {
                "id": f"rate_event_{i}",
                "kind": 1,
                "content": "hello",
                "pubkey": "b" * 64,
                "created_at": int(time.time()),
            }
            guard.check(event)
            guard.done_processing(f"rate_event_{i}")
        event = {
            "id": "rate_over_limit",
            "kind": 1,
            "content": "hello",
            "pubkey": "b" * 64,
            "created_at": int(time.time()),
        }
        reason = guard.check(event)
        assert reason == "rate limited"