354 lines
14 KiB
Python
354 lines
14 KiB
Python
import random
|
||
from datetime import datetime, timedelta, timezone
|
||
|
||
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.models.user_oauth_app import UserOAuthApp
|
||
from app.utils.id_gen import generate_id
|
||
from app.tasks.token_refresh_task import _update_redis_token
|
||
|
||
#随机获取一个可用的应用配置
|
||
async def get_available_app(open_type: int, db: AsyncSession) -> dict:
|
||
# 从 user_oauth_app 表查询可用应用
|
||
# 过滤条件：open_type 匹配、status=1（正常）、deleted_at is None
|
||
result = await db.execute(
|
||
select(UserOAuthApp).where(
|
||
UserOAuthApp.open_type == open_type,
|
||
UserOAuthApp.status == 1,
|
||
UserOAuthApp.deleted_at.is_(None),
|
||
)
|
||
)
|
||
apps = result.scalars().all()
|
||
|
||
if not apps:
|
||
raise ValueError("没有找到可用的应用配置")
|
||
|
||
available_apps = []
|
||
for app in apps:
|
||
app_id = app.app_id
|
||
|
||
# 查询该应用当前授权数量
|
||
count_result = await db.execute(
|
||
select(func.count(UserOAuth.id)).where(
|
||
UserOAuth.appid == app_id,
|
||
UserOAuth.deleted_at.is_(None),
|
||
)
|
||
)
|
||
count = count_result.scalar() or 0
|
||
|
||
# 检查是否达到最大授权数
|
||
max_users = app.max_count
|
||
if count < max_users:
|
||
available_apps.append({
|
||
"app_id": app.app_id,
|
||
"secret": app.secret,
|
||
"open_type": app.open_type,
|
||
"auth_url": app.auth_url,
|
||
"company": app.company,
|
||
})
|
||
|
||
if not available_apps:
|
||
raise ValueError("所有应用授权已超过最大数量")
|
||
|
||
return random.choice(available_apps)
|
||
|
||
|
||
async def build_oauth_url(open_type: int, user_id: str, db: AsyncSession) -> str:
|
||
if open_type == 1:
|
||
#千川
|
||
return await _build_jl_oauth_url(open_type, user_id, db)
|
||
elif open_type in [2, 3]:
|
||
#广告,本地推
|
||
return await _build_jl_oauth_url(2, user_id, db)
|
||
elif open_type == 5:
|
||
#快手代理商
|
||
return "无配置"
|
||
elif open_type == 9:
|
||
#腾讯营销K2
|
||
return "无配置"
|
||
elif open_type == 10:
|
||
#腾讯营销K3
|
||
return "无配置"
|
||
else:
|
||
raise ValueError(f"不支持的应用类型: {app_type}")
|
||
|
||
#千川授权链接构建
|
||
async def _build_jl_oauth_url(open_type: int, user_id: str, db: AsyncSession) -> str:
|
||
app = await get_available_app(open_type, db)
|
||
app_id = app.get("app_id")
|
||
auth_url = app.get("auth_url")
|
||
if not auth_url:
|
||
raise ValueError("应用授权链接不能为空")
|
||
params = {
|
||
"app_id": app_id,
|
||
"state": f"{user_id}:{app_id}",
|
||
}
|
||
query_string = "&".join(f"{k}={v}" for k, v in params.items())
|
||
return f"{auth_url}&{query_string}"
|
||
|
||
|
||
async def get_token(code: str, user_id: str, app_id: str, db: AsyncSession) -> dict:
|
||
#1.根据app_id查询应用配置
|
||
result = await db.execute(
|
||
select(UserOAuthApp).where(UserOAuthApp.app_id == app_id)
|
||
)
|
||
app = result.scalar_one_or_none()
|
||
if not app:
|
||
raise ValueError("应用配置不存在")
|
||
|
||
open_type = app.open_type
|
||
secret = app.secret
|
||
|
||
if open_type == 1 or open_type == 2:
|
||
#千川或广告
|
||
return await get_juliang_token(app_id, secret, code, open_type, user_id, db)
|
||
elif open_type == 5:
|
||
#快手代理商
|
||
return await get_kuaishou_token(app_id, secret, code, open_type, user_id, db)
|
||
elif open_type == 9 or open_type == 10:
|
||
#腾讯营销K2或腾讯营销K3
|
||
return await get_tencent_token(app_id, secret, code, open_type, user_id, db)
|
||
else:
|
||
raise ValueError(f"不支持的应用类型: {open_type}")
|
||
|
||
|
||
async def get_juliang_token(app_id: str, secret: str, code: str, open_type: int, user_id: str, db: AsyncSession) -> dict:
|
||
|
||
async with httpx.AsyncClient() as client:
|
||
#1.请求token
|
||
url = "https://api.oceanengine.com/open_api/oauth2/access_token/"
|
||
response = await client.post(
|
||
url,
|
||
json={
|
||
"app_id": app_id,
|
||
"secret": secret,
|
||
"auth_code": code,
|
||
},
|
||
)
|
||
response.raise_for_status()
|
||
data = response.json()
|
||
if data.get("code") != 0:
|
||
raise ValueError(data.get("message", "获取token失败")+f"，错误信息：{data.get('message', '')}")
|
||
data = data.get("data", {})
|
||
access_token = data.get("access_token", "")
|
||
refresh_token = data.get("refresh_token", "")
|
||
expires_in = datetime.now(timezone.utc) + timedelta(seconds=data.get("expires_in", 0))
|
||
refresh_token_expires_in = datetime.now(timezone.utc) + timedelta(seconds=data.get("refresh_token_expires_in", 0))
|
||
#2.获取已授权角色账户，一个授权可能有多个角色账户
|
||
url = "https://api.oceanengine.com/open_api/oauth2/advertiser/get/"
|
||
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", "获取已授权账户失败")+f"，错误信息：{data.get('message', '')}")
|
||
data = data.get("data", {})
|
||
account_list = data.get("list", [])
|
||
#3.获取已授权登录信息
|
||
url = "https://api.oceanengine.com/open_api/2/user/info/"
|
||
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", "获取已授权登录信息失败")+f"，错误信息：{data.get('message', '')}")
|
||
data = data.get("data", {})
|
||
account_username = data.get("email", "")
|
||
account_userid = str(data.get("id", ""))
|
||
material_auth_status = data.get("material_auth_status", False)
|
||
#4.根据登录信息判断是否新增token或更新token，不同的登录信息对应不同的token，然后更新数据库
|
||
#查询email，appid，user_id是否存在已授权记录
|
||
oauth = await db.execute(
|
||
select(UserOAuth)
|
||
.where(UserOAuth.account_username == account_username, UserOAuth.account_userid == account_userid, UserOAuth.appid == app_id, UserOAuth.user_id == user_id)
|
||
.limit(1)
|
||
)
|
||
oauth = oauth.scalar_one_or_none()
|
||
if not oauth:
|
||
new_oauth_ids = []
|
||
for account in account_list:
|
||
oauth_id = generate_id()
|
||
new_oauth_ids.append(oauth_id)
|
||
#新增授权记录
|
||
db.add(UserOAuth(
|
||
id=oauth_id,
|
||
account_id = str(account.get("account_id", "")),
|
||
account_name = account.get("account_name", ""),
|
||
account_role = account.get("account_role", ""),
|
||
account_username = account_username,
|
||
account_userid = account_userid,
|
||
user_id = user_id,
|
||
appid = app_id,
|
||
open_type = open_type,
|
||
port_type = 1,
|
||
access_token = access_token,
|
||
access_token_expired = expires_in,
|
||
refresh_token = refresh_token,
|
||
refresh_token_expired = refresh_token_expires_in,
|
||
material_auth_status = material_auth_status,
|
||
))
|
||
await db.commit()
|
||
|
||
|
||
for oauth_id in new_oauth_ids:
|
||
await _update_redis_token(oauth_id, access_token, expires_in)
|
||
else:
|
||
# 查询现有授权记录（未删除的）
|
||
existing_accounts = await db.execute(
|
||
select(UserOAuth).where(
|
||
UserOAuth.account_username == account_username,
|
||
UserOAuth.account_userid == account_userid,
|
||
UserOAuth.appid == app_id,
|
||
UserOAuth.user_id == user_id,
|
||
UserOAuth.deleted_at.is_(None),
|
||
)
|
||
)
|
||
existing_accounts = {acc.account_id: acc for acc in existing_accounts.scalars().all()}
|
||
|
||
# 新的账户列表
|
||
new_account_ids = {str(account.get("account_id")) for account in account_list}
|
||
old_account_ids = set(existing_accounts.keys())
|
||
|
||
# 需要更新Redis的oauth_id列表
|
||
update_redis_ids = []
|
||
|
||
# 1. 软删除已消失的账户
|
||
for account_id in old_account_ids - new_account_ids:
|
||
existing_accounts[account_id].deleted_at = datetime.now(timezone.utc)
|
||
|
||
# 2. 更新或新增账户
|
||
for account in account_list:
|
||
account_id = str(account.get("account_id", ""))
|
||
if account_id in existing_accounts:
|
||
# 更新现有记录
|
||
existing_oauth = existing_accounts[account_id]
|
||
existing_oauth.account_name = account.get("account_name", "")
|
||
existing_oauth.account_role = account.get("account_role", "")
|
||
existing_oauth.access_token = access_token
|
||
existing_oauth.access_token_expired = expires_in
|
||
existing_oauth.refresh_token = refresh_token
|
||
existing_oauth.refresh_token_expired = refresh_token_expires_in
|
||
existing_oauth.material_auth_status = material_auth_status
|
||
update_redis_ids.append(existing_oauth.id)
|
||
else:
|
||
# 新增记录
|
||
oauth_id = generate_id()
|
||
update_redis_ids.append(oauth_id)
|
||
db.add(UserOAuth(
|
||
id=oauth_id,
|
||
account_id=str(account_id),
|
||
account_name=account.get("account_name", ""),
|
||
account_role=account.get("account_role", ""),
|
||
account_username=account_username,
|
||
account_userid = account_userid,
|
||
user_id=user_id,
|
||
appid=app_id,
|
||
open_type=open_type,
|
||
port_type=1,
|
||
access_token = access_token,
|
||
access_token_expired=expires_in,
|
||
refresh_token=refresh_token,
|
||
refresh_token_expired=refresh_token_expires_in,
|
||
material_auth_status=material_auth_status,
|
||
))
|
||
await db.commit()
|
||
|
||
for oauth_id in update_redis_ids:
|
||
await _update_redis_token(oauth_id, access_token, expires_in)
|
||
|
||
#5.本次更新成功以后，判断是否有其他同一个appid，同一个授权登录账号的授权记录,如果有，则更新token信息
|
||
from sqlalchemy import update
|
||
|
||
related_oauths = await db.execute(
|
||
select(UserOAuth).where(
|
||
UserOAuth.account_username == account_username,
|
||
UserOAuth.account_userid == account_userid,
|
||
UserOAuth.appid == app_id,
|
||
UserOAuth.user_id != user_id,
|
||
UserOAuth.deleted_at.is_(None),
|
||
)
|
||
)
|
||
related_oauths = related_oauths.scalars().all()
|
||
|
||
if related_oauths:
|
||
await db.execute(
|
||
update(UserOAuth).where(
|
||
UserOAuth.account_username == account_username,
|
||
UserOAuth.account_userid == account_userid,
|
||
UserOAuth.appid == app_id,
|
||
UserOAuth.user_id != user_id,
|
||
UserOAuth.deleted_at.is_(None),
|
||
).values(
|
||
access_token=access_token,
|
||
access_token_expired=expires_in,
|
||
refresh_token=refresh_token,
|
||
refresh_token_expired=refresh_token_expires_in,
|
||
material_auth_status=material_auth_status,
|
||
)
|
||
)
|
||
await db.commit()
|
||
|
||
for related_oauth in related_oauths:
|
||
await _update_redis_token(related_oauth.id, access_token, expires_in)
|
||
|
||
return {"message": "授权成功"}
|
||
|
||
async def get_kuaishou_token(app_id: str, secret: str, code: str, oauth_type: int, user_id: str, db: AsyncSession) -> dict:
|
||
return "未配置"
|
||
|
||
|
||
async def get_tencent_token(app_id: str, secret: str, code: str, oauth_type: int, user_id: str, db: AsyncSession) -> str:
|
||
return "未配置"
|
||
|
||
|
||
async def get_oauth_list(
|
||
user_id: str,
|
||
db: AsyncSession,
|
||
account_userid: str | None = None,
|
||
open_type: int | None = None,
|
||
account_id: str | None = None,
|
||
page: int = 1,
|
||
page_size: int = 10,
|
||
) -> dict:
|
||
if page < 1:
|
||
page = 1
|
||
if page_size < 1:
|
||
page_size = 10
|
||
|
||
query = select(UserOAuth).where(
|
||
UserOAuth.user_id == user_id,
|
||
UserOAuth.deleted_at.is_(None),
|
||
)
|
||
|
||
if account_userid:
|
||
query = query.where(UserOAuth.account_userid == account_userid)
|
||
if open_type:
|
||
query = query.where(UserOAuth.open_type == open_type)
|
||
if account_id:
|
||
query = query.where(UserOAuth.account_id == account_id)
|
||
|
||
query = query.order_by(UserOAuth.created_at.desc())
|
||
|
||
total_result = await db.execute(query.with_only_columns(UserOAuth.id))
|
||
total = len(total_result.scalars().all())
|
||
|
||
offset = (page - 1) * page_size
|
||
query = query.offset(offset).limit(page_size)
|
||
|
||
result = await db.execute(query)
|
||
oauth_list = result.scalars().all()
|
||
|
||
return {
|
||
"data": oauth_list,
|
||
"total": total,
|
||
"page": page,
|
||
"page_size": page_size,
|
||
} |