1use std::{
2 future::Future,
3 path::Path,
4 pin::Pin,
5 sync::{Arc, Mutex},
6};
7
8use sha2::{Digest, Sha256};
9use tokio::sync::watch;
10
11use crate::{
12 api::{WeixinApiOptions, notify_start, notify_stop},
13 auth::{
14 accounts::{CDN_BASE_URL, DEFAULT_BASE_URL, resolve_weixin_account},
15 login_qr::{display_qr_code, start_weixin_login_with_qr, wait_for_weixin_login},
16 },
17 messaging::{
18 process_message::MessageHandler,
19 send::{
20 SendResult, WeixinMsgContext, get_context_token, restore_context_tokens,
21 send_message_weixin,
22 },
23 send_media::{send_media_url, send_weixin_media_file},
24 },
25 monitor::{MonitorWeixinOpts, monitor_weixin_provider},
26};
27
28#[derive(Clone, Debug)]
29pub struct WeixinBotOptions {
30 pub token: String,
31 pub base_url: Option<String>,
32 pub cdn_base_url: Option<String>,
33 pub state_dir: Option<String>,
34 pub account_id: Option<String>,
35 pub user_id: Option<String>,
36}
37
38pub struct StartOptions {
39 pub on_message: MessageHandler,
40 pub long_poll_timeout_ms: Option<u64>,
41}
42
43#[derive(Clone)]
44pub struct WeixinBot {
45 token: String,
46 base_url: String,
47 cdn_base_url: String,
48 account_id: String,
49 user_id: Option<String>,
50 stop_tx: Arc<Mutex<Option<watch::Sender<bool>>>>,
51}
52
53impl WeixinBot {
54 pub fn new(opts: WeixinBotOptions) -> Self {
55 let _state_dir = opts.state_dir;
56 let account_id = opts
57 .account_id
58 .unwrap_or_else(|| derive_account_id(&opts.token));
59 restore_context_tokens(&account_id);
60 Self {
61 token: opts.token,
62 base_url: opts.base_url.unwrap_or_else(|| DEFAULT_BASE_URL.into()),
63 cdn_base_url: opts.cdn_base_url.unwrap_or_else(|| CDN_BASE_URL.into()),
64 account_id,
65 user_id: opts.user_id,
66 stop_tx: Arc::new(Mutex::new(None)),
67 }
68 }
69
70 pub fn from_account(account_id: &str) -> crate::Result<Self> {
71 let account = resolve_weixin_account(account_id)?;
72 let token = account.token.ok_or("account is not configured")?;
73 Ok(Self::new(WeixinBotOptions {
74 token,
75 base_url: Some(account.base_url),
76 cdn_base_url: Some(account.cdn_base_url),
77 state_dir: None,
78 account_id: Some(account.account_id),
79 user_id: account.user_id,
80 }))
81 }
82
83 pub async fn login_interactive(api_base_url: Option<&str>) -> crate::Result<Self> {
84 let api_base_url = api_base_url.unwrap_or(DEFAULT_BASE_URL);
85 let start = start_weixin_login_with_qr(api_base_url, None, None, false).await?;
86 if let Some(url) = &start.qrcode_url {
87 display_qr_code(url)?;
88 }
89 let waited = wait_for_weixin_login(&start.session_key, api_base_url, None, None).await?;
90 if !waited.connected {
91 return Err(waited.message.into());
92 }
93 Ok(Self::new(WeixinBotOptions {
94 token: waited.bot_token.ok_or("login returned no token")?,
95 base_url: waited.base_url,
96 cdn_base_url: Some(CDN_BASE_URL.into()),
97 state_dir: None,
98 account_id: waited.account_id,
99 user_id: waited.user_id,
100 }))
101 }
102
103 pub async fn start(&self, opts: StartOptions) -> crate::Result<()> {
104 if self.is_running() {
105 return Ok(());
106 }
107 let (tx, rx) = watch::channel(false);
108 *self.stop_tx.lock().unwrap() = Some(tx);
109 let api_opts = self.api_opts();
110 let _ = notify_start(&api_opts).await;
111 monitor_weixin_provider(
112 MonitorWeixinOpts {
113 base_url: self.base_url.clone(),
114 cdn_base_url: self.cdn_base_url.clone(),
115 token: Some(self.token.clone()),
116 account_id: self.account_id.clone(),
117 long_poll_timeout_ms: opts.long_poll_timeout_ms,
118 on_message: opts.on_message,
119 },
120 rx,
121 )
122 .await
123 }
124
125 pub async fn stop(&self) -> crate::Result<()> {
126 if let Some(tx) = self.stop_tx.lock().unwrap().take() {
127 let _ = tx.send(true);
128 }
129 let _ = notify_stop(&self.api_opts()).await;
130 Ok(())
131 }
132
133 pub fn is_running(&self) -> bool {
134 self.stop_tx.lock().unwrap().is_some()
135 }
136 pub fn account_id(&self) -> &str {
137 &self.account_id
138 }
139 pub fn token(&self) -> &str {
140 &self.token
141 }
142 pub fn user_id(&self) -> Option<&str> {
143 self.user_id.as_deref()
144 }
145 pub fn base_url(&self) -> &str {
146 &self.base_url
147 }
148 fn api_opts(&self) -> WeixinApiOptions {
149 WeixinApiOptions::new(self.base_url.clone(), Some(self.token.clone()))
150 }
151 fn context_token(&self, to: &str) -> Option<String> {
152 get_context_token(&self.account_id, to)
153 }
154
155 pub async fn send_text(&self, to: &str, text: &str) -> crate::Result<SendResult> {
156 send_message_weixin(
157 to,
158 text,
159 &self.api_opts(),
160 self.context_token(to).as_deref(),
161 )
162 .await
163 }
164 pub async fn send_image(
165 &self,
166 to: &str,
167 file_path: impl AsRef<Path>,
168 caption: Option<&str>,
169 ) -> crate::Result<SendResult> {
170 send_weixin_media_file(
171 file_path,
172 to,
173 caption.unwrap_or(""),
174 &self.api_opts(),
175 &self.cdn_base_url,
176 self.context_token(to).as_deref(),
177 )
178 .await
179 }
180 pub async fn send_video(
181 &self,
182 to: &str,
183 file_path: impl AsRef<Path>,
184 caption: Option<&str>,
185 ) -> crate::Result<SendResult> {
186 send_weixin_media_file(
187 file_path,
188 to,
189 caption.unwrap_or(""),
190 &self.api_opts(),
191 &self.cdn_base_url,
192 self.context_token(to).as_deref(),
193 )
194 .await
195 }
196 pub async fn send_file(
197 &self,
198 to: &str,
199 file_path: impl AsRef<Path>,
200 caption: Option<&str>,
201 ) -> crate::Result<SendResult> {
202 send_weixin_media_file(
203 file_path,
204 to,
205 caption.unwrap_or(""),
206 &self.api_opts(),
207 &self.cdn_base_url,
208 self.context_token(to).as_deref(),
209 )
210 .await
211 }
212 pub async fn send_media_url(
213 &self,
214 to: &str,
215 url: &str,
216 caption: Option<&str>,
217 ) -> crate::Result<SendResult> {
218 send_media_url(
219 url,
220 to,
221 caption.unwrap_or(""),
222 &self.api_opts(),
223 &self.cdn_base_url,
224 self.context_token(to).as_deref(),
225 )
226 .await
227 }
228}
229
230pub fn handler<F, Fut>(f: F) -> MessageHandler
231where
232 F: Fn(WeixinMsgContext) -> Fut + Send + Sync + 'static,
233 Fut: Future<Output = crate::Result<Option<String>>> + Send + 'static,
234{
235 Arc::new(move |ctx| {
236 Box::pin(f(ctx)) as Pin<Box<dyn Future<Output = crate::Result<Option<String>>> + Send>>
237 })
238}
239
240fn derive_account_id(token: &str) -> String {
241 let digest = Sha256::digest(token.as_bytes());
242 format!("bot-{}", hex::encode(&digest[..8]))
243}