Harden production workflows and agent admission
This commit is contained in:
Stanislav Rossovskii
354
server/tests/test_agent_admission.py
Normal file
354
server/tests/test_agent_admission.py
Normal file
@@ -0,0 +1,354 @@
|
||||
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]
|
||||
Reference in New Issue
Block a user