200 lines
6.7 KiB
Python
200 lines
6.7 KiB
Python
"""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 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)
|
|
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)
|
|
|
|
|
|
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"
|