understanding's picture
Update bot/integrations/cf_worker2.py
640b9a5 verified
# PATH: bot/integrations/cf_worker2.py
from __future__ import annotations
from bot.config import Workers
from bot.integrations.http import post_json
from bot.core.settings import ROTATE_AFTER
def _h_hf():
return {
"Authorization": f"Bearer {Workers.HF_API_KEY}",
"Content-Type": "application/json",
}
def _h_bot():
return {
"Authorization": f"Bearer {Workers.BOT_BACKEND_KEY}",
"Content-Type": "application/json",
}
def _unwrap(r: dict) -> dict:
"""
Unwrap post_json wrapper:
{"ok": true, "status": 200, "data": {...}} -> {...}
"""
if isinstance(r, dict) and isinstance(r.get("data"), dict):
data = r["data"]
if "ok" not in data and "ok" in r:
data["ok"] = r["ok"]
return data
return r
async def allow_user(tg_id: int) -> dict:
r = await post_json(f"{Workers.WORKER2_URL}/api/allow_user", _h_hf(), {"tg_id": str(tg_id)})
return _unwrap(r)
async def disallow_user(tg_id: int) -> dict:
r = await post_json(f"{Workers.WORKER2_URL}/api/disallow_user", _h_hf(), {"tg_id": str(tg_id)})
return _unwrap(r)
async def is_allowed(tg_id: int) -> bool:
r = await post_json(f"{Workers.WORKER2_URL}/api/is_allowed", _h_hf(), {"tg_id": str(tg_id)})
j = _unwrap(r)
return bool(j.get("ok") and j.get("allowed") is True)
async def pick_profile(tg_id: int, channel_id: str, rotate_after: int = ROTATE_AFTER) -> dict:
r = await post_json(
f"{Workers.WORKER2_URL}/api/pick_profile",
_h_hf(),
{"tg_id": str(tg_id), "channel_id": str(channel_id), "rotate_after": int(rotate_after)},
)
return _unwrap(r)
async def access_token(tg_id: int, profile_id: str) -> dict:
r = await post_json(
f"{Workers.WORKER2_URL}/api/access_token",
_h_hf(),
{"tg_id": str(tg_id), "profile_id": profile_id},
)
return _unwrap(r)
async def record_upload(tg_id: int, profile_id: str) -> dict:
r = await post_json(
f"{Workers.WORKER2_URL}/api/record_upload",
_h_hf(),
{"tg_id": str(tg_id), "profile_id": profile_id},
)
return _unwrap(r)
async def stats_today() -> dict:
r = await post_json(f"{Workers.WORKER2_URL}/api/stats_today", _h_hf(), {})
return _unwrap(r)
# ✅ Added: handlers.py expects list_profiles(uid, only_connected=...)
async def list_profiles(tg_id: int, only_connected: bool = False) -> dict:
r = await post_json(
f"{Workers.WORKER2_URL}/api/list_profiles",
_h_hf(),
{"tg_id": str(tg_id), "only_connected": bool(only_connected)},
)
return _unwrap(r)
# ✅ Added: handlers.py expects get_default_profile(uid) -> {ok, profile_id, access_token}
async def get_default_profile(tg_id: int) -> dict:
# worker2 supports default profile when profile_id is "" (it uses idx.default_profile_id)
j = await access_token(tg_id, "")
if isinstance(j, dict) and j.get("ok"):
return {
"ok": True,
"profile_id": j.get("profile_id"),
"access_token": j.get("access_token"),
"expires_in": j.get("expires_in"),
"channel_id": j.get("channel_id"),
"channel_title": j.get("channel_title"),
}
return j if isinstance(j, dict) else {"ok": False, "err": "access_token_failed"}
# ✅ Added: handlers.py expects set_default_profile(uid, pid)
# Your Pages backend protects /api/profile/* with BOT_BACKEND_KEY, so we call WORKER1_URL here.
async def set_default_profile(tg_id: int, profile_id: str) -> dict:
r = await post_json(
f"{Workers.WORKER1_URL}/api/profile/set_default",
_h_bot(),
{"tg_id": str(tg_id), "profile_id": str(profile_id)},
)
return _unwrap(r)