feat: add secure agent enrollment
This commit is contained in:
lussierj committed
1 parent
82309db79a
commit
d9f0da694b
26 files changed
+2942
-76
No files matched your search
@@ -0,0 +1,625 @@
|
||||
"""PostgreSQL integration tests for secure Agent enrollment."""
|
||||
|
||||
import unittest
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from uuid import uuid4
|
||||
|
||||
from psycopg.rows import dict_row
|
||||
|
||||
from app.agent_auth.exceptions import AgentAuthenticationFailed
|
||||
from app.agent_auth.repository import AgentAuthRepository
|
||||
from app.agent_auth.service import AgentAuthService, hash_agent_secret
|
||||
from app.agent_enrollment.exceptions import (
|
||||
AgentEnrollmentConflict,
|
||||
AgentEnrollmentNotFound,
|
||||
AgentEnrollmentTokenInvalid,
|
||||
)
|
||||
from app.agent_enrollment.models import (
|
||||
AgentEnrollmentCreateRequest,
|
||||
AgentEnrollRequest,
|
||||
)
|
||||
from app.agent_enrollment.repository import (
|
||||
AgentEnrollmentRepository,
|
||||
default_connection_factory,
|
||||
)
|
||||
from app.agent_enrollment.service import AgentEnrollmentService
|
||||
from app.agent_zones.models import AgentCheckinRequest
|
||||
from app.agent_zones.repository import AgentZoneRepository
|
||||
from app.agent_zones.service import AgentZoneService
|
||||
|
||||
|
||||
class AgentEnrollmentPostgresTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.site_id = uuid4()
|
||||
self.zone_id = uuid4()
|
||||
self.site_name = f"Enrollment Test Site {self.site_id.hex[:8]}"
|
||||
self.zone_name = f"Enrollment Test Zone {self.zone_id.hex[:8]}"
|
||||
|
||||
connection = default_connection_factory()
|
||||
try:
|
||||
with connection.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO sites (id, name)
|
||||
VALUES (%s, %s)
|
||||
""",
|
||||
(self.site_id, self.site_name),
|
||||
)
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO agent_zones (
|
||||
id,
|
||||
site_id,
|
||||
name
|
||||
)
|
||||
VALUES (%s, %s, %s)
|
||||
""",
|
||||
(
|
||||
self.zone_id,
|
||||
self.site_id,
|
||||
self.zone_name,
|
||||
),
|
||||
)
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO agent_zone_network_ranges (
|
||||
agent_zone_id,
|
||||
network
|
||||
)
|
||||
VALUES (%s, '10.250.0.0/24'::cidr)
|
||||
""",
|
||||
(self.zone_id,),
|
||||
)
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
self.service = AgentEnrollmentService(
|
||||
AgentEnrollmentRepository()
|
||||
)
|
||||
|
||||
def tearDown(self):
|
||||
connection = default_connection_factory()
|
||||
try:
|
||||
with connection.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
DELETE FROM audit_logs
|
||||
WHERE entity_type = 'agent_enrollment'
|
||||
AND (
|
||||
new_value_json ->> 'site_id' = %s
|
||||
OR entity_id IN (
|
||||
SELECT id
|
||||
FROM agent_enrollments
|
||||
WHERE site_id = %s
|
||||
)
|
||||
)
|
||||
""",
|
||||
(str(self.site_id), self.site_id),
|
||||
)
|
||||
cur.execute(
|
||||
"""
|
||||
DELETE FROM agent_enrollments
|
||||
WHERE site_id = %s
|
||||
""",
|
||||
(self.site_id,),
|
||||
)
|
||||
cur.execute(
|
||||
"""
|
||||
DELETE FROM minerops_site_agents
|
||||
WHERE site_id = %s
|
||||
""",
|
||||
(self.site_id,),
|
||||
)
|
||||
cur.execute(
|
||||
"""
|
||||
DELETE FROM agent_zone_network_ranges
|
||||
WHERE agent_zone_id = %s
|
||||
""",
|
||||
(self.zone_id,),
|
||||
)
|
||||
cur.execute(
|
||||
"""
|
||||
DELETE FROM agent_zones
|
||||
WHERE id = %s
|
||||
""",
|
||||
(self.zone_id,),
|
||||
)
|
||||
cur.execute(
|
||||
"""
|
||||
DELETE FROM sites
|
||||
WHERE id = %s
|
||||
""",
|
||||
(self.site_id,),
|
||||
)
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
def create_enrollment(self, display_name="Integration Agent"):
|
||||
# created_by_user_id is nullable at the persistence layer.
|
||||
# Router-level tests separately prove that the real admin path supplies
|
||||
# the authenticated user's UUID.
|
||||
return self.service.create_enrollment(
|
||||
AgentEnrollmentCreateRequest(
|
||||
site_id=self.site_id,
|
||||
agent_zone_id=self.zone_id,
|
||||
display_name=display_name,
|
||||
),
|
||||
actor_user_id=None,
|
||||
)
|
||||
|
||||
def test_replacement_revokes_previous_open_enrollment(self):
|
||||
first = self.create_enrollment("First Agent")
|
||||
second = self.create_enrollment("Replacement Agent")
|
||||
|
||||
self.assertNotEqual(
|
||||
first["enrollment_token"],
|
||||
second["enrollment_token"],
|
||||
)
|
||||
|
||||
connection = default_connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT
|
||||
token_hash,
|
||||
used_at,
|
||||
revoked_at
|
||||
FROM agent_enrollments
|
||||
WHERE site_id = %s
|
||||
AND agent_zone_id = %s
|
||||
""",
|
||||
(self.site_id, self.zone_id),
|
||||
)
|
||||
rows = cur.fetchall()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
self.assertEqual(len(rows), 2)
|
||||
|
||||
by_hash = {
|
||||
row["token_hash"]: row
|
||||
for row in rows
|
||||
}
|
||||
|
||||
first_hash = hash_agent_secret(
|
||||
first["enrollment_token"]
|
||||
)
|
||||
second_hash = hash_agent_secret(
|
||||
second["enrollment_token"]
|
||||
)
|
||||
|
||||
self.assertIsNotNone(by_hash[first_hash]["revoked_at"])
|
||||
self.assertIsNone(by_hash[first_hash]["used_at"])
|
||||
|
||||
self.assertIsNone(by_hash[second_hash]["revoked_at"])
|
||||
self.assertIsNone(by_hash[second_hash]["used_at"])
|
||||
|
||||
with self.assertRaises(AgentEnrollmentTokenInvalid):
|
||||
self.service.redeem(
|
||||
AgentEnrollRequest(
|
||||
token=first["enrollment_token"]
|
||||
)
|
||||
)
|
||||
|
||||
listed = self.service.list_enrollments(
|
||||
site_id=self.site_id,
|
||||
agent_zone_id=self.zone_id,
|
||||
)
|
||||
|
||||
self.assertEqual(len(listed), 2)
|
||||
|
||||
for item in listed:
|
||||
self.assertNotIn("token_hash", item)
|
||||
self.assertNotIn("agent_secret_hash", item)
|
||||
self.assertNotIn("enrollment_token", item)
|
||||
self.assertNotIn("agent_secret", item)
|
||||
|
||||
def test_concurrent_creation_leaves_exactly_one_open_enrollment(self):
|
||||
def create_one(number):
|
||||
service = AgentEnrollmentService(
|
||||
AgentEnrollmentRepository()
|
||||
)
|
||||
return service.create_enrollment(
|
||||
AgentEnrollmentCreateRequest(
|
||||
site_id=self.site_id,
|
||||
agent_zone_id=self.zone_id,
|
||||
display_name=f"Concurrent Agent {number}",
|
||||
),
|
||||
actor_user_id=None,
|
||||
)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||
results = list(
|
||||
executor.map(create_one, (1, 2))
|
||||
)
|
||||
|
||||
self.assertEqual(len(results), 2)
|
||||
|
||||
connection = default_connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT
|
||||
count(*) AS total,
|
||||
count(*) FILTER (
|
||||
WHERE used_at IS NULL
|
||||
AND revoked_at IS NULL
|
||||
) AS open_count,
|
||||
count(*) FILTER (
|
||||
WHERE revoked_at IS NOT NULL
|
||||
) AS revoked_count,
|
||||
count(*) FILTER (
|
||||
WHERE revoked_at IS NOT NULL
|
||||
AND revoked_at < created_at
|
||||
) AS invalid_revocation_count
|
||||
FROM agent_enrollments
|
||||
WHERE site_id = %s
|
||||
AND agent_zone_id = %s
|
||||
""",
|
||||
(self.site_id, self.zone_id),
|
||||
)
|
||||
counts = cur.fetchone()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
self.assertEqual(counts["total"], 2)
|
||||
self.assertEqual(counts["open_count"], 1)
|
||||
self.assertEqual(counts["revoked_count"], 1)
|
||||
self.assertEqual(counts["invalid_revocation_count"], 0)
|
||||
|
||||
def test_concurrent_redemption_has_exactly_one_winner(self):
|
||||
enrollment = self.create_enrollment()
|
||||
token = enrollment["enrollment_token"]
|
||||
|
||||
def redeem_once():
|
||||
service = AgentEnrollmentService(
|
||||
AgentEnrollmentRepository()
|
||||
)
|
||||
try:
|
||||
return (
|
||||
"success",
|
||||
service.redeem(
|
||||
AgentEnrollRequest(token=token)
|
||||
),
|
||||
)
|
||||
except AgentEnrollmentTokenInvalid:
|
||||
return ("invalid", None)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||
results = list(
|
||||
executor.map(lambda _: redeem_once(), range(2))
|
||||
)
|
||||
|
||||
successes = [
|
||||
result
|
||||
for status, result in results
|
||||
if status == "success"
|
||||
]
|
||||
failures = [
|
||||
status
|
||||
for status, _ in results
|
||||
if status == "invalid"
|
||||
]
|
||||
|
||||
self.assertEqual(len(successes), 1)
|
||||
self.assertEqual(len(failures), 1)
|
||||
|
||||
credentials = successes[0]
|
||||
agent_id = credentials["agent_id"]
|
||||
agent_secret = credentials["agent_secret"]
|
||||
|
||||
connection = default_connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT
|
||||
agent_id,
|
||||
agent_secret_hash,
|
||||
credential_created_at,
|
||||
credential_revoked_at,
|
||||
last_payload::text AS last_payload
|
||||
FROM minerops_site_agents
|
||||
WHERE site_id = %s
|
||||
AND agent_zone_id = %s
|
||||
""",
|
||||
(self.site_id, self.zone_id),
|
||||
)
|
||||
agents = cur.fetchall()
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT
|
||||
used_at,
|
||||
revoked_at,
|
||||
enrolled_agent_id,
|
||||
token_hash
|
||||
FROM agent_enrollments
|
||||
WHERE id = %s
|
||||
""",
|
||||
(enrollment["id"],),
|
||||
)
|
||||
enrollment_row = cur.fetchone()
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT new_value_json::text AS payload
|
||||
FROM audit_logs
|
||||
WHERE entity_type = 'agent_enrollment'
|
||||
AND entity_id = %s
|
||||
ORDER BY created_at
|
||||
""",
|
||||
(enrollment["id"],),
|
||||
)
|
||||
audit_rows = cur.fetchall()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
self.assertEqual(len(agents), 1)
|
||||
agent = agents[0]
|
||||
|
||||
self.assertEqual(agent["agent_id"], agent_id)
|
||||
self.assertEqual(
|
||||
agent["agent_secret_hash"],
|
||||
hash_agent_secret(agent_secret),
|
||||
)
|
||||
self.assertIsNotNone(agent["credential_created_at"])
|
||||
self.assertIsNone(agent["credential_revoked_at"])
|
||||
|
||||
self.assertIsNotNone(enrollment_row["used_at"])
|
||||
self.assertIsNone(enrollment_row["revoked_at"])
|
||||
self.assertEqual(
|
||||
enrollment_row["enrolled_agent_id"],
|
||||
agent_id,
|
||||
)
|
||||
|
||||
# Neither plaintext credential may persist in operational state.
|
||||
persisted_text = "\n".join(
|
||||
[
|
||||
str(enrollment_row),
|
||||
str(agent),
|
||||
*(row["payload"] for row in audit_rows),
|
||||
]
|
||||
)
|
||||
|
||||
self.assertNotIn(token, persisted_text)
|
||||
self.assertNotIn(agent_secret, persisted_text)
|
||||
self.assertNotIn(agent_secret, agent["last_payload"])
|
||||
|
||||
# The newly issued per-Agent credential must immediately authenticate.
|
||||
auth_service = AgentAuthService(
|
||||
AgentAuthRepository()
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
auth_service.authenticate(
|
||||
agent_id,
|
||||
agent_secret,
|
||||
),
|
||||
agent_id,
|
||||
)
|
||||
|
||||
with self.assertRaises(AgentAuthenticationFailed):
|
||||
auth_service.authenticate(
|
||||
agent_id,
|
||||
"definitely-not-the-issued-secret",
|
||||
)
|
||||
|
||||
# Normal check-in must never mutate the credential.
|
||||
secret_hash_before = agent["agent_secret_hash"]
|
||||
|
||||
AgentZoneService(
|
||||
AgentZoneRepository()
|
||||
).checkin(
|
||||
agent_id,
|
||||
AgentCheckinRequest(
|
||||
site_id=self.site_id,
|
||||
agent_zone_id=self.zone_id,
|
||||
site_name=self.site_name,
|
||||
hostname="integration-test-agent",
|
||||
version="integration-test",
|
||||
),
|
||||
)
|
||||
|
||||
connection = default_connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT
|
||||
agent_secret_hash,
|
||||
credential_created_at,
|
||||
credential_revoked_at
|
||||
FROM minerops_site_agents
|
||||
WHERE agent_id = %s
|
||||
""",
|
||||
(agent_id,),
|
||||
)
|
||||
after_checkin = cur.fetchone()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
self.assertEqual(
|
||||
after_checkin["agent_secret_hash"],
|
||||
secret_hash_before,
|
||||
)
|
||||
self.assertIsNotNone(
|
||||
after_checkin["credential_created_at"]
|
||||
)
|
||||
self.assertIsNone(
|
||||
after_checkin["credential_revoked_at"]
|
||||
)
|
||||
|
||||
def test_expired_token_is_rejected_without_side_effects(self):
|
||||
enrollment = self.create_enrollment()
|
||||
|
||||
connection = default_connection_factory()
|
||||
try:
|
||||
with connection.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE agent_enrollments
|
||||
SET created_at = now() - interval '2 hours',
|
||||
expires_at = now() - interval '1 hour'
|
||||
WHERE id = %s
|
||||
""",
|
||||
(enrollment["id"],),
|
||||
)
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
with self.assertRaises(AgentEnrollmentTokenInvalid):
|
||||
self.service.redeem(
|
||||
AgentEnrollRequest(
|
||||
token=enrollment["enrollment_token"]
|
||||
)
|
||||
)
|
||||
|
||||
connection = default_connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT used_at, revoked_at, enrolled_agent_id
|
||||
FROM agent_enrollments
|
||||
WHERE id = %s
|
||||
""",
|
||||
(enrollment["id"],),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT count(*) AS count
|
||||
FROM minerops_site_agents
|
||||
WHERE agent_zone_id = %s
|
||||
""",
|
||||
(self.zone_id,),
|
||||
)
|
||||
agent_count = cur.fetchone()["count"]
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
self.assertIsNone(row["used_at"])
|
||||
self.assertIsNone(row["revoked_at"])
|
||||
self.assertIsNone(row["enrolled_agent_id"])
|
||||
self.assertEqual(agent_count, 0)
|
||||
|
||||
def test_disabled_and_occupied_zone_do_not_consume_token(self):
|
||||
enrollment = self.create_enrollment()
|
||||
token = enrollment["enrollment_token"]
|
||||
|
||||
connection = default_connection_factory()
|
||||
try:
|
||||
with connection.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE agent_zones
|
||||
SET is_enabled = false
|
||||
WHERE id = %s
|
||||
""",
|
||||
(self.zone_id,),
|
||||
)
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
with self.assertRaises(AgentEnrollmentConflict):
|
||||
self.service.redeem(
|
||||
AgentEnrollRequest(token=token)
|
||||
)
|
||||
|
||||
connection = default_connection_factory()
|
||||
try:
|
||||
with connection.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE agent_zones
|
||||
SET is_enabled = true
|
||||
WHERE id = %s
|
||||
""",
|
||||
(self.zone_id,),
|
||||
)
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO minerops_site_agents (
|
||||
agent_id,
|
||||
site_id,
|
||||
agent_zone_id,
|
||||
site_name
|
||||
)
|
||||
VALUES (%s, %s, %s, %s)
|
||||
""",
|
||||
(
|
||||
f"legacy-{uuid4().hex}",
|
||||
self.site_id,
|
||||
self.zone_id,
|
||||
self.site_name,
|
||||
),
|
||||
)
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
with self.assertRaises(AgentEnrollmentConflict):
|
||||
self.service.redeem(
|
||||
AgentEnrollRequest(token=token)
|
||||
)
|
||||
|
||||
connection = default_connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT used_at, revoked_at, enrolled_agent_id
|
||||
FROM agent_enrollments
|
||||
WHERE id = %s
|
||||
""",
|
||||
(enrollment["id"],),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
self.assertIsNone(row["used_at"])
|
||||
self.assertIsNone(row["revoked_at"])
|
||||
self.assertIsNone(row["enrolled_agent_id"])
|
||||
|
||||
def test_site_zone_mismatch_is_rejected_before_creation(self):
|
||||
wrong_site_id = uuid4()
|
||||
|
||||
with self.assertRaises(AgentEnrollmentNotFound):
|
||||
self.service.create_enrollment(
|
||||
AgentEnrollmentCreateRequest(
|
||||
site_id=wrong_site_id,
|
||||
agent_zone_id=self.zone_id,
|
||||
display_name="Wrong Site Agent",
|
||||
),
|
||||
actor_user_id=None,
|
||||
)
|
||||
|
||||
connection = default_connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT count(*) AS count
|
||||
FROM agent_enrollments
|
||||
WHERE agent_zone_id = %s
|
||||
""",
|
||||
(self.zone_id,),
|
||||
)
|
||||
count = cur.fetchone()["count"]
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
self.assertEqual(count, 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user