626 lines
20 KiB
Python
626 lines
20 KiB
Python
"""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()
|