letsbe-orchestrator/app/routes/registration_tokens.py

"""Registration token management endpoints."""

import hashlib
import uuid
from datetime import timedelta

from fastapi import APIRouter, HTTPException, status
from sqlalchemy import select

from app.db import AsyncSessionDep
from app.dependencies.admin_auth import AdminAuthDep
from app.models.base import utc_now
from app.models.registration_token import RegistrationToken
from app.models.tenant import Tenant
from app.schemas.registration_token import (
    RegistrationTokenCreate,
    RegistrationTokenCreatedResponse,
    RegistrationTokenList,
    RegistrationTokenResponse,
)

router = APIRouter(
    prefix="/tenants/{tenant_id}/registration-tokens",
    tags=["Registration Tokens"],
    dependencies=[AdminAuthDep],
)


# --- Helper functions ---


async def get_tenant_by_id(db: AsyncSessionDep, tenant_id: uuid.UUID) -> Tenant | None:
    """Retrieve a tenant by ID."""
    result = await db.execute(select(Tenant).where(Tenant.id == tenant_id))
    return result.scalar_one_or_none()


async def get_token_by_id(
    db: AsyncSessionDep, tenant_id: uuid.UUID, token_id: uuid.UUID
) -> RegistrationToken | None:
    """Retrieve a registration token by ID, scoped to tenant."""
    result = await db.execute(
        select(RegistrationToken).where(
            RegistrationToken.id == token_id,
            RegistrationToken.tenant_id == tenant_id,
        )
    )
    return result.scalar_one_or_none()


# --- Route handlers ---


@router.post(
    "",
    response_model=RegistrationTokenCreatedResponse,
    status_code=status.HTTP_201_CREATED,
    summary="Create a registration token",
    description="""
Create a new registration token for a tenant.

The token can be used by agents to register with the orchestrator.
The plaintext token is only returned once - store it securely.

**Authentication:** Requires X-Admin-Api-Key header.
""",
)
async def create_registration_token(
    tenant_id: uuid.UUID,
    request: RegistrationTokenCreate,
    db: AsyncSessionDep,
) -> RegistrationTokenCreatedResponse:
    """Create a new registration token for a tenant."""
    # Verify tenant exists
    tenant = await get_tenant_by_id(db, tenant_id)
    if tenant is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Tenant {tenant_id} not found",
        )

    # Generate token (UUID format for uniqueness)
    plaintext_token = str(uuid.uuid4())
    token_hash = hashlib.sha256(plaintext_token.encode()).hexdigest()

    # Calculate expiration if specified
    expires_at = None
    if request.expires_in_hours is not None:
        expires_at = utc_now() + timedelta(hours=request.expires_in_hours)

    # Create token record
    token_record = RegistrationToken(
        tenant_id=tenant_id,
        token_hash=token_hash,
        description=request.description,
        max_uses=request.max_uses,
        expires_at=expires_at,
    )

    db.add(token_record)
    await db.commit()
    await db.refresh(token_record)

    # Return response with plaintext token (only time it's shown)
    return RegistrationTokenCreatedResponse(
        id=token_record.id,
        tenant_id=token_record.tenant_id,
        description=token_record.description,
        max_uses=token_record.max_uses,
        use_count=token_record.use_count,
        expires_at=token_record.expires_at,
        revoked=token_record.revoked,
        created_at=token_record.created_at,
        created_by=token_record.created_by,
        token=plaintext_token,
    )


@router.get(
    "",
    response_model=RegistrationTokenList,
    summary="List registration tokens",
    description="""
List all registration tokens for a tenant.

Note: The plaintext token values are not returned.

**Authentication:** Requires X-Admin-Api-Key header.
""",
)
async def list_registration_tokens(
    tenant_id: uuid.UUID,
    db: AsyncSessionDep,
) -> RegistrationTokenList:
    """List all registration tokens for a tenant."""
    # Verify tenant exists
    tenant = await get_tenant_by_id(db, tenant_id)
    if tenant is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Tenant {tenant_id} not found",
        )

    # Get all tokens for tenant
    result = await db.execute(
        select(RegistrationToken)
        .where(RegistrationToken.tenant_id == tenant_id)
        .order_by(RegistrationToken.created_at.desc())
    )
    tokens = result.scalars().all()

    return RegistrationTokenList(
        tokens=[RegistrationTokenResponse.model_validate(t) for t in tokens],
        total=len(tokens),
    )


@router.get(
    "/{token_id}",
    response_model=RegistrationTokenResponse,
    summary="Get registration token details",
    description="""
Get details of a specific registration token.

Note: The plaintext token value is not returned.

**Authentication:** Requires X-Admin-Api-Key header.
""",
)
async def get_registration_token(
    tenant_id: uuid.UUID,
    token_id: uuid.UUID,
    db: AsyncSessionDep,
) -> RegistrationTokenResponse:
    """Get details of a specific registration token."""
    token = await get_token_by_id(db, tenant_id, token_id)
    if token is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Registration token {token_id} not found",
        )

    return RegistrationTokenResponse.model_validate(token)


@router.delete(
    "/{token_id}",
    status_code=status.HTTP_204_NO_CONTENT,
    summary="Revoke registration token",
    description="""
Revoke a registration token.

Revoked tokens cannot be used for new agent registrations.
Agents that have already registered with this token will continue to work.

**Authentication:** Requires X-Admin-Api-Key header.
""",
)
async def revoke_registration_token(
    tenant_id: uuid.UUID,
    token_id: uuid.UUID,
    db: AsyncSessionDep,
) -> None:
    """Revoke a registration token."""
    token = await get_token_by_id(db, tenant_id, token_id)
    if token is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Registration token {token_id} not found",
        )

    # Mark as revoked
    token.revoked = True
    await db.commit()