ari_assistant/backend/app/connectors/oauth.py

# backend/app/connectors/oauth.py — OAuth2 Authorization Code + PKCE (phase-13, phase-16)
# start(인증 URL) → finish(code→token·프로필조회) → valid_access_token(자동 refresh).
# phase-16: Microsoft(Outlook) provider + 멀티계정(이메일별 고유 계정 id) + MailAccount upsert.
import base64
import hashlib
import os
import secrets
import time
import urllib.parse

import httpx
from sqlmodel import Session

from ..config import get_settings
from ..crypto import decrypt_token, encrypt_token
from ..models import (
    ConnectorAccount,
    ConnectorDomain,
    ConnectorMode,
    ConnState,
    MailAccount,
    OAuthState,
)
from .mail.normalize import mail_account_id_for

PROVIDERS = {
    "google": {  # Gmail + Google Calendar
        "auth": "https://accounts.google.com/o/oauth2/v2/auth",
        "token": "https://oauth2.googleapis.com/token",
        "scopes": {
            "gmail": [
                # modify = 읽기 + 라벨변경(별표/보관/읽음) + 휴지통. send = 발송.
                "https://www.googleapis.com/auth/gmail.modify",
                "https://www.googleapis.com/auth/gmail.send",
            ],
            "google_calendar": ["https://www.googleapis.com/auth/calendar.events"],
        },
    },
    "microsoft": {  # Outlook / Microsoft 365 (Graph). {tenant} 는 호출 시 치환.
        "auth": "https://login.microsoftonline.com/{tenant}/oauth2/v2.0/authorize",
        "token": "https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token",
        "scopes": {
            "outlook": [
                # ReadWrite = 읽기 + 플래그/읽음/이동/삭제. Send = 발송.
                "https://graph.microsoft.com/Mail.ReadWrite",
                "https://graph.microsoft.com/Mail.Send",
                "https://graph.microsoft.com/User.Read",
                "offline_access",
            ],
            "outlook_calendar": [
                # ReadWrite = 일정 읽기 + 생성/수정/삭제.
                "https://graph.microsoft.com/Calendars.ReadWrite",
                "https://graph.microsoft.com/User.Read",
                "offline_access",
            ],
        },
    },
    "notion": {
        "auth": "https://api.notion.com/v1/oauth/authorize",
        "token": "https://api.notion.com/v1/oauth/token",
        "scopes": {"notion": []},
    },
}

# provider → (domain, name, kind, tone) 새 real 계정 메타
_REAL_META = {
    "gmail": ("mail", "Gmail", "Google", "blue"),
    "outlook": ("mail", "Outlook", "Microsoft", "blue"),
    "google_calendar": ("calendar", "Google 캘린더", "캘린더", "blue"),
    "outlook_calendar": ("calendar", "Outlook 캘린더", "캘린더", "blue"),
    "notion": ("knowledge", "Notion", "메모·문서", "violet"),
}


def _family(provider: str) -> str:
    if provider in ("gmail", "google_calendar"):
        return "google"
    if provider in ("outlook", "outlook_calendar"):
        return "microsoft"
    return provider


def _conf(fam: str) -> dict:
    """provider family 설정. microsoft 는 {tenant} 를 런타임 치환."""
    c = PROVIDERS[fam]
    if fam == "microsoft":
        t = get_settings().microsoft_tenant
        return {**c, "auth": c["auth"].format(tenant=t), "token": c["token"].format(tenant=t)}
    return c


def _creds(fam: str) -> tuple[str, str]:
    st = get_settings()
    if fam == "google":
        return st.google_client_id, st.google_client_secret
    if fam == "microsoft":
        return st.microsoft_client_id, st.microsoft_client_secret
    return st.notion_client_id, st.notion_client_secret


def _redirect_uri() -> str:
    st = get_settings()
    return st.oauth_redirect_uri or st.google_redirect_uri


def start_oauth(session: Session, domain: str, provider: str, redirect_after: str = "/life") -> str:
    fam = _family(provider)
    conf = _conf(fam)
    client_id, _ = _creds(fam)
    state = secrets.token_urlsafe(24)
    verifier = base64.urlsafe_b64encode(os.urandom(40)).decode().rstrip("=")
    challenge = (
        base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).decode().rstrip("=")
    )
    session.add(
        OAuthState(
            id=state,
            domain=domain,
            provider=provider,
            code_verifier=verifier,
            redirect_after=redirect_after,
        )
    )
    session.commit()
    scopes = list(conf["scopes"].get(provider, []))
    # 메일 동의에 캘린더 권한도 함께 요청 → 한 번의 동의로 메일+일정 모두 연결.
    if provider == "gmail":
        scopes += PROVIDERS["google"]["scopes"]["google_calendar"]
    if provider == "outlook":
        scopes.append("https://graph.microsoft.com/Calendars.ReadWrite")
    params = {
        "client_id": client_id,
        "redirect_uri": _redirect_uri(),
        "response_type": "code",
        "scope": " ".join(scopes),
        "state": state,
        "access_type": "offline",
        "prompt": "consent",
        "code_challenge": challenge,
        "code_challenge_method": "S256",
    }
    return f"{conf['auth']}?" + urllib.parse.urlencode(params)


