feat(agent): add durable existing-agent credential transitions
This commit is contained in:
lussierj committed
1 parent
abf5654a25
commit
3cfb0d1abe
33 files changed
+3828
-54
No files matched your search
@@ -1,5 +1,6 @@
|
||||
"""PostgreSQL persistence for secure Agent enrollment."""
|
||||
|
||||
import hmac
|
||||
from collections.abc import Callable
|
||||
from uuid import UUID
|
||||
|
||||
@@ -21,6 +22,7 @@ ConnectionFactory = Callable[[], psycopg.Connection]
|
||||
|
||||
TOKEN_HASH_CONSTRAINT = "uq_agent_enrollments_token_hash"
|
||||
OPEN_ZONE_CONSTRAINT = "uq_agent_enrollments_open_zone"
|
||||
OPEN_TARGET_CONSTRAINT = "uq_agent_enrollments_open_target_agent"
|
||||
AGENT_SECRET_HASH_CONSTRAINT = "uq_minerops_site_agents_agent_secret_hash"
|
||||
AGENT_ZONE_CONSTRAINT = "uq_minerops_site_agents_agent_zone_id"
|
||||
AGENT_ID_CONSTRAINT = "minerops_site_agents_pkey"
|
||||
@@ -214,6 +216,184 @@ class AgentEnrollmentRepository:
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
def create_credential_transition(
|
||||
self,
|
||||
*,
|
||||
token_hash: str,
|
||||
target_agent_id: str,
|
||||
actor_user_id: UUID,
|
||||
) -> dict:
|
||||
connection = self._connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT
|
||||
agent.agent_id,
|
||||
agent.site_id,
|
||||
agent.agent_zone_id,
|
||||
agent.display_name,
|
||||
agent.agent_secret_hash,
|
||||
agent.credential_created_at,
|
||||
agent.credential_revoked_at,
|
||||
agent.legacy_auth_disabled_at,
|
||||
site.name AS site_name,
|
||||
zone.name AS agent_zone_name,
|
||||
zone.is_enabled AS agent_zone_enabled,
|
||||
retirement.retired_at
|
||||
FROM minerops_site_agents AS agent
|
||||
JOIN sites AS site ON site.id = agent.site_id
|
||||
JOIN agent_zones AS zone
|
||||
ON zone.id = agent.agent_zone_id
|
||||
AND zone.site_id = agent.site_id
|
||||
LEFT JOIN agent_legacy_auth_retirements AS retirement
|
||||
ON retirement.agent_id = agent.agent_id
|
||||
WHERE agent.agent_id = %s
|
||||
FOR UPDATE OF agent, zone
|
||||
""",
|
||||
(target_agent_id,),
|
||||
)
|
||||
agent = cur.fetchone()
|
||||
|
||||
if not agent:
|
||||
raise AgentEnrollmentNotFound(
|
||||
"Existing Agent with durable Site and Zone routing was not found"
|
||||
)
|
||||
|
||||
if agent["agent_zone_enabled"] is not True:
|
||||
raise AgentEnrollmentConflict("Agent Zone is disabled")
|
||||
|
||||
if any(
|
||||
agent[field] is not None
|
||||
for field in (
|
||||
"agent_secret_hash",
|
||||
"credential_created_at",
|
||||
"credential_revoked_at",
|
||||
"legacy_auth_disabled_at",
|
||||
"retired_at",
|
||||
)
|
||||
):
|
||||
raise AgentEnrollmentConflict(
|
||||
"Existing Agent is not eligible for legacy credential transition"
|
||||
)
|
||||
|
||||
cur.execute("SELECT clock_timestamp() AS operation_now")
|
||||
operation_now = cur.fetchone()["operation_now"]
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE agent_enrollments
|
||||
SET revoked_at = %s
|
||||
WHERE agent_zone_id = %s
|
||||
AND used_at IS NULL
|
||||
AND revoked_at IS NULL
|
||||
RETURNING id
|
||||
""",
|
||||
(operation_now, agent["agent_zone_id"]),
|
||||
)
|
||||
replaced_ids = [str(row["id"]) for row in cur.fetchall()]
|
||||
|
||||
display_name = (
|
||||
str(agent["display_name"] or agent["agent_id"]).strip()[:120]
|
||||
)
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO agent_enrollments (
|
||||
token_hash,
|
||||
site_id,
|
||||
agent_zone_id,
|
||||
display_name,
|
||||
created_by_user_id,
|
||||
created_at,
|
||||
expires_at,
|
||||
target_agent_id
|
||||
)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s + interval '30 minutes', %s)
|
||||
RETURNING
|
||||
id,
|
||||
site_id,
|
||||
agent_zone_id,
|
||||
display_name,
|
||||
created_at,
|
||||
expires_at,
|
||||
target_agent_id
|
||||
""",
|
||||
(
|
||||
token_hash,
|
||||
agent["site_id"],
|
||||
agent["agent_zone_id"],
|
||||
display_name,
|
||||
actor_user_id,
|
||||
operation_now,
|
||||
operation_now,
|
||||
agent["agent_id"],
|
||||
),
|
||||
)
|
||||
transition = dict(cur.fetchone())
|
||||
transition["site_name"] = agent["site_name"]
|
||||
transition["agent_zone_name"] = agent["agent_zone_name"]
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO audit_logs (
|
||||
actor_user_id,
|
||||
action,
|
||||
entity_type,
|
||||
entity_id,
|
||||
new_value_json,
|
||||
created_at
|
||||
)
|
||||
VALUES (
|
||||
%s,
|
||||
'agent.credential_transition.created',
|
||||
'agent_enrollment',
|
||||
%s,
|
||||
%s,
|
||||
%s
|
||||
)
|
||||
""",
|
||||
(
|
||||
actor_user_id,
|
||||
transition["id"],
|
||||
Jsonb(
|
||||
{
|
||||
"target_agent_id": agent["agent_id"],
|
||||
"site_id": str(agent["site_id"]),
|
||||
"agent_zone_id": str(agent["agent_zone_id"]),
|
||||
"expires_at": transition["expires_at"].isoformat(),
|
||||
"replaced_enrollment_ids": replaced_ids,
|
||||
}
|
||||
),
|
||||
operation_now,
|
||||
),
|
||||
)
|
||||
|
||||
connection.commit()
|
||||
return transition
|
||||
|
||||
except UniqueViolation as exc:
|
||||
connection.rollback()
|
||||
constraint = exc.diag.constraint_name
|
||||
|
||||
if constraint == TOKEN_HASH_CONSTRAINT:
|
||||
raise AgentEnrollmentCredentialCollision(
|
||||
"Generated transition token collided"
|
||||
) from exc
|
||||
|
||||
if constraint in {OPEN_ZONE_CONSTRAINT, OPEN_TARGET_CONSTRAINT}:
|
||||
raise AgentEnrollmentConflict(
|
||||
"An active credential transition already exists for this Agent"
|
||||
) from exc
|
||||
|
||||
raise
|
||||
|
||||
except Exception:
|
||||
connection.rollback()
|
||||
raise
|
||||
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
def list_enrollments(
|
||||
self,
|
||||
*,
|
||||
@@ -238,7 +418,13 @@ class AgentEnrollmentRepository:
|
||||
enrollment.expires_at,
|
||||
enrollment.used_at,
|
||||
enrollment.revoked_at,
|
||||
enrollment.target_agent_id,
|
||||
enrollment.enrolled_agent_id,
|
||||
CASE
|
||||
WHEN enrollment.target_agent_id IS NULL
|
||||
THEN 'new_agent'
|
||||
ELSE 'existing_agent_credential_transition'
|
||||
END AS purpose,
|
||||
CASE
|
||||
WHEN enrollment.used_at IS NOT NULL
|
||||
THEN 'used'
|
||||
@@ -292,7 +478,8 @@ class AgentEnrollmentRepository:
|
||||
display_name,
|
||||
expires_at,
|
||||
used_at,
|
||||
revoked_at
|
||||
revoked_at,
|
||||
target_agent_id
|
||||
FROM agent_enrollments
|
||||
WHERE token_hash = %s
|
||||
FOR UPDATE
|
||||
@@ -309,6 +496,7 @@ class AgentEnrollmentRepository:
|
||||
if (
|
||||
enrollment["used_at"] is not None
|
||||
or enrollment["revoked_at"] is not None
|
||||
or enrollment["target_agent_id"] is not None
|
||||
):
|
||||
raise AgentEnrollmentTokenInvalid(
|
||||
"Invalid or expired enrollment token"
|
||||
@@ -371,6 +559,22 @@ class AgentEnrollmentRepository:
|
||||
"Invalid or expired enrollment token"
|
||||
)
|
||||
|
||||
# Retire the identity before storing its credential. Both
|
||||
# writes use this transaction, so a failed issuance cannot
|
||||
# leave a new misleading tombstone. ON CONFLICT preserves an
|
||||
# earlier retirement if this identity has credential history.
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO agent_legacy_auth_retirements (
|
||||
agent_id,
|
||||
retired_at
|
||||
)
|
||||
VALUES (%s, %s)
|
||||
ON CONFLICT (agent_id) DO NOTHING
|
||||
""",
|
||||
(agent_id, operation_now),
|
||||
)
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO minerops_site_agents (
|
||||
@@ -380,7 +584,8 @@ class AgentEnrollmentRepository:
|
||||
site_name,
|
||||
display_name,
|
||||
agent_secret_hash,
|
||||
credential_created_at
|
||||
credential_created_at,
|
||||
legacy_auth_disabled_at
|
||||
)
|
||||
VALUES (
|
||||
%s,
|
||||
@@ -389,6 +594,7 @@ class AgentEnrollmentRepository:
|
||||
%s,
|
||||
%s,
|
||||
%s,
|
||||
%s,
|
||||
%s
|
||||
)
|
||||
""",
|
||||
@@ -400,6 +606,7 @@ class AgentEnrollmentRepository:
|
||||
enrollment["display_name"],
|
||||
agent_secret_hash,
|
||||
operation_now,
|
||||
operation_now,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -504,3 +711,258 @@ class AgentEnrollmentRepository:
|
||||
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
def redeem_credential_transition(
|
||||
self,
|
||||
*,
|
||||
token_hash: str,
|
||||
agent_secret_hash: str,
|
||||
) -> dict:
|
||||
connection = self._connection_factory()
|
||||
try:
|
||||
with connection.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT
|
||||
id,
|
||||
site_id,
|
||||
agent_zone_id,
|
||||
display_name,
|
||||
expires_at,
|
||||
used_at,
|
||||
revoked_at,
|
||||
target_agent_id,
|
||||
enrolled_agent_id
|
||||
FROM agent_enrollments
|
||||
WHERE token_hash = %s
|
||||
FOR UPDATE
|
||||
""",
|
||||
(token_hash,),
|
||||
)
|
||||
enrollment = cur.fetchone()
|
||||
|
||||
if (
|
||||
not enrollment
|
||||
or enrollment["target_agent_id"] is None
|
||||
or enrollment["revoked_at"] is not None
|
||||
):
|
||||
raise AgentEnrollmentTokenInvalid(
|
||||
"Invalid or expired credential transition token"
|
||||
)
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT
|
||||
agent.agent_id,
|
||||
agent.site_id,
|
||||
agent.agent_zone_id,
|
||||
agent.display_name,
|
||||
agent.agent_secret_hash,
|
||||
agent.credential_created_at,
|
||||
agent.credential_revoked_at,
|
||||
agent.legacy_auth_disabled_at,
|
||||
site.name AS site_name,
|
||||
zone.name AS agent_zone_name,
|
||||
zone.is_enabled AS agent_zone_enabled,
|
||||
retirement.retired_at
|
||||
FROM minerops_site_agents AS agent
|
||||
JOIN sites AS site ON site.id = agent.site_id
|
||||
JOIN agent_zones AS zone
|
||||
ON zone.id = agent.agent_zone_id
|
||||
AND zone.site_id = agent.site_id
|
||||
LEFT JOIN agent_legacy_auth_retirements AS retirement
|
||||
ON retirement.agent_id = agent.agent_id
|
||||
WHERE agent.agent_id = %s
|
||||
FOR UPDATE OF agent, zone
|
||||
""",
|
||||
(enrollment["target_agent_id"],),
|
||||
)
|
||||
agent = cur.fetchone()
|
||||
|
||||
if not agent:
|
||||
raise AgentEnrollmentConflict(
|
||||
"Credential transition target no longer exists with durable routing"
|
||||
)
|
||||
|
||||
if (
|
||||
agent["site_id"] != enrollment["site_id"]
|
||||
or agent["agent_zone_id"] != enrollment["agent_zone_id"]
|
||||
):
|
||||
raise AgentEnrollmentConflict(
|
||||
"Existing Agent routing changed after transition authorization"
|
||||
)
|
||||
|
||||
if agent["agent_zone_enabled"] is not True:
|
||||
raise AgentEnrollmentConflict("Agent Zone is disabled")
|
||||
|
||||
def transition_result() -> dict:
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT network::text AS network
|
||||
FROM agent_zone_network_ranges
|
||||
WHERE agent_zone_id = %s
|
||||
ORDER BY network, id
|
||||
""",
|
||||
(enrollment["agent_zone_id"],),
|
||||
)
|
||||
return {
|
||||
"agent_id": agent["agent_id"],
|
||||
"site_id": agent["site_id"],
|
||||
"site_name": agent["site_name"],
|
||||
"agent_zone_id": agent["agent_zone_id"],
|
||||
"agent_zone_name": agent["agent_zone_name"],
|
||||
"display_name": str(
|
||||
agent["display_name"] or enrollment["display_name"]
|
||||
),
|
||||
"ip_ranges": [row["network"] for row in cur.fetchall()],
|
||||
}
|
||||
|
||||
# A retry after an ambiguous network response is successful only
|
||||
# for the exact locally generated credential that won the first
|
||||
# transaction. It never rotates or reveals a credential.
|
||||
if enrollment["used_at"] is not None:
|
||||
stored_hash = str(agent["agent_secret_hash"] or "")
|
||||
if (
|
||||
enrollment["enrolled_agent_id"]
|
||||
== enrollment["target_agent_id"]
|
||||
and stored_hash
|
||||
and hmac.compare_digest(stored_hash, agent_secret_hash)
|
||||
and agent["credential_created_at"] is not None
|
||||
and agent["credential_revoked_at"] is None
|
||||
and agent["legacy_auth_disabled_at"] is not None
|
||||
and agent["retired_at"] is not None
|
||||
):
|
||||
result = transition_result()
|
||||
connection.commit()
|
||||
return result
|
||||
|
||||
raise AgentEnrollmentTokenInvalid(
|
||||
"Credential transition token has already been consumed"
|
||||
)
|
||||
|
||||
cur.execute("SELECT clock_timestamp() AS operation_now")
|
||||
operation_now = cur.fetchone()["operation_now"]
|
||||
|
||||
if enrollment["expires_at"] <= operation_now:
|
||||
raise AgentEnrollmentTokenInvalid(
|
||||
"Invalid or expired credential transition token"
|
||||
)
|
||||
|
||||
if any(
|
||||
agent[field] is not None
|
||||
for field in (
|
||||
"agent_secret_hash",
|
||||
"credential_created_at",
|
||||
"credential_revoked_at",
|
||||
"legacy_auth_disabled_at",
|
||||
"retired_at",
|
||||
)
|
||||
):
|
||||
raise AgentEnrollmentConflict(
|
||||
"Existing Agent is no longer eligible for credential transition"
|
||||
)
|
||||
|
||||
# The durable retirement, Agent credential, enrollment terminal
|
||||
# state, and audit record commit together. Any later failure
|
||||
# rolls all of them back.
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO agent_legacy_auth_retirements (
|
||||
agent_id,
|
||||
retired_at
|
||||
)
|
||||
VALUES (%s, %s)
|
||||
""",
|
||||
(agent["agent_id"], operation_now),
|
||||
)
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE minerops_site_agents
|
||||
SET agent_secret_hash = %s,
|
||||
credential_created_at = %s,
|
||||
legacy_auth_disabled_at = %s
|
||||
WHERE agent_id = %s
|
||||
RETURNING agent_id
|
||||
""",
|
||||
(
|
||||
agent_secret_hash,
|
||||
operation_now,
|
||||
operation_now,
|
||||
agent["agent_id"],
|
||||
),
|
||||
)
|
||||
if not cur.fetchone():
|
||||
raise AgentEnrollmentConflict(
|
||||
"Credential transition target disappeared"
|
||||
)
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE agent_enrollments
|
||||
SET used_at = %s,
|
||||
enrolled_agent_id = target_agent_id
|
||||
WHERE id = %s
|
||||
AND used_at IS NULL
|
||||
AND revoked_at IS NULL
|
||||
RETURNING used_at, enrolled_agent_id
|
||||
""",
|
||||
(operation_now, enrollment["id"]),
|
||||
)
|
||||
completed = cur.fetchone()
|
||||
if not completed:
|
||||
raise AgentEnrollmentConflict(
|
||||
"Credential transition state changed concurrently"
|
||||
)
|
||||
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO audit_logs (
|
||||
actor_user_id,
|
||||
action,
|
||||
entity_type,
|
||||
entity_id,
|
||||
new_value_json,
|
||||
created_at
|
||||
)
|
||||
VALUES (
|
||||
NULL,
|
||||
'agent.credential_transition.redeemed',
|
||||
'agent_enrollment',
|
||||
%s,
|
||||
%s,
|
||||
%s
|
||||
)
|
||||
""",
|
||||
(
|
||||
enrollment["id"],
|
||||
Jsonb(
|
||||
{
|
||||
"agent_id": agent["agent_id"],
|
||||
"site_id": str(agent["site_id"]),
|
||||
"agent_zone_id": str(agent["agent_zone_id"]),
|
||||
"used_at": completed["used_at"].isoformat(),
|
||||
}
|
||||
),
|
||||
operation_now,
|
||||
),
|
||||
)
|
||||
|
||||
result = transition_result()
|
||||
|
||||
connection.commit()
|
||||
return result
|
||||
|
||||
except UniqueViolation as exc:
|
||||
connection.rollback()
|
||||
if exc.diag.constraint_name == AGENT_SECRET_HASH_CONSTRAINT:
|
||||
raise AgentEnrollmentCredentialCollision(
|
||||
"Agent credential is already in use"
|
||||
) from exc
|
||||
raise
|
||||
|
||||
except Exception:
|
||||
connection.rollback()
|
||||
raise
|
||||
|
||||
finally:
|
||||
connection.close()
|
||||
Reference in new issue
Block a user