"""Pre-launch hardening: re-encrypt every integration row's secrets
blob with AAD bound to the row id.

Migration 124 introduced the AAD binding in code — but pre-2026-05-06
rows were encrypted without AAD and the new ``decrypt_fields``
silently falls back to the no-AAD code path for them. Those rows
are still vulnerable to the exact ciphertext-swap attack the AAD
binding was meant to defeat (an attacker with DB write but no key
access can copy a low-priv Slack-test integration's
``(secrets_encrypted, secrets_nonce)`` onto a high-priv production
AWS row, then call test_connection — the AWS provider class would
happily use whichever credentials the swapped blob decrypts to).

This migration walks every integration row, decrypts with the legacy
no-AAD path, re-encrypts with AAD = ``f"integration:{row.id}"``, and
writes the new ciphertext + nonce back. Idempotent — re-running on
already-AAD-bound rows is a no-op (the AAD-decrypt path succeeds, no
fallback warning fires, no re-encryption needed).

Migration runs out-of-tx so a single bad row (corrupted blob,
mismatched key, etc.) doesn't roll the whole upgrade back; we log
the failed row and continue, leaving operator with a list of rows
to repair manually.

We deliberately use the same AESGCM helper as the application code
(via the credential_crypto module) so any future format change to the
ciphertext envelope flows through automatically.
"""
from __future__ import annotations

import base64
import binascii
import json
import logging
import os

import sqlalchemy as sa
from cryptography.exceptions import InvalidTag
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

from alembic import op

logger = logging.getLogger(__name__)

revision = "125"
down_revision = "124"
branch_labels = None
depends_on = None


def _table_exists(bind, table: str) -> bool:
    return bool(bind.execute(
        sa.text(
            "SELECT 1 FROM information_schema.tables "
            "WHERE table_schema = DATABASE() AND table_name = :t LIMIT 1",
        ),
        {"t": table},
    ).scalar())


def _get_key() -> bytes | None:
    raw = os.environ.get("CREDENTIAL_ENCRYPTION_KEY") or ""
    raw = raw.strip()
    if not raw:
        return None
    padded = raw + "=" * (4 - len(raw) % 4)
    try:
        return base64.urlsafe_b64decode(padded)
    except (binascii.Error, ValueError):
        return None


def upgrade() -> None:
    """Re-encrypt every integration row's secrets with AAD."""
    bind = op.get_bind()
    if not _table_exists(bind, "integrations"):
        return

    key = _get_key()
    if key is None:
        # No key configured — likely a fresh-deploy migration on a CI
        # box with no real data. Skip cleanly so the migration chain
        # advances; the AAD binding takes effect on next encrypt anyway.
        logger.info(
            "migration 125: CREDENTIAL_ENCRYPTION_KEY not set — "
            "skipping re-encryption pass (chain advances; new writes "
            "will encrypt with AAD).",
        )
        return

    aesgcm = AESGCM(key)

    rows = bind.execute(sa.text(
        "SELECT id, secrets_encrypted, secrets_nonce "
        "FROM integrations "
        "WHERE secrets_encrypted IS NOT NULL "
        "AND secrets_nonce IS NOT NULL"
    )).fetchall()

    upgraded = 0
    skipped_already_bound = 0
    failed: list[tuple[str, str]] = []

    for row in rows:
        rid = str(row.id)
        ct_b64 = row.secrets_encrypted
        nonce_b64 = row.secrets_nonce
        try:
            ct = base64.urlsafe_b64decode(
                ct_b64 + "=" * (4 - len(ct_b64) % 4),
            )
            nonce = base64.urlsafe_b64decode(
                nonce_b64 + "=" * (4 - len(nonce_b64) % 4),
            )
        except (binascii.Error, ValueError) as exc:
            failed.append((rid, f"base64 decode failed: {exc}"))
            continue

        aad = f"integration:{rid}".encode()

        # Try AAD-bound decrypt first. Idempotency: if it works, the
        # row is already on the new format; skip.
        try:
            aesgcm.decrypt(nonce, ct, aad)
        except InvalidTag:
            pass
        else:
            skipped_already_bound += 1
            continue

        # Try legacy no-AAD decrypt. If that works, re-encrypt with AAD.
        try:
            plaintext = aesgcm.decrypt(nonce, ct, None)
        except InvalidTag as exc:
            failed.append((rid, f"both AAD and no-AAD decrypt failed: {exc}"))
            continue

        # Re-encrypt with AAD bound to the row id.
        new_nonce = os.urandom(12)
        try:
            new_ct = aesgcm.encrypt(new_nonce, plaintext, aad)
        except Exception as exc:  # noqa: BLE001
            failed.append((rid, f"re-encrypt failed: {exc}"))
            continue

        new_ct_b64 = base64.urlsafe_b64encode(new_ct).decode()
        new_nonce_b64 = base64.urlsafe_b64encode(new_nonce).decode()
        try:
            bind.execute(
                sa.text(
                    "UPDATE integrations "
                    "SET secrets_encrypted = :ct, "
                    "    secrets_nonce = :nonce "
                    "WHERE id = :id"
                ),
                {"ct": new_ct_b64, "nonce": new_nonce_b64, "id": rid},
            )
            upgraded += 1
        except Exception as exc:  # noqa: BLE001
            failed.append((rid, f"DB update failed: {exc}"))
            continue

        # Sanity-check that the new blob decrypts correctly with AAD.
        try:
            roundtrip = aesgcm.decrypt(new_nonce, new_ct, aad)
            if roundtrip != plaintext:
                failed.append((
                    rid,
                    "post-rotate sanity check: roundtrip plaintext mismatch",
                ))
        except InvalidTag as exc:
            failed.append((rid, f"post-rotate sanity check failed: {exc}"))

    logger.info(
        "migration 125: AAD-rotate complete — upgraded=%d, "
        "already_bound=%d, failed=%d",
        upgraded, skipped_already_bound, len(failed),
    )
    for rid, reason in failed:
        logger.warning(
            "migration 125: integration_id=%s NOT upgraded: %s. "
            "Re-upload credentials via the admin UI to repair.",
            rid, reason,
        )


def downgrade() -> None:
    """No-op. Downgrading would re-encrypt without AAD which weakens
    the security posture; if you really need to roll back, also
    revert migration 124 first and run a manual rotation script."""
    pass