返回 JoyAI-Echo
whatsapp.py
1 """WhatsApp channel implementation using Node.js bridge."""
2
3 import asyncio
4 import json
5 import mimetypes
6 import os
7 import secrets
8 import shutil
9 import subprocess
10 from collections import OrderedDict
11 from pathlib import Path
12 from typing import Any, Literal
13
14 from loguru import logger
15 from pydantic import Field
16
17 from nanobot.bus.events import OutboundMessage
18 from nanobot.bus.queue import MessageBus
19 from nanobot.channels.base import BaseChannel
20 from nanobot.config.schema import Base
21
22
23 class WhatsAppConfig(Base):
24 """WhatsApp channel configuration."""
25
26 enabled: bool = False
27 bridge_url: str = "ws://localhost:3001"
28 bridge_token: str = ""
29 allow_from: list[str] = Field(default_factory=list)
30 group_policy: Literal["open", "mention"] = "open" # "open" responds to all, "mention" only when @mentioned
31
32
33 def _bridge_token_path() -> Path:
34 from nanobot.config.paths import get_runtime_subdir
35
36 return get_runtime_subdir("whatsapp-auth") / "bridge-token"
37
38
39 def _load_or_create_bridge_token(path: Path) -> str:
40 """Load a persisted bridge token or create one on first use."""
41 if path.exists():
42 token = path.read_text(encoding="utf-8").strip()
43 if token:
44 return token
45
46 path.parent.mkdir(parents=True, exist_ok=True)
47 token = secrets.token_urlsafe(32)
48 path.write_text(token, encoding="utf-8")
49 try:
50 path.chmod(0o600)
51 except OSError:
52 pass
53 return token
54
55
56 class WhatsAppChannel(BaseChannel):
57 """
58 WhatsApp channel that connects to a Node.js bridge.
59
60 The bridge uses @whiskeysockets/baileys to handle the WhatsApp Web protocol.
61 Communication between Python and Node.js is via WebSocket.
62 """
63
64 name = "whatsapp"
65 display_name = "WhatsApp"
66
67 @classmethod
68 def default_config(cls) -> dict[str, Any]:
69 return WhatsAppConfig().model_dump(by_alias=True)
70
71 def __init__(self, config: Any, bus: MessageBus):
72 if isinstance(config, dict):
73 config = WhatsAppConfig.model_validate(config)
74 super().__init__(config, bus)
75 self._ws = None
76 self._connected = False
77 self._processed_message_ids: OrderedDict[str, None] = OrderedDict()
78 self._lid_to_phone: dict[str, str] = {}
79 self._bridge_token: str | None = None
80
81 def _effective_bridge_token(self) -> str:
82 """Resolve the bridge token, generating a local secret when needed."""
83 if self._bridge_token is not None:
84 return self._bridge_token
85 configured = self.config.bridge_token.strip()
86 if configured:
87 self._bridge_token = configured
88 else:
89 self._bridge_token = _load_or_create_bridge_token(_bridge_token_path())
90 return self._bridge_token
91
92 async def login(self, force: bool = False) -> bool:
93 """
94 Set up and run the WhatsApp bridge for QR code login.
95
96 This spawns the Node.js bridge process which handles the WhatsApp
97 authentication flow. The process blocks until the user scans the QR code
98 or interrupts with Ctrl+C.
99 """
100 try:
101 bridge_dir = _ensure_bridge_setup()
102 except RuntimeError as e:
103 logger.error("{}", e)
104 return False
105
106 env = {**os.environ}
107 env["BRIDGE_TOKEN"] = self._effective_bridge_token()
108 env["AUTH_DIR"] = str(_bridge_token_path().parent)
109
110 logger.info("Starting WhatsApp bridge for QR login...")
111 try:
112 subprocess.run(
113 [shutil.which("npm"), "start"], cwd=bridge_dir, check=True, env=env
114 )
115 except subprocess.CalledProcessError:
116 return False
117
118 return True
119
120 async def start(self) -> None:
121 """Start the WhatsApp channel by connecting to the bridge."""
122 import websockets
123
124 bridge_url = self.config.bridge_url
125
126 logger.info("Connecting to WhatsApp bridge at {}...", bridge_url)
127
128 self._running = True
129
130 while self._running:
131 try:
132 async with websockets.connect(bridge_url) as ws:
133 self._ws = ws
134 await ws.send(
135 json.dumps({"type": "auth", "token": self._effective_bridge_token()})
136 )
137 self._connected = True
138 logger.info("Connected to WhatsApp bridge")
139
140 # Listen for messages
141 async for message in ws:
142 try:
143 await self._handle_bridge_message(message)
144 except Exception as e:
145 logger.error("Error handling bridge message: {}", e)
146
147 except asyncio.CancelledError:
148 break
149 except Exception as e:
150 self._connected = False
151 self._ws = None
152 logger.warning("WhatsApp bridge connection error: {}", e)
153
154 if self._running:
155 logger.info("Reconnecting in 5 seconds...")
156 await asyncio.sleep(5)
157
158 async def stop(self) -> None:
159 """Stop the WhatsApp channel."""
160 self._running = False
161 self._connected = False
162
163 if self._ws:
164 await self._ws.close()
165 self._ws = None
166
167 async def send(self, msg: OutboundMessage) -> None:
168 """Send a message through WhatsApp."""
169 if not self._ws or not self._connected:
170 logger.warning("WhatsApp bridge not connected")
171 return
172
173 chat_id = msg.chat_id
174
175 if msg.content:
176 try:
177 payload = {"type": "send", "to": chat_id, "text": msg.content}
178 await self._ws.send(json.dumps(payload, ensure_ascii=False))
179 except Exception as e:
180 logger.error("Error sending WhatsApp message: {}", e)
181 raise
182
183 for media_path in msg.media or []:
184 try:
185 mime, _ = mimetypes.guess_type(media_path)
186 payload = {
187 "type": "send_media",
188 "to": chat_id,
189 "filePath": media_path,
190 "mimetype": mime or "application/octet-stream",
191 "fileName": media_path.rsplit("/", 1)[-1],
192 }
193 await self._ws.send(json.dumps(payload, ensure_ascii=False))
194 except Exception as e:
195 logger.error("Error sending WhatsApp media {}: {}", media_path, e)
196 raise
197
198 async def _handle_bridge_message(self, raw: str) -> None:
199 """Handle a message from the bridge."""
200 try:
201 data = json.loads(raw)
202 except json.JSONDecodeError:
203 logger.warning("Invalid JSON from bridge: {}", raw[:100])
204 return
205
206 msg_type = data.get("type")
207
208 if msg_type == "message":
209 # Incoming message from WhatsApp
210 # Deprecated by whatsapp: old phone number style typically: <phone>@s.whatspp.net
211 pn = data.get("pn", "")
212 # New LID sytle typically:
213 sender = data.get("sender", "")
214 content = data.get("content", "")
215 message_id = data.get("id", "")
216
217 if message_id:
218 if message_id in self._processed_message_ids:
219 return
220 self._processed_message_ids[message_id] = None
221 while len(self._processed_message_ids) > 1000:
222 self._processed_message_ids.popitem(last=False)
223
224 # Extract just the phone number or lid as chat_id
225 is_group = data.get("isGroup", False)
226 was_mentioned = data.get("wasMentioned", False)
227
228 if is_group and getattr(self.config, "group_policy", "open") == "mention":
229 if not was_mentioned:
230 return
231
232 # Classify by JID suffix: @s.whatsapp.net = phone, @lid.whatsapp.net = LID
233 # The bridge's pn/sender fields don't consistently map to phone/LID across versions.
234 raw_a = pn or ""
235 raw_b = sender or ""
236 id_a = raw_a.split("@")[0] if "@" in raw_a else raw_a
237 id_b = raw_b.split("@")[0] if "@" in raw_b else raw_b
238
239 phone_id = ""
240 lid_id = ""
241 for raw, extracted in [(raw_a, id_a), (raw_b, id_b)]:
242 if "@s.whatsapp.net" in raw:
243 phone_id = extracted
244 elif "@lid.whatsapp.net" in raw:
245 lid_id = extracted
246 elif extracted and not phone_id:
247 phone_id = extracted # best guess for bare values
248
249 if phone_id and lid_id:
250 self._lid_to_phone[lid_id] = phone_id
251 sender_id = phone_id or self._lid_to_phone.get(lid_id, "") or lid_id or id_a or id_b
252
253 logger.info("Sender phone={} lid={} → sender_id={}", phone_id or "(empty)", lid_id or "(empty)", sender_id)
254
255 # Extract media paths (images/documents/videos downloaded by the bridge)
256 media_paths = data.get("media") or []
257
258 # Handle voice transcription if it's a voice message
259 if content == "[Voice Message]":
260 if media_paths:
261 logger.info("Transcribing voice message from {}...", sender_id)
262 transcription = await self.transcribe_audio(media_paths[0])
263 if transcription:
264 content = transcription
265 logger.info("Transcribed voice from {}: {}...", sender_id, transcription[:50])
266 else:
267 content = "[Voice Message: Transcription failed]"
268 else:
269 content = "[Voice Message: Audio not available]"
270
271 # Build content tags matching Telegram's pattern: [image: /path] or [file: /path]
272 if media_paths:
273 for p in media_paths:
274 mime, _ = mimetypes.guess_type(p)
275 media_type = "image" if mime and mime.startswith("image/") else "file"
276 media_tag = f"[{media_type}: {p}]"
277 content = f"{content}\n{media_tag}" if content else media_tag
278
279 await self._handle_message(
280 sender_id=sender_id,
281 chat_id=sender, # Use full LID for replies
282 content=content,
283 media=media_paths,
284 metadata={
285 "message_id": message_id,
286 "timestamp": data.get("timestamp"),
287 "is_group": data.get("isGroup", False),
288 },
289 )
290
291 elif msg_type == "status":
292 # Connection status update
293 status = data.get("status")
294 logger.info("WhatsApp status: {}", status)
295
296 if status == "connected":
297 self._connected = True
298 elif status == "disconnected":
299 self._connected = False
300
301 elif msg_type == "qr":
302 # QR code for authentication
303 logger.info("Scan QR code in the bridge terminal to connect WhatsApp")
304
305 elif msg_type == "error":
306 logger.error("WhatsApp bridge error: {}", data.get("error"))
307
308
309 def _ensure_bridge_setup() -> Path:
310 """
311 Ensure the WhatsApp bridge is set up and built.
312
313 Returns the bridge directory. Raises RuntimeError if npm is not found
314 or bridge cannot be built.
315 """
316 from nanobot.config.paths import get_bridge_install_dir
317
318 user_bridge = get_bridge_install_dir()
319
320 if (user_bridge / "dist" / "index.js").exists():
321 return user_bridge
322
323 npm_path = shutil.which("npm")
324 if not npm_path:
325 raise RuntimeError("npm not found. Please install Node.js >= 18.")
326
327 # Find source bridge
328 current_file = Path(__file__)
329 pkg_bridge = current_file.parent.parent / "bridge"
330 src_bridge = current_file.parent.parent.parent / "bridge"
331
332 source = None
333 if (pkg_bridge / "package.json").exists():
334 source = pkg_bridge
335 elif (src_bridge / "package.json").exists():
336 source = src_bridge
337
338 if not source:
339 raise RuntimeError(
340 "WhatsApp bridge source not found. "
341 "Try reinstalling: pip install --force-reinstall nanobot"
342 )
343
344 logger.info("Setting up WhatsApp bridge...")
345 user_bridge.parent.mkdir(parents=True, exist_ok=True)
346 if user_bridge.exists():
347 shutil.rmtree(user_bridge)
348 shutil.copytree(source, user_bridge, ignore=shutil.ignore_patterns("node_modules", "dist"))
349
350 logger.info(" Installing dependencies...")
351 subprocess.run([npm_path, "install"], cwd=user_bridge, check=True, capture_output=True)
352
353 logger.info(" Building...")
354 subprocess.run([npm_path, "run", "build"], cwd=user_bridge, check=True, capture_output=True)
355
356 logger.info("Bridge ready")
357 return user_bridge
358
358 lines PYTHON