import pytest from sqlalchemy import func, select from monlet_server.models import Agent, AgentBlacklist, Check from monlet_server.services.detector import _prune_history from monlet_server.settings import reset_settings_cache from ._helpers import agent_headers, make_event, make_heartbeat, register_agent @pytest.mark.asyncio async def test_unknown_heartbeat_creates_pending_agent(app_client, auth_headers, session): r = await app_client.post( "/api/v1/heartbeat", json=make_heartbeat("agent-pending"), headers=auth_headers ) assert r.status_code == 202 agent = ( await session.execute(select(Agent).where(Agent.agent_id == "agent-pending")) ).scalar_one() assert agent.pending_key_hash assert agent.accepted_key_hash is None @pytest.mark.asyncio async def test_pending_events_are_dropped_until_accept(app_client, auth_headers, session): await app_client.post("/api/v1/heartbeat", json=make_heartbeat(), headers=auth_headers) ev = make_event(status="critical", exit_code=2) r = await app_client.post( "/api/v1/events", json={"agent_id": "agent-1", "events": [ev]}, headers=auth_headers ) assert r.status_code == 202 assert r.json() == {"accepted": 0, "deduplicated": 0} n = (await session.execute(select(func.count()).select_from(Check))).scalar_one() assert n == 0 @pytest.mark.asyncio async def test_accept_enables_events(app_client, ui_auth_headers, session): auth_headers = await register_agent(app_client, ui_auth_headers) ev = make_event(status="critical", exit_code=2) r = await app_client.post( "/api/v1/events", json={"agent_id": "agent-1", "events": [ev]}, headers=auth_headers ) assert r.json() == {"accepted": 1, "deduplicated": 0} n = (await session.execute(select(func.count()).select_from(Check))).scalar_one() assert n == 1 @pytest.mark.asyncio async def test_block_pending_key(app_client, auth_headers, ui_auth_headers, session): await app_client.post( "/api/v1/heartbeat", json=make_heartbeat("agent-block"), headers=auth_headers ) r = await app_client.post("/api/v1/agents/agent-block/block", headers=ui_auth_headers) assert r.status_code == 200 assert r.json()["count"] == 1 assert (await session.execute(select(func.count()).select_from(Agent))).scalar_one() == 0 assert ( await session.execute(select(func.count()).select_from(AgentBlacklist)) ).scalar_one() == 1 r = await app_client.post( "/api/v1/heartbeat", json=make_heartbeat("agent-block"), headers=auth_headers ) assert r.status_code == 403 assert r.json()["error"]["code"] == "agent_blocked" @pytest.mark.asyncio async def test_blacklist_list_delete_and_clear(app_client, auth_headers, ui_auth_headers): await app_client.post( "/api/v1/heartbeat", json=make_heartbeat("agent-block-a"), headers=auth_headers ) await app_client.post("/api/v1/agents/agent-block-a/block", headers=ui_auth_headers) other_headers = agent_headers("agent-block-b") await app_client.post( "/api/v1/heartbeat", json=make_heartbeat("agent-block-b"), headers=other_headers, ) await app_client.post("/api/v1/agents/agent-block-b/block", headers=ui_auth_headers) r = await app_client.get("/api/v1/agents/admission/blacklist", headers=ui_auth_headers) assert r.status_code == 200 items = r.json()["items"] assert len(items) == 2 blocked_id = items[0]["id"] r = await app_client.delete( f"/api/v1/agents/admission/blacklist/{blocked_id}", headers=ui_auth_headers, ) assert r.json()["count"] == 1 r = await app_client.delete("/api/v1/agents/admission/blacklist", headers=ui_auth_headers) assert r.json()["count"] == 1 @pytest.mark.asyncio async def test_remove_accepted_agent_reappears_pending(app_client, ui_auth_headers): auth_headers = await register_agent(app_client, ui_auth_headers, "agent-remove") r = await app_client.delete("/api/v1/agents/agent-remove", headers=ui_auth_headers) assert r.status_code == 200 r = await app_client.post( "/api/v1/heartbeat", json=make_heartbeat("agent-remove"), headers=auth_headers ) assert r.status_code == 202 r = await app_client.get("/api/v1/agents/agent-remove", headers=ui_auth_headers) assert r.json()["admission_state"] == "pending" @pytest.mark.asyncio async def test_new_key_same_hostname_is_pending_replacement(app_client, ui_auth_headers, session): old_headers = await register_agent(app_client, ui_auth_headers, "old-id") hb = make_heartbeat("new-id") hb["hostname"] = "old-id" new_headers = agent_headers("new-id") r = await app_client.post("/api/v1/heartbeat", json=hb, headers=new_headers) assert r.status_code == 202 r = await app_client.get("/api/v1/agents/old-id", headers=ui_auth_headers) assert r.json()["admission_state"] == "pending" assert r.json()["rotation_pending"] is True ev = make_event(status="ok") r = await app_client.post( "/api/v1/events", json={"agent_id": "old-id", "events": [ev]}, headers=old_headers ) assert r.json()["accepted"] == 1 r = await app_client.post("/api/v1/agents/old-id/block", headers=ui_auth_headers) assert r.json()["count"] == 1 agent = (await session.execute(select(Agent).where(Agent.agent_id == "old-id"))).scalar_one() assert agent.accepted_key_hash assert agent.pending_key_hash is None @pytest.mark.asyncio async def test_accept_all_skips_rotations_without_explicit_flag( app_client, ui_auth_headers, session ): await register_agent(app_client, ui_auth_headers, "stable-host") hb = make_heartbeat("replacement-id") hb["hostname"] = "stable-host" await app_client.post("/api/v1/heartbeat", json=hb, headers=agent_headers("replacement")) await app_client.post( "/api/v1/heartbeat", json=make_heartbeat("new-host"), headers=agent_headers("new-host") ) r = await app_client.post("/api/v1/agents/admission/accept-all", headers=ui_auth_headers) assert r.status_code == 200 assert r.json()["count"] == 1 stable = ( await session.execute(select(Agent).where(Agent.agent_id == "stable-host")) ).scalar_one() assert stable.pending_key_hash is not None new_host = ( await session.execute(select(Agent).where(Agent.agent_id == "new-host")) ).scalar_one() assert new_host.pending_key_hash is None assert new_host.accepted_key_hash is not None r = await app_client.post( "/api/v1/agents/admission/accept-all?include_rotation=true", headers=ui_auth_headers, ) assert r.status_code == 200 assert r.json()["count"] == 1 await session.refresh(stable) assert stable.pending_key_hash is None @pytest.mark.asyncio async def test_new_key_cannot_hijack_existing_pending_hostname(app_client, auth_headers): await app_client.post( "/api/v1/heartbeat", json=make_heartbeat("pending-host"), headers=auth_headers ) r = await app_client.post( "/api/v1/heartbeat", json=make_heartbeat("pending-host"), headers=agent_headers("other-key"), ) assert r.status_code == 409 assert r.json()["error"]["code"] == "conflict" @pytest.mark.asyncio async def test_new_key_agent_id_collision_is_rejected(app_client, ui_auth_headers): await register_agent(app_client, ui_auth_headers, "agent-a") await register_agent(app_client, ui_auth_headers, "agent-b") hb = make_heartbeat("agent-a") hb["hostname"] = "new-host" r = await app_client.post("/api/v1/heartbeat", json=hb, headers=agent_headers("new-key")) assert r.status_code == 409 assert r.json()["error"]["code"] == "conflict" @pytest.mark.asyncio async def test_accepted_key_cannot_take_another_hostname(app_client, ui_auth_headers): headers = await register_agent(app_client, ui_auth_headers, "agent-a") await register_agent(app_client, ui_auth_headers, "agent-b") hb = make_heartbeat("agent-a") hb["hostname"] = "agent-b" r = await app_client.post("/api/v1/heartbeat", json=hb, headers=headers) assert r.status_code == 409 assert r.json()["error"]["code"] == "conflict" @pytest.mark.asyncio async def test_hostname_replacement_still_wins_over_body_agent_id(app_client, ui_auth_headers): await register_agent(app_client, ui_auth_headers, "agent-a") await register_agent(app_client, ui_auth_headers, "agent-b") hb = make_heartbeat("agent-a") hb["hostname"] = "agent-b" r = await app_client.post("/api/v1/heartbeat", json=hb, headers=agent_headers("new-key")) assert r.status_code == 202 assert (await app_client.get("/api/v1/agents/agent-b", headers=ui_auth_headers)).json()[ "admission_state" ] == "pending" assert (await app_client.get("/api/v1/agents/agent-a", headers=ui_auth_headers)).json()[ "admission_state" ] == "accepted" @pytest.mark.asyncio async def test_remove_pending_replacement_keeps_accepted_agent( app_client, ui_auth_headers, session ): old_headers = await register_agent(app_client, ui_auth_headers, "stable-host") hb = make_heartbeat("replacement-id") hb["hostname"] = "stable-host" r = await app_client.post("/api/v1/heartbeat", json=hb, headers=agent_headers("replacement")) assert r.status_code == 202 r = await app_client.delete("/api/v1/agents/stable-host", headers=ui_auth_headers) assert r.status_code == 200 agent = ( await session.execute(select(Agent).where(Agent.agent_id == "stable-host")) ).scalar_one() assert agent.accepted_key_hash assert agent.pending_key_hash is None ev = make_event(status="ok") r = await app_client.post( "/api/v1/events", json={"agent_id": "stable-host", "events": [ev]}, headers=old_headers ) assert r.json()["accepted"] == 1 @pytest.mark.asyncio async def test_pending_agent_cap_rejects_new_pending(app_client, auth_headers, monkeypatch): monkeypatch.setenv("MONLET_MAX_PENDING_AGENTS", "1") reset_settings_cache() r = await app_client.post("/api/v1/heartbeat", json=make_heartbeat("one"), headers=auth_headers) assert r.status_code == 202 r = await app_client.post( "/api/v1/heartbeat", json=make_heartbeat("two"), headers=agent_headers("two") ) assert r.status_code == 429 assert r.json()["error"]["code"] == "rate_limited" monkeypatch.delenv("MONLET_MAX_PENDING_AGENTS", raising=False) reset_settings_cache() @pytest.mark.asyncio async def test_pending_prune_deletes_pending_only_and_clears_replacement( app_client, ui_auth_headers, session, monkeypatch ): from datetime import UTC, datetime, timedelta await app_client.post( "/api/v1/heartbeat", json=make_heartbeat("pending-only"), headers=agent_headers("p1") ) await register_agent(app_client, ui_auth_headers, "accepted") hb = make_heartbeat("replacement") hb["hostname"] = "accepted" await app_client.post("/api/v1/heartbeat", json=hb, headers=agent_headers("p2")) old = datetime.now(UTC) - timedelta(days=2) rows = ( (await session.execute(select(Agent).where(Agent.pending_key_hash.is_not(None)))) .scalars() .all() ) for row in rows: row.pending_seen_at = old await session.commit() monkeypatch.setenv("MONLET_PENDING_AGENT_TTL_DAYS", "1") reset_settings_cache() await _prune_history(session) await session.commit() assert ( await session.execute(select(Agent).where(Agent.agent_id == "pending-only")) ).scalar_one_or_none() is None accepted = ( await session.execute(select(Agent).where(Agent.agent_id == "accepted")) ).scalar_one() assert accepted.accepted_key_hash assert accepted.pending_key_hash is None monkeypatch.delenv("MONLET_PENDING_AGENT_TTL_DAYS", raising=False) reset_settings_cache() @pytest.mark.asyncio async def test_blacklist_delete_uses_unique_row_id(app_client, ui_auth_headers, session): from datetime import UTC, datetime now = datetime.now(UTC) same_fingerprint = "0123456789abcdef" first_hash = "a" * 64 second_hash = "b" * 64 session.add_all( [ AgentBlacklist( key_hash=first_hash, key_fingerprint=same_fingerprint, agent_id="agent-a", hostname="host-a", last_seen_at=now, ), AgentBlacklist( key_hash=second_hash, key_fingerprint=same_fingerprint, agent_id="agent-b", hostname="host-b", last_seen_at=now, ), ] ) await session.commit() r = await app_client.get("/api/v1/agents/admission/blacklist", headers=ui_auth_headers) assert r.status_code == 200 items = r.json()["items"] assert [i["id"] for i in items] == [first_hash, second_hash] r = await app_client.delete( f"/api/v1/agents/admission/blacklist/{first_hash}", headers=ui_auth_headers, ) assert r.json()["count"] == 1 r = await app_client.get("/api/v1/agents/admission/blacklist", headers=ui_auth_headers) assert [i["id"] for i in r.json()["items"]] == [second_hash]