311 lines
11 KiB
Python
311 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import hmac
|
|
import json
|
|
import logging
|
|
import ssl
|
|
import time
|
|
from collections import defaultdict
|
|
|
|
from yuxi.channel.extensions.feishu.config import FeishuConfigAdapter
|
|
from yuxi.channel.extensions.feishu.dedup import get_deduplicator
|
|
from yuxi.channel.extensions.feishu.monitor import FeishuMonitor
|
|
from yuxi.channel.extensions.feishu.probe import probe_account
|
|
from yuxi.channel.extensions.feishu.client import get_client, get_bot_info
|
|
from yuxi.channel.extensions.feishu.types import FeishuAccount, FeishuConnectionMode
|
|
from yuxi.channel.extensions.feishu.ws_client import FeishuWSClient
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
MAX_WEBHOOK_BODY_BYTES = 2 * 1024 * 1024
|
|
WEBHOOK_READ_TIMEOUT = 10
|
|
|
|
|
|
class WebhookRateLimiter:
|
|
def __init__(self, max_per_second: int = 50):
|
|
self._max = max_per_second
|
|
self._buckets: dict[str, list[float]] = defaultdict(list)
|
|
|
|
def allow(self, key: str) -> bool:
|
|
now = time.monotonic()
|
|
bucket = self._buckets[key]
|
|
bucket[:] = [t for t in bucket if now - t < 1.0]
|
|
if len(bucket) >= self._max:
|
|
return False
|
|
bucket.append(now)
|
|
return True
|
|
|
|
|
|
class FeishuGatewayAdapter:
|
|
def __init__(self):
|
|
self._config_adapter = FeishuConfigAdapter()
|
|
self._monitor = FeishuMonitor()
|
|
self._ws_client: FeishuWSClient | None = None
|
|
self._running = False
|
|
self._tasks: list[asyncio.Task] = []
|
|
self._webhook_handlers: dict[str, dict] = {}
|
|
self._rate_limiter = WebhookRateLimiter()
|
|
|
|
async def start(self, ctx) -> object:
|
|
account_id = getattr(ctx, "account_id", "default")
|
|
config = getattr(ctx, "config", {}) or {}
|
|
|
|
account = await self._config_adapter.resolve_account(account_id, config)
|
|
|
|
if not account.is_configured():
|
|
logger.warning("Feishu account %s not configured, skipping", account_id)
|
|
return asyncio.Queue()
|
|
|
|
ctx.logger.info("Starting Feishu gateway for account %s (mode=%s)", account_id, account.connection_mode)
|
|
|
|
self._running = True
|
|
|
|
queue = getattr(ctx, "queue", asyncio.Queue())
|
|
cancel_event = getattr(ctx, "cancel_event", asyncio.Event())
|
|
|
|
dedup = get_deduplicator(account_id)
|
|
dedup.warmup()
|
|
|
|
try:
|
|
client = get_client(account.app_id, account.app_secret, account.domain, account.http_timeout_ms)
|
|
info = await get_bot_info(client)
|
|
bot_open_id = info.get("open_id", "")
|
|
if bot_open_id:
|
|
self._monitor.set_bot_open_id(account_id, bot_open_id)
|
|
ctx.logger.info("Bot open_id resolved: %s", bot_open_id)
|
|
except Exception:
|
|
ctx.logger.debug("Failed to fetch bot open_id", exc_info=True)
|
|
|
|
if account.connection_mode == FeishuConnectionMode.WEBSOCKET:
|
|
task = asyncio.create_task(
|
|
self._run_websocket_mode(account, queue, cancel_event, config)
|
|
)
|
|
self._tasks.append(task)
|
|
elif account.connection_mode == FeishuConnectionMode.WEBHOOK:
|
|
self._register_webhook_handler(account)
|
|
logger.info("Feishu webhook mode registered for account %s at path %s", account_id, account.webhook_path)
|
|
|
|
probe_task = asyncio.create_task(probe_account(
|
|
{
|
|
"app_id": account.app_id,
|
|
"app_secret": account.app_secret,
|
|
"domain": account.domain,
|
|
"http_timeout_ms": account.http_timeout_ms,
|
|
}
|
|
))
|
|
self._tasks.append(probe_task)
|
|
|
|
return queue
|
|
|
|
async def stop(self, ctx) -> None:
|
|
self._running = False
|
|
account_id = getattr(ctx, "account_id", "default")
|
|
|
|
if self._ws_client:
|
|
await self._ws_client.stop()
|
|
self._ws_client = None
|
|
|
|
for task in self._tasks:
|
|
if task and not task.done():
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
self._tasks.clear()
|
|
|
|
self._webhook_handlers.pop(account_id, None)
|
|
|
|
dedup = get_deduplicator(account_id)
|
|
dedup.reset()
|
|
|
|
logger.info("Feishu gateway stopped for account %s", account_id)
|
|
|
|
async def _run_websocket_mode(
|
|
self,
|
|
account: FeishuAccount,
|
|
queue: asyncio.Queue,
|
|
cancel_event: asyncio.Event,
|
|
config: dict,
|
|
) -> None:
|
|
self._ws_client = FeishuWSClient(account)
|
|
|
|
async def handle_message_event(event_type: str, event_data: dict):
|
|
try:
|
|
await self._monitor.handle_event(
|
|
{"type": event_type, "event": event_data},
|
|
account,
|
|
queue,
|
|
)
|
|
except Exception:
|
|
logger.exception("Error handling feishu event %s", event_type)
|
|
|
|
event_types = [
|
|
"im.message.receive_v1",
|
|
"im.message.reaction.created_v1",
|
|
"im.message.reaction.deleted_v1",
|
|
"card.action.trigger",
|
|
"im.chat.member.bot.added_v1",
|
|
"im.chat.member.bot.deleted_v1",
|
|
"im.chat.member.user.added_v1",
|
|
"im.chat.disbanded_v1",
|
|
"im.chat.updated_v1",
|
|
]
|
|
for et in event_types:
|
|
self._ws_client.on_event(et, handle_message_event)
|
|
|
|
try:
|
|
await self._ws_client.start()
|
|
|
|
while self._running and not cancel_event.is_set():
|
|
await asyncio.sleep(1)
|
|
|
|
except (ssl.SSLError, OSError) as e:
|
|
logger.warning(
|
|
"Feishu WebSocket error for account %s on first attempt: %s, retrying without SSL verification",
|
|
account.account_id, e,
|
|
)
|
|
if self._ws_client:
|
|
await self._ws_client.stop()
|
|
|
|
ssl._create_default_https_context = ssl._create_unverified_context
|
|
try:
|
|
self._ws_client = FeishuWSClient(account)
|
|
for et in event_types:
|
|
self._ws_client.on_event(et, handle_message_event)
|
|
await self._ws_client.start()
|
|
|
|
while self._running and not cancel_event.is_set():
|
|
await asyncio.sleep(1)
|
|
except Exception as retry_error:
|
|
logger.error(
|
|
"Feishu WebSocket retry also failed for account %s: %s",
|
|
account.account_id, retry_error,
|
|
)
|
|
except Exception as e:
|
|
logger.error("Feishu WebSocket error for account %s: %s", account.account_id, e)
|
|
finally:
|
|
if self._ws_client:
|
|
await self._ws_client.stop()
|
|
|
|
def _register_webhook_handler(self, account: FeishuAccount) -> None:
|
|
self._webhook_handlers[account.account_id] = {
|
|
"account": account,
|
|
"path": account.webhook_path,
|
|
"encrypt_key": account.encrypt_key,
|
|
"verification_token": account.verification_token,
|
|
}
|
|
|
|
async def handle_webhook_request(
|
|
self,
|
|
account_id: str,
|
|
headers: dict,
|
|
body: bytes,
|
|
client_ip: str = "",
|
|
) -> dict:
|
|
handler = self._webhook_handlers.get(account_id)
|
|
if not handler:
|
|
return {"status": 404, "body": json.dumps({"error": "Account not found"})}
|
|
|
|
account = handler["account"]
|
|
|
|
if len(body) > MAX_WEBHOOK_BODY_BYTES:
|
|
return {"status": 413, "body": json.dumps({"error": "Body too large"})}
|
|
|
|
if not self._rate_limiter.allow(f"{account_id}:{client_ip}"):
|
|
return {"status": 429, "body": json.dumps({"error": "Rate limit exceeded"})}
|
|
|
|
if not self._verify_signature(handler, headers, body):
|
|
return {"status": 401, "body": json.dumps({"error": "Invalid signature"})}
|
|
|
|
event = self._decrypt_event(handler, body)
|
|
if event is None:
|
|
return {"status": 400, "body": json.dumps({"error": "Invalid event data"})}
|
|
|
|
queue = asyncio.Queue()
|
|
|
|
try:
|
|
await self._monitor.handle_event(event, account, queue)
|
|
except Exception:
|
|
logger.exception("Error handling webhook event for account %s", account_id)
|
|
|
|
return {"status": 200, "body": json.dumps({"code": 0})}
|
|
|
|
def _verify_signature(self, handler: dict, headers: dict, body: bytes) -> bool:
|
|
verification_token = handler.get("verification_token")
|
|
if not verification_token:
|
|
logger.warning("Webhook verification token not configured for account %s (skipping signature check)",
|
|
handler.get("account", {}).get("account_id", "unknown"))
|
|
return True
|
|
|
|
signature = headers.get("x-lark-signature", "")
|
|
timestamp = headers.get("x-lark-request-timestamp", "")
|
|
nonce = headers.get("x-lark-request-nonce", "")
|
|
|
|
if not signature or not timestamp or not nonce:
|
|
return False
|
|
|
|
sign_str = f"{timestamp}{nonce}{verification_token}{body.decode('utf-8', errors='replace')}"
|
|
expected = hmac.new(
|
|
verification_token.encode("utf-8"),
|
|
sign_str.encode("utf-8"),
|
|
hashlib.sha256,
|
|
).hexdigest()
|
|
|
|
return hmac.compare_digest(signature, expected)
|
|
|
|
def _decrypt_event(self, handler: dict, body: bytes) -> dict | None:
|
|
try:
|
|
data = json.loads(body)
|
|
|
|
encrypt_key = handler.get("encrypt_key")
|
|
if encrypt_key and "encrypt" in data:
|
|
event_data = self._decrypt_with_key(encrypt_key, data.get("encrypt", ""))
|
|
if event_data:
|
|
return event_data
|
|
logger.warning("Failed to decrypt event data")
|
|
return None
|
|
|
|
if "challenge" in data:
|
|
return {"type": "challenge", "challenge": data.get("challenge", "")}
|
|
|
|
return data
|
|
except json.JSONDecodeError:
|
|
logger.warning("Invalid JSON in webhook body")
|
|
return None
|
|
|
|
def _decrypt_with_key(self, encrypt_key: str, encrypted: str) -> dict | None:
|
|
import base64
|
|
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
|
from cryptography.hazmat.primitives import padding
|
|
|
|
try:
|
|
raw = base64.b64decode(encrypted)
|
|
iv = raw[:16]
|
|
ciphertext = raw[16:]
|
|
|
|
key = hashlib.sha256(encrypt_key.encode("utf-8")).digest()
|
|
cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
|
|
decryptor = cipher.decryptor()
|
|
plaintext = decryptor.update(ciphertext) + decryptor.finalize()
|
|
|
|
unpadder = padding.PKCS7(128).unpadder()
|
|
plaintext = unpadder.update(plaintext) + unpadder.finalize()
|
|
|
|
return json.loads(plaintext)
|
|
except Exception:
|
|
return None
|
|
|
|
async def resolve_gateway_auth_bypass_paths(self, config: dict) -> list[str]:
|
|
return ["/feishu/events", "/lark/events"]
|
|
|
|
@property
|
|
def webhook_paths(self) -> list[str]:
|
|
return [h["path"] for h in self._webhook_handlers.values()]
|
|
|
|
@property
|
|
def has_webhook_handlers(self) -> bool:
|
|
return len(self._webhook_handlers) > 0
|