Orchestrator/shared/AI-Orchestrator/scripts/pat_manager/rotator.py

"""GitLab Rotator module for the PAT Manager system.

Handles automatic token rotation for GitLab PATs using the
GitLab Personal Access Tokens API.
"""

from __future__ import annotations

import os
from datetime import date

import httpx

from .errors import RotationError
from .models import RotationResult


class GitLabRotator:
    """Rotates GitLab Personal Access Tokens via the GitLab API.

    Uses POST /api/v4/personal_access_tokens/{id}/rotate to obtain a new
    token with a fresh expiry date. Implements retry logic (one retry on
    failure) and skips rotation for already-expired tokens.

    Attributes:
        timeout: HTTP request timeout in seconds.
    """

    def __init__(self, timeout: int = 30) -> None:
        self.timeout = timeout
        self._gitlab_url = os.environ.get("GITLAB_URL", "https://gitlab.com")

    async def rotate(
        self, token_id: str, current_token: str, expiry_date: str | None = None
    ) -> RotationResult:
        """Rotate a GitLab PAT to obtain a new token and expiry date.

        If the token is already expired (based on the provided expiry_date),
        rotation is skipped and an error result is returned indicating manual
        intervention is needed.

        On HTTP error or timeout, retries once. If the retry also fails,
        returns an error RotationResult.

        Args:
            token_id: The GitLab personal access token ID (numeric string).
            current_token: The current token value used for authentication.
            expiry_date: Optional ISO 8601 date string (YYYY-MM-DD) of the
                         token's current expiry. Used to skip expired tokens.

        Returns:
            RotationResult with success=True and new token/expiry on success,
            or success=False with an error message on failure.
        """
        # Skip rotation for expired tokens
        if expiry_date is not None:
            try:
                parsed_expiry = date.fromisoformat(expiry_date)
                if parsed_expiry < date.today():
                    return RotationResult(
                        success=False,
                        new_token=None,
                        new_expiry_date=None,
                        error_message=(
                            f"Token {token_id} is expired (expiry: {expiry_date}). "
                            "Manual intervention required — cannot rotate an expired token."
                        ),
                    )
            except ValueError:
                # Invalid date format — proceed with rotation attempt
                pass

        url = f"{self._gitlab_url}/api/v4/personal_access_tokens/{token_id}/rotate"
        headers = {"PRIVATE-TOKEN": current_token}

        # First attempt
        result = await self._attempt_rotation(url, headers, token_id)
        if result.success:
            return result

        # Retry once on failure
        result = await self._attempt_rotation(url, headers, token_id)
        return result

    async def _attempt_rotation(
        self, url: str, headers: dict[str, str], token_id: str
    ) -> RotationResult:
        """Make a single rotation API call.

        Args:
            url: The full rotation endpoint URL.
            headers: Request headers including PRIVATE-TOKEN.
            token_id: The token ID for error messages.

        Returns:
            RotationResult indicating success or failure.
        """
        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.post(url, headers=headers)
        except httpx.TimeoutException:
            return RotationResult(
                success=False,
                new_token=None,
                new_expiry_date=None,
                error_message=f"Rotation request timed out for token {token_id}",
            )
        except httpx.ConnectError as exc:
            return RotationResult(
                success=False,
                new_token=None,
                new_expiry_date=None,
                error_message=f"Connection error during rotation for token {token_id}: {exc}",
            )
        except httpx.HTTPError as exc:
            return RotationResult(
                success=False,
                new_token=None,
                new_expiry_date=None,
                error_message=f"HTTP error during rotation for token {token_id}: {exc}",
            )

        if not (200 <= response.status_code < 300):
            return RotationResult(
                success=False,
                new_token=None,
                new_expiry_date=None,
                error_message=(
                    f"GitLab API returned HTTP {response.status_code} "
                    f"for token {token_id} rotation"
                ),
            )

        # Parse successful response
        try:
            data = response.json()
            new_token = data["token"]
            new_expiry_date = data["expires_at"]
        except (KeyError, ValueError) as exc:
            return RotationResult(
                success=False,
                new_token=None,
                new_expiry_date=None,
                error_message=(
                    f"Failed to parse rotation response for token {token_id}: {exc}"
                ),
            )

        return RotationResult(
            success=True,
            new_token=new_token,
            new_expiry_date=new_expiry_date,
            error_message=None,
        )