o2o-site-AEO/solution/backend/services/social_account_service.py

"""위임 토큰은 암호문만 저장한다."""

import json
from config import social_config as config
import secrets
import hashlib
from datetime import datetime, timedelta, timezone
from uuid import UUID

import httpx
from cryptography.fernet import Fernet, InvalidToken
from sqlalchemy import select, update, text

from common.database.db_session_manager import DB_SESSION_MNG
from common.database.model.models import owner_social_accounts as Account
from services.external.social import adapter, SocialError


def cipher():
    try:
        return Fernet(config.required("SOCIAL_TOKEN_SECRET").encode())
    except (KeyError, ValueError) as ex:
        raise SocialError("SOCIAL_CONNECTION_DISABLED") from ex


def configured(provider=2):
    try:
        cipher()
        return adapter(provider).is_configured()
    except SocialError:
        return False


def encrypt(value):
    return cipher().encrypt(value.encode()).decode()


def decrypt(value):
    try:
        return cipher().decrypt(value.encode()).decode()
    except (InvalidToken, AttributeError) as ex:
        raise SocialError("TOKEN_KEY_CHANGED", reauth=True) from ex


def begin(user_id, provider):
    if not configured(provider):
        raise SocialError("SOCIAL_CONNECTION_DISABLED")
    browser = secrets.token_urlsafe(32)
    verifier = secrets.token_urlsafe(32)
    # 서명만 된 state는 PKCE verifier를 URL로 공개한다.
    state = encrypt(
        json.dumps(
            {
                "u": str(user_id),
                "p": provider,
                "v": verifier,
                "b": hashlib.sha256(browser.encode()).hexdigest(),
            }
        )
    )
    return adapter(provider).authorize_url(state, verifier), browser


async def finish(state, browser, code):
    try:
        data = json.loads(cipher().decrypt(state.encode(), ttl=600))
        if not browser or not secrets.compare_digest(
            data["b"], hashlib.sha256(browser.encode()).hexdigest()
        ):
            raise ValueError()
        user_id = UUID(data["u"])
        provider = int(data["p"])
    except (InvalidToken, ValueError, KeyError, TypeError) as ex:
        raise SocialError("INVALID_OAUTH_STATE") from ex
    async with httpx.AsyncClient(timeout=20) as client:
        token = await adapter(provider).exchange(code, data["v"], client=client)
        profile = await adapter(provider).me(token["access_token"], client=client)

    async def save(s):
        # 업장별 연결 해제로 다른 업장의 토큰이 살아남지 않도록 사람+플랫폼 단위로 잠근다.
        await lock_user(s, user_id, provider)
        await s.execute(
            update(Account)
            .where(
                Account.user_id == user_id,
                Account.provider == provider,
                Account.deleted == False,
            )
            .values(status="revoked", access_token=None, refresh_token=None)
        )  # noqa: E712
        s.add(
            Account(
                user_id=user_id,
                provider=provider,
                provider_user_id=profile["id"],
                handle=profile["handle"],
                profile_url=profile["profile_url"],
                access_token=encrypt(token["access_token"]),
                refresh_token=encrypt(token["refresh_token"])
                if token.get("refresh_token")
                else None,
                scopes=token["scope"].split(),
                status="linked",
                access_expires_at=datetime.now(timezone.utc)
                + timedelta(seconds=int(token["expires_in"])),
            )
        )

    await DB_SESSION_MNG.execute_lambda_write(Account.DBType(), save)


async def lock_user(s, user_id, provider):
    await s.execute(
        text("SELECT pg_advisory_xact_lock(hashtextextended(:key, 0))"),
        {"key": f"social:{user_id}:{provider}"},
    )


async def account(s, user_id, provider):
    return (
        await s.execute(
            select(Account).where(
                Account.user_id == user_id,
                Account.provider == provider,
                Account.deleted == False,
                Account.status.in_(["linked", "needs_reauth"]),
            )
        )
    ).scalar_one_or_none()  # noqa: E712


async def get_usable_token(s, row, client):
    if row.status != "linked":
        raise SocialError("ACCOUNT_NEEDS_REAUTH", reauth=True)
    token = decrypt(row.access_token)
    if row.access_expires_at and row.access_expires_at > datetime.now(
        timezone.utc
    ) + timedelta(minutes=5):
        return token
    result = await adapter(row.provider).refresh(
        decrypt(row.refresh_token) if row.refresh_token else token, client=client
    )
    # 회전 토큰 저장은 한 UPDATE.
    await s.execute(
        update(Account)
        .where(Account.account_id == row.account_id)
        .values(
            access_token=encrypt(result["access_token"]),
            refresh_token=encrypt(result["refresh_token"])
            if result.get("refresh_token")
            else None,
            access_expires_at=datetime.now(timezone.utc)
            + timedelta(seconds=int(result["expires_in"])),
            scopes=result["scope"].split(),
            updated_at=datetime.now(timezone.utc),
            last_error=None,
        )
    )
    return result["access_token"]


async def disconnect(user_id, provider):
    async def run(s):
        await lock_user(s, user_id, provider)
        await s.execute(
            update(Account)
            .where(
                Account.user_id == user_id,
                Account.provider == provider,
                Account.deleted == False,
            )
            .values(
                status="revoked",
                access_token=None,
                refresh_token=None,  # noqa: E712
                updated_at=datetime.now(timezone.utc),
            )
        )

    await DB_SESSION_MNG.execute_lambda_write(Account.DBType(), run)