353 lines
12 KiB
Python
353 lines
12 KiB
Python
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 |