Skip to main content

douyin_cli/
auth.rs

1use std::collections::HashMap;
2use std::io::{Read, Write};
3use std::net::{TcpListener, TcpStream};
4use std::thread;
5use std::time::{Duration, Instant};
6
7use base64::Engine;
8use base64::engine::general_purpose::URL_SAFE_NO_PAD;
9use clap::{Args, Subcommand};
10use qrcode::{Color, QrCode};
11use ring::rand::{SecureRandom, SystemRandom};
12use serde_json::{Map, Value, json};
13
14use crate::cookie;
15use crate::err;
16use crate::openapi::{OpenApiClient, RequestSpec};
17use crate::settings;
18
19#[derive(Debug, Args)]
20pub struct AuthArgs {
21    #[command(subcommand)]
22    command: AuthCommand,
23}
24
25#[derive(Debug, Subcommand)]
26enum AuthCommand {
27    /// 通过官方 OAuth 授权接入账号
28    Login {
29        #[arg(long, env = "DOUYIN_CLIENT_KEY")]
30        client_key: Option<String>,
31        #[arg(long, env = "DOUYIN_CLIENT_SECRET")]
32        client_secret: Option<String>,
33        #[arg(long)]
34        redirect_uri: Option<String>,
35        #[arg(long)]
36        scope: Vec<String>,
37        #[arg(long)]
38        code: Option<String>,
39        #[arg(long, conflicts_with = "no_qr")]
40        qr: bool,
41        #[arg(long, conflicts_with = "qr")]
42        no_qr: bool,
43        #[arg(long)]
44        listen: bool,
45        #[arg(long, default_value = "127.0.0.1")]
46        callback_host: String,
47        #[arg(long, default_value_t = 8787, value_parser = clap::value_parser!(u16).range(1..))]
48        callback_port: u16,
49        #[arg(long, default_value_t = 300, value_parser = clap::value_parser!(u64).range(1..=3600))]
50        timeout: u64,
51    },
52    /// 用官方 OAuth code 换取并保存 token
53    Code {
54        #[arg(long)]
55        code: String,
56        #[arg(long, env = "DOUYIN_CLIENT_SECRET")]
57        client_secret: Option<String>,
58    },
59    /// 刷新已保存的官方 access_token
60    Refresh,
61    /// 检查官方授权状态
62    Status {
63        #[arg(long)]
64        json: bool,
65    },
66    /// 删除已保存的官方 OAuth token
67    Logout,
68    /// 保存网页端 Cookie,用于搜索、评论和下载等网页端采集
69    CookieLogin {
70        #[arg(long, env = "DOUYIN_COOKIE")]
71        cookie: String,
72    },
73    /// 检查已保存 Cookie 格式,并尝试确认网页登录态
74    CookieStatus {
75        /// 只检查本地 Cookie 格式,不访问网络
76        #[arg(long)]
77        offline: bool,
78    },
79    /// 删除已保存的网页端 Cookie
80    CookieLogout,
81}
82
83pub fn run(args: AuthArgs) -> Result<(), String> {
84    match args.command {
85        AuthCommand::Login {
86            client_key,
87            client_secret,
88            redirect_uri,
89            scope,
90            code,
91            qr: _,
92            no_qr,
93            listen,
94            callback_host,
95            callback_port,
96            timeout,
97        } => login(LoginOptions {
98            client_key,
99            client_secret,
100            redirect_uri,
101            scopes: scope,
102            code,
103            show_qr: !no_qr,
104            listen,
105            callback_host,
106            callback_port,
107            timeout,
108        }),
109        AuthCommand::Code {
110            code,
111            client_secret,
112        } => exchange_code(&code, client_secret),
113        AuthCommand::Refresh => refresh(),
114        AuthCommand::Status { json } => status(json),
115        AuthCommand::Logout => logout(),
116        AuthCommand::CookieLogin { cookie } => cookie_login(&cookie),
117        AuthCommand::CookieStatus { offline } => cookie_status(offline),
118        AuthCommand::CookieLogout => cookie_logout(),
119    }
120}
121
122struct LoginOptions {
123    client_key: Option<String>,
124    client_secret: Option<String>,
125    redirect_uri: Option<String>,
126    scopes: Vec<String>,
127    code: Option<String>,
128    show_qr: bool,
129    listen: bool,
130    callback_host: String,
131    callback_port: u16,
132    timeout: u64,
133}
134
135fn login(options: LoginOptions) -> Result<(), String> {
136    let data = settings::load().map_err(err)?;
137    let saved = settings::openapi(&data);
138    let client_key = options
139        .client_key
140        .or_else(|| saved_string(&saved, "clientKey"))
141        .ok_or_else(|| missing_client_key(options.show_qr))?;
142    let client_secret = options
143        .client_secret
144        .or_else(|| saved_string(&saved, "clientSecret"));
145    let mut redirect_uri = options
146        .redirect_uri
147        .or_else(|| saved_string(&saved, "redirectUri"));
148    let scopes = if options.scopes.is_empty() {
149        saved
150            .get("scopes")
151            .and_then(Value::as_array)
152            .map(|values| {
153                values
154                    .iter()
155                    .filter_map(Value::as_str)
156                    .map(str::to_owned)
157                    .collect()
158            })
159            .filter(|values: &Vec<String>| !values.is_empty())
160            .unwrap_or_else(|| vec!["user_info".to_owned()])
161    } else {
162        options.scopes
163    };
164    if options.listen {
165        redirect_uri = Some(format!(
166            "http://{}:{}/callback",
167            options.callback_host, options.callback_port
168        ));
169    }
170    let redirect_uri =
171        redirect_uri.ok_or_else(|| "缺少 redirect_uri,请传入 --redirect-uri".to_owned())?;
172    let state = (options.listen && options.code.is_none())
173        .then(random_state)
174        .transpose()?;
175    let client = OpenApiClient::new()?;
176    let url = client.authorize_url(&client_key, &redirect_uri, &scopes, state.as_deref())?;
177    println!("请在浏览器打开以下官方授权链接:\n{url}");
178    if options.show_qr {
179        print_qr(&url)?;
180    }
181
182    let mut code = options.code;
183    if options.listen && code.is_none() {
184        println!("正在等待授权回调: {redirect_uri}");
185        code = Some(wait_for_code(
186            &options.callback_host,
187            options.callback_port,
188            state.as_deref(),
189            Duration::from_secs(options.timeout),
190        )?);
191    }
192    let mut updates = Map::from_iter([
193        ("clientKey".to_owned(), json!(client_key)),
194        (
195            "clientSecret".to_owned(),
196            json!(client_secret.clone().unwrap_or_default()),
197        ),
198        ("redirectUri".to_owned(), json!(redirect_uri)),
199        ("scopes".to_owned(), json!(scopes)),
200    ]);
201    if let Some(code) = code {
202        let secret =
203            client_secret.ok_or_else(|| "使用 code 换 token 需要 --client-secret".to_owned())?;
204        let response = client.access_token(&client_key, &secret, &code)?;
205        updates.extend(extract_token_fields(&response));
206        print_json(&response)?;
207    } else {
208        println!("授权完成后运行:douyin auth code --code 授权码");
209    }
210    save_openapi(updates)?;
211    println!(
212        "官方授权配置已保存: {}",
213        settings::settings_file().display()
214    );
215    Ok(())
216}
217
218fn exchange_code(code: &str, client_secret: Option<String>) -> Result<(), String> {
219    let data = settings::load().map_err(err)?;
220    let saved = settings::openapi(&data);
221    let client_key = saved_string(&saved, "clientKey")
222        .ok_or_else(|| "缺少 client_key,请先运行 douyin auth login".to_owned())?;
223    let secret = client_secret
224        .or_else(|| saved_string(&saved, "clientSecret"))
225        .ok_or_else(|| "缺少 client_secret,请传入 --client-secret".to_owned())?;
226    let response = OpenApiClient::new()?.access_token(&client_key, &secret, code)?;
227    let mut updates = extract_token_fields(&response);
228    updates.insert("clientSecret".to_owned(), json!(secret));
229    save_openapi(updates)?;
230    print_json(&response)?;
231    println!("官方 token 已保存: {}", settings::settings_file().display());
232    Ok(())
233}
234
235fn refresh() -> Result<(), String> {
236    let data = settings::load().map_err(err)?;
237    let saved = settings::openapi(&data);
238    let client_key = saved_string(&saved, "clientKey");
239    let refresh_token = saved_string(&saved, "refreshToken");
240    let (Some(client_key), Some(refresh_token)) = (client_key, refresh_token) else {
241        return Err("缺少 client_key 或 refresh_token,请重新授权".to_owned());
242    };
243    let response = OpenApiClient::new()?.refresh_token(&client_key, &refresh_token)?;
244    save_openapi(extract_token_fields(&response))?;
245    print_json(&response)?;
246    println!("官方 token 已刷新");
247    Ok(())
248}
249
250fn status(json_output: bool) -> Result<(), String> {
251    let data = settings::load().map_err(err)?;
252    let saved = settings::openapi(&data);
253    let token = saved_string(&saved, "accessToken");
254    let open_id = saved_string(&saved, "openId");
255    let authorized = token.is_some() && open_id.is_some();
256    let mut output = json!({
257        "authorized": authorized,
258        "connected": false,
259        "configFile": settings::settings_file(),
260        "openId": open_id.clone().unwrap_or_default(),
261        "scopes": saved.get("scopes").cloned().unwrap_or_else(|| json!([]))
262    });
263    let (Some(token), Some(open_id)) = (token, open_id) else {
264        if json_output {
265            return print_json(&output);
266        }
267        println!("未完成官方授权");
268        return Ok(());
269    };
270    if !json_output {
271        println!("已保存官方授权: {}", settings::settings_file().display());
272        println!("open_id: {open_id}");
273        println!("正在检查官方 OpenAPI 连通性...");
274    }
275    match OpenApiClient::new()?.request(RequestSpec {
276        method: "GET",
277        path: "/oauth/userinfo/",
278        token: Some(&token),
279        params: Some(HashMap::from([("open_id".to_owned(), open_id)])),
280        auth_required: true,
281        ..RequestSpec::default()
282    }) {
283        Ok(userinfo) => {
284            output["connected"] = json!(true);
285            output["userinfo"] = userinfo.clone();
286            print_json(if json_output { &output } else { &userinfo })
287        }
288        Err(error) if json_output => {
289            output["error"] = json!(error);
290            print_json(&output)?;
291            Err("官方 OpenAPI 连通性检查失败".to_owned())
292        }
293        Err(error) => Err(format!("官方 OpenAPI 连通性检查失败: {error}")),
294    }
295}
296
297fn logout() -> Result<(), String> {
298    save_openapi(Map::from_iter([
299        ("accessToken".to_owned(), json!("")),
300        ("refreshToken".to_owned(), json!("")),
301        ("openId".to_owned(), json!("")),
302        ("expiresIn".to_owned(), json!(0)),
303    ]))?;
304    println!("已清除官方授权 token");
305    Ok(())
306}
307
308fn cookie_login(value: &str) -> Result<(), String> {
309    let value = value.trim();
310    if !cookie::validate(value) {
311        return Err("Cookie 格式校验失败,未保存".to_owned());
312    }
313    let mut data = settings::load().map_err(err)?;
314    data["cookie"] = json!(value);
315    settings::save(&data).map_err(err)?;
316    println!("Cookie 已保存: {}", settings::settings_file().display());
317    Ok(())
318}
319
320fn cookie_status(offline: bool) -> Result<(), String> {
321    let data = settings::load().map_err(err)?;
322    let value = data
323        .get("cookie")
324        .and_then(Value::as_str)
325        .unwrap_or("")
326        .trim();
327    if value.is_empty() {
328        println!("未保存 Cookie");
329        return Ok(());
330    }
331    if !cookie::validate(value) {
332        return Err(format!(
333            "已保存 Cookie,但格式无效: {}",
334            settings::settings_file().display()
335        ));
336    }
337    if offline {
338        println!("Cookie 格式有效: {}", settings::settings_file().display());
339        return Ok(());
340    }
341    println!("正在确认网页登录态...");
342    match cookie::probe(value) {
343        Ok(true) => {
344            println!(
345                "Cookie 网页登录态有效: {}",
346                settings::settings_file().display()
347            );
348            Ok(())
349        }
350        Ok(false) => Err("Cookie 已保存,但网页登录态无效或已过期".to_owned()),
351        Err(error) => Err(format!(
352            "Cookie 已保存且格式有效,但无法确认网页登录态: {error}\n可运行 douyin auth cookie-status --offline 仅检查本地格式"
353        )),
354    }
355}
356
357fn cookie_logout() -> Result<(), String> {
358    let mut data = settings::load().map_err(err)?;
359    data["cookie"] = json!("");
360    settings::save(&data).map_err(err)?;
361    println!("已清除 Cookie");
362    Ok(())
363}
364
365fn save_openapi(updates: Map<String, Value>) -> Result<(), String> {
366    let mut data = settings::load().map_err(err)?;
367    let openapi = data
368        .get_mut("openapi")
369        .and_then(Value::as_object_mut)
370        .ok_or_else(|| "openapi 配置格式无效".to_owned())?;
371    openapi.extend(updates);
372    settings::save(&data).map_err(err)
373}
374
375fn extract_token_fields(data: &Value) -> Map<String, Value> {
376    let source = data
377        .get("data")
378        .filter(|value| value.is_object())
379        .unwrap_or(data);
380    [
381        ("access_token", "accessToken"),
382        ("refresh_token", "refreshToken"),
383        ("open_id", "openId"),
384        ("expires_in", "expiresIn"),
385    ]
386    .into_iter()
387    .filter_map(|(source_key, target_key)| {
388        source
389            .get(source_key)
390            .cloned()
391            .map(|value| (target_key.to_owned(), value))
392    })
393    .collect()
394}
395
396fn saved_string(values: &Map<String, Value>, key: &str) -> Option<String> {
397    values
398        .get(key)
399        .and_then(Value::as_str)
400        .filter(|value| !value.is_empty())
401        .map(str::to_owned)
402}
403
404fn missing_client_key(show_qr: bool) -> String {
405    let qr_line =
406        show_qr.then_some("--qr 只会把官方 OAuth 授权链接渲染成二维码,仍然需要 client_key。\n");
407    format!(
408        "当前命令是官方 OpenAPI OAuth 授权,需要开放平台 client_key。\n{}这不是网页端 Cookie 扫码登录,不能直接生成可保存 Cookie 的登录二维码。\n\n可选方案:\n  1. 官方 OpenAPI:传入 --client-key,或设置 DOUYIN_CLIENT_KEY。\n  2. 网页端采集:从浏览器复制 Cookie 后运行:\n     douyin auth cookie-login --cookie 'sessionid=...; ttwid=...'",
409        qr_line.unwrap_or("")
410    )
411}
412
413fn random_state() -> Result<String, String> {
414    let mut bytes = [0_u8; 18];
415    SystemRandom::new()
416        .fill(&mut bytes)
417        .map_err(|_| "无法生成 OAuth state".to_owned())?;
418    Ok(URL_SAFE_NO_PAD.encode(bytes))
419}
420
421fn print_qr(value: &str) -> Result<(), String> {
422    let code = QrCode::new(value.as_bytes()).map_err(err)?;
423    let width = code.width();
424    println!();
425    for y in (0..width).step_by(2) {
426        let mut line = String::new();
427        for x in 0..width {
428            let top = code[(x, y)] == Color::Dark;
429            let bottom = y + 1 < width && code[(x, y + 1)] == Color::Dark;
430            line.push(match (top, bottom) {
431                (true, true) => '█',
432                (true, false) => '▀',
433                (false, true) => '▄',
434                (false, false) => ' ',
435            });
436        }
437        println!(" {line} ");
438    }
439    println!();
440    Ok(())
441}
442
443fn wait_for_code(
444    host: &str,
445    port: u16,
446    expected_state: Option<&str>,
447    timeout: Duration,
448) -> Result<String, String> {
449    let listener = TcpListener::bind((host, port))
450        .map_err(|_| format!("无法监听 {host}:{port},请换一个 --callback-port"))?;
451    listener.set_nonblocking(true).map_err(err)?;
452    let started = Instant::now();
453    while started.elapsed() < timeout {
454        match listener.accept() {
455            Ok((mut stream, _)) => return handle_callback(&mut stream, expected_state),
456            Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
457                thread::sleep(Duration::from_millis(50));
458            }
459            Err(error) => return Err(error.to_string()),
460        }
461    }
462    Err("等待授权回调超时,未获取到 code".to_owned())
463}
464
465const MAX_CALLBACK_REQUEST_BYTES: usize = 16 * 1024;
466
467fn handle_callback(stream: &mut TcpStream, expected_state: Option<&str>) -> Result<String, String> {
468    stream
469        .set_read_timeout(Some(Duration::from_secs(5)))
470        .map_err(err)?;
471    let request = read_request_head(stream)?;
472    let target = request
473        .lines()
474        .next()
475        .and_then(|line| line.split_whitespace().nth(1))
476        .ok_or_else(|| "授权回调请求无效".to_owned())?;
477    let url = reqwest::Url::parse(&format!("http://localhost{target}"))
478        .map_err(|error| format!("授权回调 URL 无效: {error}"))?;
479    if url.path() != "/callback" {
480        send_html(stream, 404, "未找到回调路径")?;
481        return Err("授权回调路径无效".to_owned());
482    }
483    let params: HashMap<_, _> = url.query_pairs().into_owned().collect();
484    if let Some(error) = params
485        .get("error")
486        .or_else(|| params.get("error_description"))
487    {
488        send_html(stream, 400, "授权失败,可以关闭此页面并返回终端。")?;
489        return Err(format!("授权失败: {error}"));
490    }
491    if expected_state
492        .is_some_and(|expected| params.get("state").map(String::as_str) != Some(expected))
493    {
494        send_html(stream, 400, "state 不匹配。")?;
495        return Err("授权回调 state 不匹配,已拒绝".to_owned());
496    }
497    let Some(code) = params.get("code") else {
498        send_html(stream, 400, "回调缺少 code。")?;
499        return Err("授权回调缺少 code".to_owned());
500    };
501    send_html(stream, 200, "授权完成,可以关闭此页面并返回终端。")?;
502    Ok(code.to_owned())
503}
504
505/// Reads until the end of the HTTP request headers (`\r\n\r\n`) instead of assuming
506/// they arrive in a single `read`, which TCP does not guarantee.
507fn read_request_head(stream: &mut TcpStream) -> Result<String, String> {
508    let mut buffer = Vec::new();
509    let mut chunk = [0_u8; 4096];
510    while !buffer.windows(4).any(|window| window == b"\r\n\r\n") {
511        if buffer.len() >= MAX_CALLBACK_REQUEST_BYTES {
512            return Err("授权回调请求过大".to_owned());
513        }
514        let count = stream.read(&mut chunk).map_err(err)?;
515        if count == 0 {
516            break;
517        }
518        buffer.extend_from_slice(&chunk[..count]);
519    }
520    Ok(String::from_utf8_lossy(&buffer).into_owned())
521}
522
523fn send_html(stream: &mut TcpStream, status: u16, body: &str) -> Result<(), String> {
524    let content =
525        format!("<!doctype html><meta charset='utf-8'><title>Douyin CLI</title><p>{body}</p>");
526    let reason = if status == 200 { "OK" } else { "Bad Request" };
527    write!(
528        stream,
529        "HTTP/1.1 {status} {reason}\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{content}",
530        content.len()
531    )
532    .map_err(err)
533}
534
535fn print_json(value: &Value) -> Result<(), String> {
536    println!("{}", serde_json::to_string_pretty(value).map_err(err)?);
537    Ok(())
538}
539
540#[cfg(test)]
541mod tests {
542    use super::{extract_token_fields, missing_client_key};
543    use serde_json::json;
544
545    #[test]
546    fn extracts_nested_token_fields() {
547        let fields = extract_token_fields(&json!({"data": {
548            "access_token": "access", "refresh_token": "refresh", "open_id": "open", "expires_in": 1
549        }}));
550        assert_eq!(fields["accessToken"], "access");
551        assert_eq!(fields["openId"], "open");
552    }
553
554    #[test]
555    fn missing_key_message_explains_cookie_alternative() {
556        let message = missing_client_key(true);
557        assert!(message.contains("官方 OpenAPI OAuth 授权"));
558        assert!(message.contains("--qr 只会把官方 OAuth 授权链接渲染成二维码"));
559        assert!(message.contains("douyin auth cookie-login --cookie"));
560    }
561}