FabledCurator/backend/app/services/credential_service.py

"""FC-3b: CRUD over Credential rows + on-demand cookies-file
materialisation for gallery-dl (consumed by FC-3c).
"""

from __future__ import annotations

import json
import os
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from ..models import Credential
from .credential_crypto import CredentialCrypto
from .platforms import PLATFORMS, auth_type_for


class CredentialServiceError(Exception):
    """Base."""


class UnknownPlatformError(CredentialServiceError):
    pass


class WrongAuthTypeError(CredentialServiceError):
    def __init__(self, expected: str):
        super().__init__(f"wrong credential_type for platform; expected {expected!r}")
        self.expected = expected


class EmptyDataError(CredentialServiceError):
    pass


@dataclass(frozen=True)
class CredentialRecord:
    platform: str
    credential_type: str
    captured_at: str
    expires_at: str | None
    last_verified: str | None

    def to_dict(self) -> dict:
        return {
            "platform": self.platform,
            "credential_type": self.credential_type,
            "captured_at": self.captured_at,
            "expires_at": self.expires_at,
            "last_verified": self.last_verified,
        }


def _to_record(row: Credential) -> CredentialRecord:
    return CredentialRecord(
        platform=row.platform,
        credential_type=row.credential_type,
        captured_at=row.captured_at.isoformat() if row.captured_at else "",
        expires_at=row.expires_at.isoformat() if row.expires_at else None,
        last_verified=row.last_verified.isoformat() if row.last_verified else None,
    )


# Default cookies dir under the image root; tests override.
_DEFAULT_COOKIES_DIR = Path("/images/cookies")


class CredentialService:
    def __init__(
        self,
        session: AsyncSession,
        crypto: CredentialCrypto,
        cookies_dir: Path | None = None,
    ):
        self.session = session
        self.crypto = crypto
        self.cookies_dir = Path(cookies_dir) if cookies_dir else _DEFAULT_COOKIES_DIR

    async def list(self) -> list[CredentialRecord]:
        rows = (await self.session.execute(
            select(Credential).order_by(Credential.platform.asc())
        )).scalars().all()
        return [_to_record(r) for r in rows]

    async def get(self, platform: str) -> CredentialRecord | None:
        row = (await self.session.execute(
            select(Credential).where(Credential.platform == platform)
        )).scalar_one_or_none()
        return _to_record(row) if row else None

    async def upsert(
        self,
        *,
        platform: str,
        credential_type: str,
        data: str,
        expires_at: datetime | None = None,
    ) -> CredentialRecord:
        if platform not in PLATFORMS:
            raise UnknownPlatformError(f"unknown platform: {platform!r}")
        expected = auth_type_for(platform)
        if credential_type != expected:
            raise WrongAuthTypeError(expected=expected or "")
        cleaned = (data or "").strip()
        if not cleaned:
            raise EmptyDataError("data must be a non-empty string")
        blob = self.crypto.encrypt(cleaned)

        row = (await self.session.execute(
            select(Credential).where(Credential.platform == platform)
        )).scalar_one_or_none()
        if row is None:
            row = Credential(
                platform=platform,
                credential_type=credential_type,
                encrypted_blob=blob,
                expires_at=expires_at,
            )
            self.session.add(row)
        else:
            row.credential_type = credential_type
            row.encrypted_blob = blob
            row.expires_at = expires_at
            row.last_verified = None  # invalidate prior verification
        await self.session.commit()
        await self.session.refresh(row)
        return _to_record(row)

    async def delete(self, platform: str) -> None:
        row = (await self.session.execute(
            select(Credential).where(Credential.platform == platform)
        )).scalar_one_or_none()
        if row is None:
            raise LookupError(f"no credential for platform {platform!r}")
        await self.session.delete(row)
        await self.session.commit()

    async def get_cookies_path(self, platform: str) -> Path | None:
        """Decrypt the credential and write a Netscape cookies.txt for
        gallery-dl. Returns None if no credential or wrong kind."""
        row = (await self.session.execute(
            select(Credential).where(Credential.platform == platform)
        )).scalar_one_or_none()
        if row is None or row.credential_type != "cookies":
            return None
        plaintext = self.crypto.decrypt(row.encrypted_blob)
        netscape = _to_netscape(plaintext)
        netscape = _augment_cookies(platform, netscape)
        self.cookies_dir.mkdir(parents=True, exist_ok=True)
        out = self.cookies_dir / f"{platform}_cookies.txt"
        out.write_text(netscape)
        os.chmod(out, 0o600)
        return out

    async def get_token(self, platform: str) -> str | None:
        row = (await self.session.execute(
            select(Credential).where(Credential.platform == platform)
        )).scalar_one_or_none()
        if row is None or row.credential_type != "token":
            return None
        return self.crypto.decrypt(row.encrypted_blob)

    async def mark_verified(self, platform: str) -> datetime | None:
        """Stamp last_verified=now after a successful verify. Returns the
        timestamp, or None if the credential is gone."""
        row = (await self.session.execute(
            select(Credential).where(Credential.platform == platform)
        )).scalar_one_or_none()
        if row is None:
            return None
        ts = datetime.now(UTC)
        row.last_verified = ts
        await self.session.commit()
        return ts


def _augment_cookies(platform: str, netscape: str) -> str:
    """Delegate to the platform's `augment_cookies` hook if one is
    registered (subscribestar, hentaifoundry, etc. — see
    `services/platforms/<name>.py`). No-op when the platform doesn't
    register a hook (Patreon, Discord). Centralizing the
    quirks-per-platform in the platforms package means adding a new
    platform's cookie quirks doesn't require touching this file."""
    info = PLATFORMS.get(platform)
    if info is None or info.augment_cookies is None:
        return netscape
    return info.augment_cookies(netscape)


def _to_netscape(plaintext: str) -> str:
    """Accept either Netscape-format text (the extension's output) or a
    JSON array of cookie dicts (a manual-paste edge case); produce
    Netscape-format text suitable for gallery-dl --cookies."""
    stripped = plaintext.strip()
    if not stripped:
        return ""
    if stripped.startswith("#") or "\t" in stripped:
        return plaintext  # already Netscape
    try:
        loaded = json.loads(stripped)
    except json.JSONDecodeError:
        return plaintext  # write as-is and hope; gallery-dl will complain if invalid
    if isinstance(loaded, dict):
        loaded = [loaded]
    if not isinstance(loaded, list):
        return plaintext
    lines = ["# Netscape HTTP Cookie File"]
    for c in loaded:
        domain = str(c.get("domain", ""))
        if domain and not domain.startswith("."):
            domain = "." + domain
        flag = "TRUE"
        path = str(c.get("path", "/"))
        secure = "TRUE" if c.get("secure", False) else "FALSE"
        expiration = c.get("expiration") or c.get("expirationDate") or 0
        try:
            expiration_str = str(int(float(expiration)))
        except (TypeError, ValueError):
            expiration_str = "0"
        name = str(c.get("name", ""))
        value = str(c.get("value", ""))
        lines.append("\t".join([domain, flag, path, secure, expiration_str, name, value]))
    return "\n".join(lines) + "\n"