DataOps-platform/app/core/data_source/credentials.py

"""AES-GCM codec and platform PostgreSQL repository for data-source secrets."""

import base64
import binascii
import json
from typing import Optional

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

from app.core.common.identifiers import new_governance_uid
from app.core.data_source.errors import DataSourceCredentialUnavailable
from app.core.data_source.models import DataSourceCredential, SealedCredential


class CredentialCodec:
    """Encrypt credentials with data-source identity bound as AES-GCM AAD."""

    def __init__(self, master_key: bytes, key_version: str):
        if len(master_key) != 32:
            raise DataSourceCredentialUnavailable(
                "master key must decode to 32 bytes"
            )
        if not str(key_version).strip():
            raise DataSourceCredentialUnavailable(
                "credential key version is not configured"
            )
        self._master_key = bytes(master_key)
        self.key_version = str(key_version).strip()

    @classmethod
    def from_base64(cls, encoded_key: str, key_version: str):
        if not str(encoded_key or "").strip():
            raise DataSourceCredentialUnavailable(
                "master key is not configured"
            )
        value = str(encoded_key).strip()
        padding = "=" * (-len(value) % 4)
        try:
            decoded = base64.b64decode(
                value + padding,
                altchars=b"-_",
                validate=True,
            )
        except (binascii.Error, ValueError) as exc:
            raise DataSourceCredentialUnavailable(
                "master key is not valid base64"
            ) from exc
        return cls(decoded, key_version)

    @staticmethod
    def _aad(data_source_uid: str, credential_version: int) -> bytes:
        return f"{data_source_uid}:{int(credential_version)}".encode("utf-8")

    def encrypt(
        self,
        data_source_uid: str,
        credential_version: int,
        credential: DataSourceCredential,
    ) -> SealedCredential:
        import os

        nonce = os.urandom(12)
        payload = json.dumps(
            {
                "username": credential.username,
                "password": credential.password,
                "options": dict(credential.options),
            },
            sort_keys=True,
            separators=(",", ":"),
        ).encode("utf-8")
        encrypted_payload = AESGCM(self._master_key).encrypt(
            nonce,
            payload,
            self._aad(data_source_uid, credential_version),
        )
        return SealedCredential(
            id=new_governance_uid(),
            data_source_uid=str(data_source_uid),
            credential_version=int(credential_version),
            encrypted_payload=encrypted_payload,
            nonce=nonce,
            key_version=self.key_version,
        )

    def decrypt(self, sealed: SealedCredential) -> DataSourceCredential:
        if sealed.key_version != self.key_version:
            raise DataSourceCredentialUnavailable(
                "credential key version is unavailable"
            )
        try:
            plaintext = AESGCM(self._master_key).decrypt(
                bytes(sealed.nonce),
                bytes(sealed.encrypted_payload),
                self._aad(
                    sealed.data_source_uid,
                    sealed.credential_version,
                ),
            )
            payload = json.loads(plaintext.decode("utf-8"))
            username = payload["username"]
            password = payload["password"]
            options = payload.get("options", {})
            if (
                not isinstance(username, str)
                or not isinstance(password, str)
                or not isinstance(options, dict)
            ):
                raise ValueError("malformed credential payload")
            return DataSourceCredential(username, password, options)
        except (
            InvalidTag,
            UnicodeDecodeError,
            json.JSONDecodeError,
            KeyError,
            TypeError,
            ValueError,
        ) as exc:
            raise DataSourceCredentialUnavailable(
                "credential cannot be decrypted"
            ) from exc


