minerops/backend/alembic/versions/0017_agent_enrollment.py

"""add secure Agent enrollment and per-Agent credential storage

Revision ID: 0017_agent_enrollment
Revises: 0016_agent_zones
Create Date: 2026-08-26
"""

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql


revision: str = "0017_agent_enrollment"
down_revision: Union[str, None] = "0016_agent_zones"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # Per-Agent credentials are additive. Existing Agents retain NULL credential
    # fields and therefore remain eligible for the temporary legacy shared-secret
    # authentication path until explicitly migrated.
    op.add_column(
        "minerops_site_agents",
        sa.Column("display_name", sa.Text(), nullable=True),
    )
    op.add_column(
        "minerops_site_agents",
        sa.Column("agent_secret_hash", sa.Text(), nullable=True),
    )
    op.add_column(
        "minerops_site_agents",
        sa.Column(
            "credential_created_at",
            sa.DateTime(timezone=True),
            nullable=True,
        ),
    )
    op.add_column(
        "minerops_site_agents",
        sa.Column(
            "credential_revoked_at",
            sa.DateTime(timezone=True),
            nullable=True,
        ),
    )

    op.create_check_constraint(
        "ck_minerops_site_agents_display_name",
        "minerops_site_agents",
        """
        display_name IS NULL
        OR (
            btrim(display_name) <> ''
            AND char_length(display_name) <= 120
        )
        """,
    )

    op.create_check_constraint(
        "ck_minerops_site_agents_secret_hash",
        "minerops_site_agents",
        """
        agent_secret_hash IS NULL
        OR agent_secret_hash ~ '^[0-9a-f]{64}$'
        """,
    )

    op.create_check_constraint(
        "ck_minerops_site_agents_credential_pair",
        "minerops_site_agents",
        """
        (
            agent_secret_hash IS NULL
            AND credential_created_at IS NULL
        )
        OR (
            agent_secret_hash IS NOT NULL
            AND credential_created_at IS NOT NULL
        )
        """,
    )

    op.create_check_constraint(
        "ck_minerops_site_agents_credential_revoked_time",
        "minerops_site_agents",
        """
        credential_revoked_at IS NULL
        OR (
            credential_created_at IS NOT NULL
            AND credential_revoked_at >= credential_created_at
        )
        """,
    )

    op.create_index(
        "uq_minerops_site_agents_agent_secret_hash",
        "minerops_site_agents",
        ["agent_secret_hash"],
        unique=True,
        postgresql_where=sa.text("agent_secret_hash IS NOT NULL"),
    )

    op.create_table(
        "agent_enrollments",
        sa.Column(
            "id",
            postgresql.UUID(as_uuid=True),
            primary_key=True,
            server_default=sa.text("gen_random_uuid()"),
        ),
        sa.Column("token_hash", sa.Text(), nullable=False),
        sa.Column(
            "site_id",
            postgresql.UUID(as_uuid=True),
            nullable=False,
        ),
        sa.Column(
            "agent_zone_id",
            postgresql.UUID(as_uuid=True),
            nullable=False,
        ),
        sa.Column("display_name", sa.Text(), nullable=False),
        sa.Column(
            "created_by_user_id",
            postgresql.UUID(as_uuid=True),
            sa.ForeignKey("users.id", ondelete="SET NULL"),
            nullable=True,
        ),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.text("now()"),
        ),
        sa.Column(
            "expires_at",
            sa.DateTime(timezone=True),
            nullable=False,
        ),
        sa.Column(
            "used_at",
            sa.DateTime(timezone=True),
            nullable=True,
        ),
        sa.Column(
            "revoked_at",
            sa.DateTime(timezone=True),
            nullable=True,
        ),
        sa.Column("enrolled_agent_id", sa.Text(), nullable=True),
        sa.ForeignKeyConstraint(
            ["agent_zone_id", "site_id"],
            ["agent_zones.id", "agent_zones.site_id"],
            name="fk_agent_enrollments_zone_site",
            ondelete="RESTRICT",
        ),
        sa.ForeignKeyConstraint(
            ["enrolled_agent_id"],
            ["minerops_site_agents.agent_id"],
            name="fk_agent_enrollments_agent",
            ondelete="RESTRICT",
        ),
        sa.CheckConstraint(
            "token_hash ~ '^[0-9a-f]{64}$'",
            name="ck_agent_enrollments_token_hash",
        ),
        sa.CheckConstraint(
            "btrim(display_name) <> '' AND char_length(display_name) <= 120",
            name="ck_agent_enrollments_display_name",
        ),
        sa.CheckConstraint(
            "expires_at > created_at",
            name="ck_agent_enrollments_expiry",
        ),
        sa.CheckConstraint(
            "used_at IS NULL OR used_at >= created_at",
            name="ck_agent_enrollments_used_time",
        ),
        sa.CheckConstraint(
            "used_at IS NULL OR used_at <= expires_at",
            name="ck_agent_enrollments_used_before_expiry",
        ),
        sa.CheckConstraint(
            "revoked_at IS NULL OR revoked_at >= created_at",
            name="ck_agent_enrollments_revoked_time",
        ),
        sa.CheckConstraint(
            "NOT (used_at IS NOT NULL AND revoked_at IS NOT NULL)",
            name="ck_agent_enrollments_terminal_state",
        ),
        sa.CheckConstraint(
            """
            (
                used_at IS NULL
                AND enrolled_agent_id IS NULL
            )
            OR (
                used_at IS NOT NULL
                AND enrolled_agent_id IS NOT NULL
            )
            """,
            name="ck_agent_enrollments_agent_when_used",
        ),
    )

    op.create_index(
        "uq_agent_enrollments_token_hash",
        "agent_enrollments",
        ["token_hash"],
        unique=True,
    )

    # PostgreSQL cannot use now() in a partial-index predicate because now() is
    # not immutable. Therefore "open" means unused and unrevoked. Creation code
    # will revoke an old expired/open enrollment before inserting a replacement.
    op.create_index(
        "uq_agent_enrollments_open_zone",
        "agent_enrollments",
        ["agent_zone_id"],
        unique=True,
        postgresql_where=sa.text(
            "used_at IS NULL AND revoked_at IS NULL"
        ),
    )

    op.create_index(
        "uq_agent_enrollments_enrolled_agent",
        "agent_enrollments",
        ["enrolled_agent_id"],
        unique=True,
        postgresql_where=sa.text("enrolled_agent_id IS NOT NULL"),
    )

    op.create_index(
        "idx_agent_enrollments_site_id",
        "agent_enrollments",
        ["site_id"],
    )

    op.create_index(
        "idx_agent_enrollments_created_at",
        "agent_enrollments",
        ["created_at"],
    )


