提交初始文件，增加授权应用路由

2026-06-10 10:47:18 +08:00
parent d8bfebdcfb
commit b24c176b88
6 changed files with 590 additions and 1 deletions
@@ -0,0 +1,353 @@
+import random
+from datetime import datetime
+
+import httpx
+from sqlalchemy import select, func
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.config import settings
+from app.models.user_oauth import UserOAuth
+from app.utils.id_gen import generate_id
+
+
+OAUTH_TYPE_CONFIG = {
+    1: {"port_type": 1, "name": "千川", "app_type": "juliang_qianchuan"},
+    2: {"port_type": 1, "name": "广告", "app_type": "juliang_ad"},
+    3: {"port_type": 1, "name": "本地推", "app_type": "juliang_ad"},
+    4: {"port_type": 1, "name": "星图", "app_type": "juliang_ad"},
+    5: {"port_type": 2, "name": "快手代理商", "app_type": "kuaishou"},
+    6: {"port_type": 3, "name": "巨量星图", "app_type": "juliang_ad"},
+    7: {"port_type": 4, "name": "巨量服务单", "app_type": "juliang_ad"},
+    8: {"port_type": 4, "name": "腾讯服务单", "app_type": "tencent"},
+    9: {"port_type": 5, "name": "腾讯营销K2", "app_type": "tencent"},
+    10: {"port_type": 5, "name": "腾讯营销K3", "app_type": "tencent"},
+}
+
+
+async def get_available_app(app_type: str, db: AsyncSession) -> dict:
+    if app_type == "juliang_ad":
+        apps = settings.JULIANG_AD_APPS
+    elif app_type == "juliang_qianchuan":
+        apps = settings.JULIANG_QIANCHUAN_APPS
+    elif app_type == "kuaishou":
+        apps = settings.KUAISHOU_APPS
+    elif app_type == "tencent":
+        apps = settings.TENCENT_APPS
+    else:
+        raise ValueError(f"不支持的应用类型: {app_type}")
+
+    if not apps:
+        raise ValueError(f"{app_type}未配置应用")
+
+    available_apps = []
+    for app in apps:
+        app_id = app.get("app_id")
+        if not app_id:
+            continue
+
+        result = await db.execute(
+            select(func.count(UserOAuth.id)).where(UserOAuth.appid == app_id)
+        )
+        count = result.scalar() or 0
+
+        if count < 5000:
+            available_apps.append(app)
+
+    if not available_apps:
+        raise ValueError("所有应用授权已超过最大数量")
+
+    return random.choice(available_apps)
+
+
+async def build_oauth_url(oauth_type: int, user_id: str) -> str:
+    if oauth_type == 1:
+        return await _build_juliang_oauth_url(oauth_type, user_id, app_type)
+    elif oauth_type == 2:
+        return await _build_kuaishou_oauth_url(oauth_type, user_id)
+    elif oauth_type == 3:
+        return await _build_tencent_oauth_url(oauth_type, user_id)
+    elif oauth_type == 4:
+        return await _build_tencent_oauth_url(oauth_type, user_id)    
+    else:
+        raise ValueError(f"不支持的应用类型: {app_type}")
+
+
+async def _build_juliang_oauth_url(oauth_type: int, user_id: str, app_type: str) -> str:
+    async with AsyncSession() as db:
+        app = await get_available_app(app_type, db)
+        app_id = app.get("app_id")
+
+    redirect_uri = "https://open.oceanengine.com/audit/oauth.html"
+    rid = "ktm0cl7napb"
+    if oauth_type == 1:
+        redirect_uri = "https://qianchuan.jinritemai.com/openapi/qc/audit/oauth.html"
+        rid = "vr7kclvmvs9"
+
+    params = {
+        "app_id": app_id,
+        "state": f"{oauth_type}:{user_id}:{app_id}:{app_type}",
+        "material_auth": 1,
+        "rid": rid,
+    }
+    query_string = "&".join(f"{k}={v}" for k, v in params.items())
+    return f"{redirect_uri}?{query_string}"
+
+
+async def _build_kuaishou_oauth_url(oauth_type: int, user_id: str) -> str:
+    async with AsyncSession() as db:
+        app = await get_available_app("kuaishou", db)
+        app_id = app.get("app_id")
+
+    redirect_uri = f"{settings.BASE_URL}/api/user-oauth/callback"
+    params = {
+        "app_id": app_id,
+        "redirect_uri": redirect_uri,
+        "response_type": "code",
+        "scope": "basic",
+        "state": f"{oauth_type}:{user_id}:{app_id}:kuaishou",
+    }
+    query_string = "&".join(f"{k}={v}" for k, v in params.items())
+    return f"https://open.kuaishou.com/oauth2/authorize?{query_string}"
+
+
+async def _build_tencent_oauth_url(oauth_type: int, user_id: str) -> str:
+    async with AsyncSession() as db:
+        app = await get_available_app("tencent", db)
+        app_id = app.get("app_id")
+
+    redirect_uri = f"{settings.BASE_URL}/api/user-oauth/callback"
+    params = {
+        "app_id": app_id,
+        "redirect_uri": redirect_uri,
+        "response_type": "code",
+        "scope": "get_user_info",
+        "state": f"{oauth_type}:{user_id}:{app_id}:tencent",
+    }
+    query_string = "&".join(f"{k}={v}" for k, v in params.items())
+    return f"https://api.e.qq.com/oauth/authorize?{query_string}"
+
+
+async def get_token_by_type(code: str, oauth_type: int, app_id: str, app_type: str) -> dict:
+    if app_type in ("juliang_ad", "juliang_qianchuan"):
+        return await get_juliang_token(code, oauth_type, app_id, app_type)
+    elif app_type == "kuaishou":
+        return await get_kuaishou_token(code, oauth_type, app_id)
+    elif app_type == "tencent":
+        return await get_tencent_token(code, oauth_type, app_id)
+    else:
+        raise ValueError(f"不支持的应用类型: {app_type}")
+
+
+async def get_juliang_token(code: str, oauth_type: int, app_id: str, app_type: str) -> dict:
+    url = "https://api.oceanengine.com/open_api/oauth2/access_token/"
+
+    if app_type == "juliang_ad":
+        apps = settings.JULIANG_AD_APPS
+    else:
+        apps = settings.JULIANG_QIANCHUAN_APPS
+
+    app = next((a for a in apps if a.get("app_id") == app_id), None)
+    if not app:
+        raise ValueError("应用配置不存在")
+
+    async with httpx.AsyncClient() as client:
+        response = await client.post(
+            url,
+            data={
+                "app_id": app_id,
+                "secret": app.get("secret"),
+                "auth_code": code,
+            },
+        )
+        response.raise_for_status()
+        content = response.json()
+        if content.get("code") != 0:
+            raise ValueError(content.get("message", "获取token失败"))
+        return content.get("data", {})
+
+
+async def get_kuaishou_token(code: str, oauth_type: int, app_id: str) -> dict:
+    url = "https://open.kuaishou.com/oauth2/token"
+    redirect_uri = f"{settings.BASE_URL}/api/user-oauth/callback"
+
+    app = next((a for a in settings.KUAISHOU_APPS if a.get("app_id") == app_id), None)
+    if not app:
+        raise ValueError("应用配置不存在")
+
+    async with httpx.AsyncClient() as client:
+        response = await client.post(
+            url,
+            data={
+                "app_id": app_id,
+                "secret": app.get("secret"),
+                "code": code,
+                "grant_type": "authorization_code",
+                "redirect_uri": redirect_uri,
+            },
+        )
+        response.raise_for_status()
+        return response.json()
+
+
+async def get_tencent_token(code: str, oauth_type: int, app_id: str) -> dict:
+    url = "https://api.e.qq.com/oauth/token"
+    redirect_uri = f"{settings.BASE_URL}/api/user-oauth/callback"
+
+    app = next((a for a in settings.TENCENT_APPS if a.get("app_id") == app_id), None)
+    if not app:
+        raise ValueError("应用配置不存在")
+
+    async with httpx.AsyncClient() as client:
+        response = await client.post(
+            url,
+            data={
+                "app_id": app_id,
+                "secret": app.get("secret"),
+                "code": code,
+                "grant_type": "authorization_code",
+                "redirect_uri": redirect_uri,
+            },
+        )
+        response.raise_for_status()
+        return response.json()
+
+
+async def get_account_info_by_type(token: dict, oauth_type: int, app_type: str) -> dict:
+    if app_type in ("juliang_ad", "juliang_qianchuan"):
+        return await _get_juliang_account_info(token)
+    elif app_type == "kuaishou":
+        return await _get_kuaishou_account_info(token)
+    elif app_type == "tencent":
+        return await _get_tencent_account_info(token)
+    else:
+        raise ValueError(f"不支持的应用类型: {app_type}")
+
+
+async def _get_juliang_account_info(token: dict) -> dict:
+    access_token = token.get("access_token")
+    url = "https://ad.oceanengine.com/openapi/oauth/user/info/"
+
+    async with httpx.AsyncClient() as client:
+        response = await client.get(
+            url,
+            headers={"Access-Token": access_token},
+        )
+        response.raise_for_status()
+        data = response.json()
+        if data.get("code") != 0:
+            raise ValueError(data.get("message", "获取账户信息失败"))
+        data = data.get("data", {})
+        return {
+            "account_id": data.get("advertiser_id", data.get("account_id", "")),
+            "account_name": data.get("advertiser_name", data.get("account_name", "")),
+            "account_role": data.get("role", ""),
+            "account_username": data.get("username", ""),
+        }
+
+
+async def _get_kuaishou_account_info(token: dict) -> dict:
+    access_token = token.get("access_token")
+    url = "https://open.kuaishou.com/api/user/info"
+
+    async with httpx.AsyncClient() as client:
+        response = await client.get(
+            url,
+            headers={"Authorization": f"Bearer {access_token}"},
+        )
+        response.raise_for_status()
+        data = response.json()
+        return {
+            "account_id": data.get("account_id", ""),
+            "account_name": data.get("account_name", ""),
+            "account_role": data.get("role", ""),
+            "account_username": data.get("username", ""),
+        }
+
+
+async def _get_tencent_account_info(token: dict) -> dict:
+    access_token = token.get("access_token")
+    url = "https://api.e.qq.com/user/info"
+
+    async with httpx.AsyncClient() as client:
+        response = await client.get(
+            url,
+            headers={"Authorization": f"Bearer {access_token}"},
+        )
+        response.raise_for_status()
+        data = response.json()
+        return {
+            "account_id": data.get("account_id", ""),
+            "account_name": data.get("account_name", ""),
+            "account_role": data.get("role", ""),
+            "account_username": data.get("username", ""),
+        }
+
+
+async def save_oauth_token(
+    db: AsyncSession,
+    user_id: str,
+    oauth_type: int,
+    token: dict,
+    account_info: dict,
+    app_id: str,
+) -> UserOAuth:
+    config = OAUTH_TYPE_CONFIG.get(oauth_type)
+    if not config:
+        raise ValueError(f"不支持的oauth_type: {oauth_type}")
+
+    access_token = token.get("access_token")
+    access_token_expired = token.get("expires_in")
+    refresh_token = token.get("refresh_token")
+    refresh_token_expired = token.get("refresh_token_expires_in")
+
+    expires_at = None
+    if access_token_expired:
+        expires_at = datetime.now().timestamp() + int(access_token_expired)
+        expires_at = datetime.fromtimestamp(expires_at)
+
+    refresh_expires_at = None
+    if refresh_token_expired:
+        refresh_expires_at = datetime.now().timestamp() + int(refresh_token_expired)
+        refresh_expires_at = datetime.fromtimestamp(refresh_expires_at)
+
+    existing = await db.execute(
+        select(UserOAuth).where(
+            UserOAuth.user_id == user_id,
+            UserOAuth.open_type == oauth_type,
+            UserOAuth.account_id == account_info.get("account_id", ""),
+        ).limit(1)
+    )
+    existing_oauth = existing.scalar_one_or_none()
+
+    if existing_oauth:
+        existing_oauth.access_token = access_token
+        existing_oauth.access_token_expired = expires_at
+        existing_oauth.refresh_token = refresh_token
+        existing_oauth.refresh_token_expired = refresh_expires_at
+        existing_oauth.account_name = account_info.get("account_name", "")
+        existing_oauth.account_role = account_info.get("account_role", "")
+        existing_oauth.account_username = account_info.get("account_username", "")
+        existing_oauth.appid = app_id
+        await db.flush()
+        return existing_oauth
+
+    user_oauth = UserOAuth(
+        id=generate_id(),
+        account_id=account_info.get("account_id", ""),
+        account_name=account_info.get("account_name", ""),
+        account_role=account_info.get("account_role", ""),
+        account_username=account_info.get("account_username", ""),
+        user_id=user_id,
+        open_type=oauth_type,
+        port_type=config["port_type"],
+        appid=app_id,
+        access_token=access_token,
+        access_token_expired=expires_at,
+        refresh_token=refresh_token,
+        refresh_token_expired=refresh_expires_at,
+        material_auth_status=True,
+    )
+
+    db.add(user_oauth)
+    await db.flush()
+    return user_oauth