Bryan Van Deusen
114 lines
3.8 KiB
Python
114 lines
3.8 KiB
Python
"""Credential management service.
|
|
|
|
Handles retrieving, decrypting, and writing cookie files for gallery-dl.
|
|
"""
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.config import get_settings
|
|
from app.models.credential import Credential
|
|
from app.utils.encryption import decrypt_data
|
|
from app.utils.cookies import parse_netscape_cookies, write_cookie_file
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class CredentialManager:
|
|
"""Manages platform credentials for gallery-dl authentication."""
|
|
|
|
def __init__(self, session: AsyncSession):
|
|
self.session = session
|
|
self.settings = get_settings()
|
|
|
|
async def get_cookies_path(self, platform: str) -> Optional[str]:
|
|
"""Get or create a cookies file for a platform.
|
|
|
|
If credentials exist for the platform, decrypts them and writes
|
|
to a temporary cookies file that gallery-dl can use.
|
|
|
|
Args:
|
|
platform: Platform name (patreon, subscribestar, etc.)
|
|
|
|
Returns:
|
|
Path to cookies file, or None if no credentials stored
|
|
"""
|
|
result = await self.session.execute(
|
|
select(Credential).where(Credential.platform == platform)
|
|
)
|
|
credential = result.scalar()
|
|
|
|
if not credential:
|
|
logger.debug(f"No credentials stored for platform: {platform}")
|
|
return None
|
|
|
|
try:
|
|
# Decrypt credential data
|
|
decrypted = decrypt_data(credential.data, self.settings.secret_key)
|
|
|
|
# Write to cookies file
|
|
cookies_dir = Path(self.settings.cookies_path)
|
|
cookies_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
cookies_path = cookies_dir / f"{platform}_cookies.txt"
|
|
|
|
# If it's already in Netscape format, write directly
|
|
if decrypted.strip().startswith("#") or "\t" in decrypted:
|
|
cookies_path.write_text(decrypted)
|
|
else:
|
|
# Assume it's JSON and needs conversion
|
|
import json
|
|
try:
|
|
cookies_list = json.loads(decrypted)
|
|
if isinstance(cookies_list, list):
|
|
write_cookie_file(cookies_list, cookies_path)
|
|
else:
|
|
# Single cookie object
|
|
write_cookie_file([cookies_list], cookies_path)
|
|
except json.JSONDecodeError:
|
|
# Not JSON, write as-is
|
|
cookies_path.write_text(decrypted)
|
|
|
|
logger.debug(f"Wrote cookies file for {platform}: {cookies_path}")
|
|
return str(cookies_path)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to prepare cookies for {platform}: {e}")
|
|
return None
|
|
|
|
async def get_token(self, platform: str) -> Optional[str]:
|
|
"""Get a decrypted token for a platform.
|
|
|
|
Used for platforms like Discord that use token auth instead of cookies.
|
|
|
|
Args:
|
|
platform: Platform name
|
|
|
|
Returns:
|
|
Decrypted token string, or None if not stored
|
|
"""
|
|
result = await self.session.execute(
|
|
select(Credential).where(Credential.platform == platform)
|
|
)
|
|
credential = result.scalar()
|
|
|
|
if not credential or credential.credential_type != "token":
|
|
return None
|
|
|
|
try:
|
|
return decrypt_data(credential.data, self.settings.secret_key)
|
|
except Exception as e:
|
|
logger.error(f"Failed to decrypt token for {platform}: {e}")
|
|
return None
|
|
|
|
async def has_credentials(self, platform: str) -> bool:
|
|
"""Check if credentials exist for a platform."""
|
|
result = await self.session.execute(
|
|
select(Credential).where(Credential.platform == platform)
|
|
)
|
|
return result.scalar() is not None
|