def downgrade() -> None:
    # 0016 has nowhere to represent enrollment history or per-Agent
    # credentials. Refuse before destructive DDL if either has been used.
    connection = op.get_bind()

    enrollment_count = connection.execute(
        sa.text("SELECT COUNT(*) FROM agent_enrollments")
    ).scalar_one()

    credential_count = connection.execute(
        sa.text(
            """
            SELECT COUNT(*)
            FROM minerops_site_agents
            WHERE display_name IS NOT NULL
               OR agent_secret_hash IS NOT NULL
               OR credential_created_at IS NOT NULL
               OR credential_revoked_at IS NOT NULL
            """
        )
    ).scalar_one()

    if enrollment_count:
        raise RuntimeError(
            "Cannot downgrade 0017: Agent enrollment history exists"
        )

    if credential_count:
        raise RuntimeError(
            "Cannot downgrade 0017: per-Agent credentials exist"
        )

    op.drop_index(
        "idx_agent_enrollments_created_at",
        table_name="agent_enrollments",
    )
    op.drop_index(
        "idx_agent_enrollments_site_id",
        table_name="agent_enrollments",
    )
    op.drop_index(
        "uq_agent_enrollments_enrolled_agent",
        table_name="agent_enrollments",
    )
    op.drop_index(
        "uq_agent_enrollments_open_zone",
        table_name="agent_enrollments",
    )
    op.drop_index(
        "uq_agent_enrollments_token_hash",
        table_name="agent_enrollments",
    )
    op.drop_table("agent_enrollments")

    op.drop_index(
        "uq_minerops_site_agents_agent_secret_hash",
        table_name="minerops_site_agents",
    )

    op.drop_constraint(
        "ck_minerops_site_agents_credential_revoked_time",
        "minerops_site_agents",
        type_="check",
    )
    op.drop_constraint(
        "ck_minerops_site_agents_credential_pair",
        "minerops_site_agents",
        type_="check",
    )
    op.drop_constraint(
        "ck_minerops_site_agents_secret_hash",
        "minerops_site_agents",
        type_="check",
    )
    op.drop_constraint(
        "ck_minerops_site_agents_display_name",
        "minerops_site_agents",
        type_="check",
    )

    op.drop_column("minerops_site_agents", "credential_revoked_at")
    op.drop_column("minerops_site_agents", "credential_created_at")
    op.drop_column("minerops_site_agents", "agent_secret_hash")
    op.drop_column("minerops_site_agents", "display_name")