240 lines
8.5 KiB
Python
240 lines
8.5 KiB
Python
|
import asyncio
|
||
|
|
import json
|
||
|
|
import logging
|
||
|
|
import os
|
||
|
|
import time
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
|
||
|
|
from yuxi.channel.extensions.jd.config import JDConfig
|
||
|
|
from yuxi.channel.extensions.jd.types import JDAccount, JOSToken
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
OAUTH_TOKEN_URL = "https://oauth.jd.com/oauth/token"
|
||
|
|
TOKEN_REFRESH_MARGIN = 300
|
||
|
|
REFRESH_INTERVAL = 1800
|
||
|
|
RETRY_MAX = 3
|
||
|
|
RETRY_BASE_DELAY = 2
|
||
|
|
|
||
|
|
_TOKEN_STORE_DIR = os.path.join(os.path.dirname(__file__), ".token_cache")
|
||
|
|
|
||
|
|
|
||
|
|
class JDGateway:
|
||
|
|
def __init__(self):
|
||
|
|
self._account: JDAccount | None = None
|
||
|
|
self._queue: asyncio.Queue | None = None
|
||
|
|
self._cancel_event = asyncio.Event()
|
||
|
|
self._access_token: str | None = None
|
||
|
|
self._token_obj: JOSToken | None = None
|
||
|
|
self._token_expires_at: float = 0
|
||
|
|
self._refresh_task: asyncio.Task | None = None
|
||
|
|
self._http: httpx.AsyncClient | None = None
|
||
|
|
self._running = False
|
||
|
|
self._token_lock = asyncio.Lock()
|
||
|
|
self._last_chat_id: dict[str, str] = {}
|
||
|
|
|
||
|
|
async def start(self, ctx) -> object:
|
||
|
|
config = JDConfig()
|
||
|
|
account = config.resolve_account(ctx.account_id if hasattr(ctx, "account_id") else "default")
|
||
|
|
|
||
|
|
if not account.is_configured():
|
||
|
|
logger.warning("JD account not configured, skipping start")
|
||
|
|
return {"running": False, "reason": "not-configured"}
|
||
|
|
|
||
|
|
self._account = account
|
||
|
|
self._queue = asyncio.Queue(maxsize=1000)
|
||
|
|
self._cancel_event.clear()
|
||
|
|
self._http = httpx.AsyncClient(timeout=15.0)
|
||
|
|
|
||
|
|
try:
|
||
|
|
await self._restore_or_fetch_token()
|
||
|
|
self._refresh_task = asyncio.create_task(self._token_refresh_loop())
|
||
|
|
self._running = True
|
||
|
|
logger.info(
|
||
|
|
"JD gateway started: account=%s shop=%s token_expires=%d",
|
||
|
|
account.account_id,
|
||
|
|
account.shop_id,
|
||
|
|
int(self._token_expires_at - time.time()) if self._token_expires_at else 0,
|
||
|
|
)
|
||
|
|
return {
|
||
|
|
"running": True,
|
||
|
|
"account_id": account.account_id,
|
||
|
|
"shop_id": account.shop_id,
|
||
|
|
"queue": self._queue,
|
||
|
|
}
|
||
|
|
except Exception:
|
||
|
|
logger.exception("JD gateway failed to start")
|
||
|
|
await self.stop(ctx)
|
||
|
|
raise
|
||
|
|
|
||
|
|
async def stop(self, ctx) -> None:
|
||
|
|
self._running = False
|
||
|
|
self._cancel_event.set()
|
||
|
|
|
||
|
|
if self._refresh_task and not self._refresh_task.done():
|
||
|
|
self._refresh_task.cancel()
|
||
|
|
try:
|
||
|
|
await self._refresh_task
|
||
|
|
except asyncio.CancelledError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
if self._http:
|
||
|
|
await self._http.aclose()
|
||
|
|
self._http = None
|
||
|
|
|
||
|
|
self._queue = None
|
||
|
|
self._last_chat_id.clear()
|
||
|
|
logger.info("JD gateway stopped")
|
||
|
|
|
||
|
|
@property
|
||
|
|
def account(self) -> JDAccount | None:
|
||
|
|
return self._account
|
||
|
|
|
||
|
|
@property
|
||
|
|
def queue(self) -> asyncio.Queue | None:
|
||
|
|
return self._queue
|
||
|
|
|
||
|
|
def get_token(self, account_id: str = "default") -> str | None:
|
||
|
|
if self._access_token and time.time() < self._token_expires_at:
|
||
|
|
return self._access_token
|
||
|
|
return None
|
||
|
|
|
||
|
|
def get_account(self, account_id: str = "default") -> JDAccount | None:
|
||
|
|
return self._account
|
||
|
|
|
||
|
|
def cache_chat_id(self, user_id: str, chat_id: str):
|
||
|
|
self._last_chat_id[user_id] = chat_id
|
||
|
|
|
||
|
|
def get_cached_chat_id(self, user_id: str) -> str | None:
|
||
|
|
return self._last_chat_id.get(user_id)
|
||
|
|
|
||
|
|
async def _restore_or_fetch_token(self):
|
||
|
|
account = self._account
|
||
|
|
|
||
|
|
loaded = await self._load_token_from_disk()
|
||
|
|
if loaded and loaded.get("expires_at", 0) > time.time() + TOKEN_REFRESH_MARGIN:
|
||
|
|
self._access_token = loaded["access_token"]
|
||
|
|
self._token_expires_at = loaded["expires_at"]
|
||
|
|
if loaded.get("refresh_token"):
|
||
|
|
self._account.refresh_token = loaded["refresh_token"]
|
||
|
|
logger.info("JD token restored from disk")
|
||
|
|
return
|
||
|
|
|
||
|
|
if account.access_token and account.token_expires_at:
|
||
|
|
if account.token_expires_at > time.time() + TOKEN_REFRESH_MARGIN:
|
||
|
|
self._access_token = account.access_token
|
||
|
|
self._token_expires_at = account.token_expires_at
|
||
|
|
return
|
||
|
|
|
||
|
|
if account.refresh_token:
|
||
|
|
try:
|
||
|
|
await self._refresh_token()
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("JD token refresh failed on start: %s", e)
|
||
|
|
|
||
|
|
if not self._access_token:
|
||
|
|
logger.warning("JD account %s has no valid access token", account.account_id)
|
||
|
|
|
||
|
|
async def _token_refresh_loop(self):
|
||
|
|
while not self._cancel_event.is_set():
|
||
|
|
try:
|
||
|
|
await asyncio.sleep(REFRESH_INTERVAL)
|
||
|
|
if self._cancel_event.is_set():
|
||
|
|
return
|
||
|
|
|
||
|
|
remaining = self._token_expires_at - time.time()
|
||
|
|
if remaining < TOKEN_REFRESH_MARGIN:
|
||
|
|
logger.info("JD token refresh triggered: remaining=%ds", int(remaining))
|
||
|
|
await self._refresh_token()
|
||
|
|
|
||
|
|
except asyncio.CancelledError:
|
||
|
|
break
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("JD token refresh loop error: %s", e)
|
||
|
|
|
||
|
|
async def _refresh_token(self):
|
||
|
|
async with self._token_lock:
|
||
|
|
if not self._account.refresh_token:
|
||
|
|
return
|
||
|
|
|
||
|
|
params = {
|
||
|
|
"grant_type": "refresh_token",
|
||
|
|
"client_id": self._account.app_key,
|
||
|
|
"client_secret": self._account.app_secret,
|
||
|
|
"refresh_token": self._account.refresh_token,
|
||
|
|
}
|
||
|
|
|
||
|
|
for attempt in range(RETRY_MAX):
|
||
|
|
try:
|
||
|
|
resp = await self._http.post(OAUTH_TOKEN_URL, data=params)
|
||
|
|
resp.raise_for_status()
|
||
|
|
data = resp.json()
|
||
|
|
|
||
|
|
expires_in = data.get("expires_in", 7200)
|
||
|
|
self._access_token = data["access_token"]
|
||
|
|
self._token_expires_at = time.time() + expires_in - TOKEN_REFRESH_MARGIN
|
||
|
|
|
||
|
|
new_refresh = data.get("refresh_token", "")
|
||
|
|
if new_refresh:
|
||
|
|
self._account.refresh_token = new_refresh
|
||
|
|
|
||
|
|
self._token_obj = JOSToken(
|
||
|
|
access_token=data["access_token"],
|
||
|
|
refresh_token=data.get("refresh_token", self._account.refresh_token or ""),
|
||
|
|
expires_at=time.time() + expires_in,
|
||
|
|
uid=data.get("uid", ""),
|
||
|
|
user_nick=data.get("user_nick", ""),
|
||
|
|
)
|
||
|
|
|
||
|
|
await self._save_token_to_disk()
|
||
|
|
|
||
|
|
logger.info(
|
||
|
|
"JD token refreshed: expires_in=%ds uid=%s",
|
||
|
|
expires_in,
|
||
|
|
data.get("uid", ""),
|
||
|
|
)
|
||
|
|
return
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.warning("JD token refresh attempt %d failed: %s", attempt + 1, e)
|
||
|
|
if attempt == RETRY_MAX - 1:
|
||
|
|
raise
|
||
|
|
delay = RETRY_BASE_DELAY * (2**attempt)
|
||
|
|
await asyncio.sleep(delay)
|
||
|
|
|
||
|
|
async def _save_token_to_disk(self):
|
||
|
|
if not self._token_obj:
|
||
|
|
return
|
||
|
|
try:
|
||
|
|
os.makedirs(_TOKEN_STORE_DIR, exist_ok=True)
|
||
|
|
token_file = os.path.join(_TOKEN_STORE_DIR, f"{self._account.shop_id}.json")
|
||
|
|
with open(token_file, "w") as f:
|
||
|
|
json.dump(
|
||
|
|
{
|
||
|
|
"access_token": self._access_token,
|
||
|
|
"refresh_token": self._account.refresh_token,
|
||
|
|
"expires_at": self._token_expires_at,
|
||
|
|
"uid": self._token_obj.uid,
|
||
|
|
"user_nick": self._token_obj.user_nick,
|
||
|
|
},
|
||
|
|
f,
|
||
|
|
)
|
||
|
|
except Exception as e:
|
||
|
|
logger.warning("Failed to save JD token to disk: %s", e)
|
||
|
|
|
||
|
|
async def _load_token_from_disk(self) -> dict | None:
|
||
|
|
if not self._account:
|
||
|
|
return None
|
||
|
|
token_file = os.path.join(_TOKEN_STORE_DIR, f"{self._account.shop_id}.json")
|
||
|
|
if not os.path.exists(token_file):
|
||
|
|
return None
|
||
|
|
try:
|
||
|
|
with open(token_file) as f:
|
||
|
|
data = json.load(f)
|
||
|
|
if data.get("expires_at", 0) > time.time():
|
||
|
|
return data
|
||
|
|
except Exception as e:
|
||
|
|
logger.warning("Failed to load JD token from disk: %s", e)
|
||
|
|
return None
|