class DataSourceCredentialRepository:
    """Store immutable encrypted credential versions in the platform DB."""

    def __init__(self, codec: CredentialCodec):
        self.codec = codec

    def create_version(
        self,
        session,
        *,
        data_source_uid: str,
        credential: DataSourceCredential,
        actor_uid: Optional[str],
    ) -> SealedCredential:
        session.execute(
            text("SELECT pg_advisory_xact_lock(hashtext(:key))"),
            {"key": f"datasource-credential:{data_source_uid}"},
        )
        version = int(
            session.execute(
                text(
                    """
                    SELECT COALESCE(MAX(credential_version), 0) + 1
                    FROM public.datasource_credentials
                    WHERE data_source_uid = CAST(:uid AS uuid)
                    """
                ),
                {"uid": data_source_uid},
            ).scalar_one()
        )
        sealed = self.codec.encrypt(data_source_uid, version, credential)
        session.execute(
            text(
                """
                UPDATE public.datasource_credentials
                SET status = 'retired', retired_at = CURRENT_TIMESTAMP
                WHERE data_source_uid = CAST(:uid AS uuid)
                  AND status = 'active'
                """
            ),
            {"uid": data_source_uid},
        )
        session.execute(
            text(
                """
                INSERT INTO public.datasource_credentials (
                    id, data_source_uid, credential_version,
                    encrypted_payload, nonce, key_version, status
                ) VALUES (
                    CAST(:id AS uuid), CAST(:uid AS uuid), :version,
                    :encrypted_payload, :nonce, :key_version, 'active'
                )
                """
            ),
            {
                "id": sealed.id,
                "uid": sealed.data_source_uid,
                "version": sealed.credential_version,
                "encrypted_payload": sealed.encrypted_payload,
                "nonce": sealed.nonce,
                "key_version": sealed.key_version,
            },
        )
        self._audit(
            session,
            data_source_uid=data_source_uid,
            credential_version=version,
            actor_uid=actor_uid,
            event_type="credential_version_created",
            success=True,
            safe_detail="encrypted credential version created",
        )
        return sealed

    def get_active(
        self,
        session,
        data_source_uid: str,
        credential_version: Optional[int] = None,
    ) -> DataSourceCredential:
        version_clause = (
            "AND credential_version = :version"
            if credential_version is not None
            else ""
        )
        parameters = {"uid": data_source_uid}
        if credential_version is not None:
            parameters["version"] = int(credential_version)
        row = (
            session.execute(
                text(
                    f"""
                    SELECT id::text, data_source_uid::text,
                           credential_version, encrypted_payload,
                           nonce, key_version, status
                    FROM public.datasource_credentials
                    WHERE data_source_uid = CAST(:uid AS uuid)
                      AND status = 'active'
                      {version_clause}
                    ORDER BY credential_version DESC
                    LIMIT 1
                    """
                ),
                parameters,
            )
            .mappings()
            .one_or_none()
        )
        if row is None:
            raise DataSourceCredentialUnavailable()
        sealed = SealedCredential(
            id=row["id"],
            data_source_uid=row["data_source_uid"],
            credential_version=int(row["credential_version"]),
            encrypted_payload=bytes(row["encrypted_payload"]),
            nonce=bytes(row["nonce"]),
            key_version=row["key_version"],
            status=row["status"],
        )
        return self.codec.decrypt(sealed)

    def revoke_all(
        self,
        session,
        *,
        data_source_uid: str,
        actor_uid: Optional[str],
    ) -> int:
        result = session.execute(
            text(
                """
                UPDATE public.datasource_credentials
                SET status = 'revoked',
                    retired_at = COALESCE(retired_at, CURRENT_TIMESTAMP)
                WHERE data_source_uid = CAST(:uid AS uuid)
                  AND status <> 'revoked'
                """
            ),
            {"uid": data_source_uid},
        )
        count = int(result.rowcount or 0)
        self._audit(
            session,
            data_source_uid=data_source_uid,
            credential_version=None,
            actor_uid=actor_uid,
            event_type="credentials_revoked",
            success=True,
            safe_detail=f"revoked credential versions: {count}",
        )
        return count

    def compensate_failed_activation(
        self,
        session,
        *,
        data_source_uid: str,
        failed_version: int,
        restore_version: Optional[int],
        actor_uid: Optional[str],
    ) -> None:
        session.execute(
            text("SELECT pg_advisory_xact_lock(hashtext(:key))"),
            {"key": f"datasource-credential:{data_source_uid}"},
        )
        session.execute(
            text(
                """
                UPDATE public.datasource_credentials
                SET status = 'revoked',
                    retired_at = COALESCE(retired_at, CURRENT_TIMESTAMP)
                WHERE data_source_uid = CAST(:uid AS uuid)
                  AND credential_version = :failed_version
                """
            ),
            {
                "uid": data_source_uid,
                "failed_version": int(failed_version),
            },
        )
        if restore_version is not None:
            session.execute(
                text(
                    """
                    UPDATE public.datasource_credentials
                    SET status = 'active', retired_at = NULL
                    WHERE data_source_uid = CAST(:uid AS uuid)
                      AND credential_version = :restore_version
                      AND status = 'retired'
                    """
                ),
                {
                    "uid": data_source_uid,
                    "restore_version": int(restore_version),
                },
            )
        self._audit(
            session,
            data_source_uid=data_source_uid,
            credential_version=failed_version,
            actor_uid=actor_uid,
            event_type="credential_activation_compensated",
            success=True,
            safe_detail="failed version revoked and prior version restored",
        )

    def record_pool_event(
        self,
        session,
        *,
        data_source_uid: str,
        actor_uid: Optional[str],
        event_type: str,
        safe_detail: str,
    ) -> None:
        self._audit(
            session,
            data_source_uid=data_source_uid,
            credential_version=None,
            actor_uid=actor_uid,
            event_type=event_type,
            success=True,
            safe_detail=safe_detail,
        )

    @staticmethod
    def _audit(
        session,
        *,
        data_source_uid: str,
        credential_version: Optional[int],
        actor_uid: Optional[str],
        event_type: str,
        success: bool,
        safe_detail: str,
    ) -> None:
        session.execute(
            text(
                """
                INSERT INTO public.datasource_credential_audit_events (
                    data_source_uid, credential_version, actor_uid,
                    event_type, success, safe_detail
                ) VALUES (
                    CAST(:uid AS uuid), :version, CAST(:actor_uid AS uuid),
                    :event_type, :success, :safe_detail
                )
                """
            ),
            {
                "uid": data_source_uid,
                "version": credential_version,
                "actor_uid": actor_uid,
                "event_type": event_type,
                "success": bool(success),
                "safe_detail": safe_detail[:500],
            },
        )