def _fetch_identity(fam: str, access_token: str) -> str:
    """연결된 계정의 이메일 주소를 조회(멀티계정 식별용). 실패해도 빈 문자열로 graceful."""
    st = get_settings()
    headers = {"Authorization": f"Bearer {access_token}"}
    try:
        if fam == "google":
            r = httpx.get(
                "https://gmail.googleapis.com/gmail/v1/users/me/profile",
                headers=headers,
                timeout=st.connector_http_timeout,
            )
            r.raise_for_status()
            return r.json().get("emailAddress", "") or ""
        if fam == "microsoft":
            r = httpx.get(
                "https://graph.microsoft.com/v1.0/me",
                headers=headers,
                timeout=st.connector_http_timeout,
            )
            r.raise_for_status()
            j = r.json()
            return j.get("mail") or j.get("userPrincipalName", "") or ""
    except Exception:
        return ""
    return ""


def _ensure_account(session: Session, domain, provider: str, email: str = "") -> ConnectorAccount:
    """이메일이 있으면 계정별 고유 id(멀티계정). 없으면 provider 단일 폴백."""
    dom = domain.value if hasattr(domain, "value") else str(domain)
    meta = _REAL_META.get(provider, (dom, provider, "", "ink"))
    if email:
        acct_id = f"ca-{dom}-{provider}-{hashlib.sha1(email.lower().encode()).hexdigest()[:8]}"
    else:
        acct_id = f"ca-{dom}-{provider}"
    acct = session.get(ConnectorAccount, acct_id)
    if not acct:
        acct = ConnectorAccount(
            id=acct_id,
            user_id="jiwoo",
            domain=domain,
            provider=provider,
            name=email or meta[1],
            kind=meta[2],
            tone=meta[3],
        )
        session.add(acct)
    if email:
        acct.external_account_id = email
        acct.name = email
    # 메일 계정은 메일 페이지 계정 탭/필터에 노출되도록 MailAccount 도 보장.
    if dom == "mail" and email:
        ma_id = mail_account_id_for(email)
        if not session.get(MailAccount, ma_id):
            session.add(
                MailAccount(
                    id=ma_id,
                    name=email.split("@")[0],
                    email=email,
                    tone=meta[3],
                    kind=meta[2],
                    connector=provider,
                    sort_order=100,
                )
            )
    return acct


def _connect(session: Session, domain, provider: str, email: str, tok: dict) -> ConnectorAccount:
    """provider 계정을 real+connected 로 만들고 토큰 저장(여러 도메인이 토큰 공유 가능)."""
    acct = _ensure_account(session, domain, provider, email)
    acct.mode = ConnectorMode.real
    acct.state = ConnState.connected
    acct.token_enc = encrypt_token(tok)
    acct.scopes = tok.get("scope", "")
    acct.last_label = "방금 연결됨"
    session.add(acct)
    # 캘린더 계정은 연결 즉시 Calendar 행 보장 → 동기화 전에도 일정 화면에 노출.
    dom = domain.value if hasattr(domain, "value") else str(domain)
    if dom == "calendar":
        from .calendar.normalize import ensure_calendar_row  # 지연 import(순환 회피)

        ensure_calendar_row(session, acct)
    return acct


def finish_oauth(session: Session, code: str, state: str) -> ConnectorAccount:
    os_row = session.get(OAuthState, state)
    if not os_row:
        raise ValueError("invalid oauth state")
    fam = _family(os_row.provider)
    conf = _conf(fam)
    client_id, client_secret = _creds(fam)
    data = {
        "code": code,
        "grant_type": "authorization_code",
        "redirect_uri": _redirect_uri(),
        "code_verifier": os_row.code_verifier,
        "client_id": client_id,
        "client_secret": client_secret,
    }
    r = httpx.post(conf["token"], data=data, timeout=get_settings().connector_http_timeout)
    r.raise_for_status()
    tok = r.json()
    tok["expires_at"] = int(time.time()) + int(tok.get("expires_in", 3600))
    email = _fetch_identity(fam, tok.get("access_token", ""))
    acct = _connect(session, os_row.domain, os_row.provider, email, tok)
    # 메일 연결 시, 같은 동의에 포함된 캘린더 권한으로 캘린더 계정도 함께 연결.
    if os_row.provider == "gmail":
        _connect(session, ConnectorDomain.calendar, "google_calendar", email, tok)
    elif os_row.provider == "outlook":
        _connect(session, ConnectorDomain.calendar, "outlook_calendar", email, tok)
    session.delete(os_row)
    session.commit()
    return acct


def valid_access_token(session: Session, account: ConnectorAccount) -> str:
    """만료되면 refresh_token 으로 갱신. 갱신 불가면 token_expired 로 표시 후 예외."""
    tok = decrypt_token(account.token_enc)
    if not tok:
        account.state = ConnState.token_expired
        session.add(account)
        session.commit()
        raise PermissionError("token missing/corrupt")
    if tok.get("expires_at", 0) > int(time.time()) + 60:
        return tok["access_token"]
    fam = _family(account.provider)
    conf = _conf(fam)
    client_id, client_secret = _creds(fam)
    rt = tok.get("refresh_token")
    if not rt:
        account.state = ConnState.token_expired
        session.add(account)
        session.commit()
        raise PermissionError("no refresh_token")
    r = httpx.post(
        conf["token"],
        data={
            "grant_type": "refresh_token",
            "refresh_token": rt,
            "client_id": client_id,
            "client_secret": client_secret,
        },
        timeout=get_settings().connector_http_timeout,
    )
    if r.status_code != 200:
        account.state = ConnState.token_expired
        session.add(account)
        session.commit()
        raise PermissionError("refresh failed")
    new = r.json()
    new.setdefault("refresh_token", rt)  # MS 는 새 refresh_token 을 줄 수도, 안 줄 수도 있음
    new["expires_at"] = int(time.time()) + int(new.get("expires_in", 3600))
    account.token_enc = encrypt_token(new)
    session.add(account)
    session.commit()
    return new["access_token"]