Skip to main content

dingtalk_stream/client/stream_/
download_resources.rs

1use crate::frames::down_message::callback_message::{PayloadFile, PayloadPicture, PayloadVideo};
2use crate::DingTalkStream;
3use anyhow::anyhow;
4use async_trait::async_trait;
5use log::info;
6use serde_json::json;
7use std::ops::Deref;
8use std::path::PathBuf;
9use url::Url;
10
11#[async_trait]
12pub trait DingtalkResource {
13    type T;
14    async fn fetch(
15        &self,
16        dingtalk: &DingTalkStream,
17        save_to: PathBuf,
18    ) -> crate::Result<(PathBuf, Self::T)>;
19}
20
21#[cfg(feature = "image")]
22type Picture = image::DynamicImage;
23
24#[cfg(not(feature = "image"))]
25type Picture = Vec<u8>;
26
27#[async_trait]
28impl DingtalkResource for PayloadPicture {
29    type T = Picture;
30
31    async fn fetch(
32        &self,
33        dingtalk: &DingTalkStream,
34        save_to_dir: PathBuf,
35    ) -> crate::Result<(PathBuf, Self::T)> {
36        if !save_to_dir.exists() {
37            tokio::fs::create_dir_all(&save_to_dir).await?;
38        }
39        if save_to_dir.is_file() {
40            return Err(anyhow!("save_to_dir is a file"));
41        }
42        let filepath = save_to_dir.join(format!(
43            "{}.png",
44            format!("{:x}", md5::compute(&self.download_code))
45        ));
46        if filepath.exists() {
47            let bytes = tokio::fs::read(&filepath).await?;
48            #[cfg(feature = "image")]
49            return Ok((filepath, image::load_from_memory(&bytes)?));
50            #[cfg(not(feature = "image"))]
51            return Ok((filepath, bytes));
52        }
53        let download_url = fetch_download_url(dingtalk, &self.download_code).await?;
54        let bytes = dingtalk
55            .http_client
56            .get(download_url)
57            .send()
58            .await?
59            .bytes()
60            .await?;
61
62        #[cfg(feature = "image")]
63        let image = {
64            use std::io::Cursor;
65            let image = image::load_from_memory(&bytes)?;
66            let mut cursor = Cursor::new(vec![]);
67            image.write_to(&mut cursor, image::ImageFormat::Png)?;
68            tokio::fs::write(&filepath, cursor.into_inner()).await?;
69            image
70        };
71        #[cfg(not(feature = "image"))]
72        let image = {
73            tokio::fs::write(&filepath, bytes.as_ref()).await?;
74            bytes.to_vec()
75        };
76        info!("Downloaded image to {}", filepath.display());
77        Ok((filepath, image))
78    }
79}
80
81#[async_trait]
82impl DingtalkResource for PayloadVideo {
83    type T = Vec<u8>;
84
85    async fn fetch(
86        &self,
87        dingtalk: &DingTalkStream,
88        save_to_dir: PathBuf,
89    ) -> crate::Result<(PathBuf, Self::T)> {
90        if !save_to_dir.exists() {
91            tokio::fs::create_dir_all(&save_to_dir).await?;
92        }
93        if save_to_dir.is_file() {
94            return Err(anyhow!("save_to_dir is a file"));
95        }
96        let filepath = save_to_dir.join(format!(
97            "{}.{}",
98            format!("{:x}", md5::compute(&self.download_code)),
99            self.video_type
100        ));
101        if filepath.exists() {
102            let bytes = tokio::fs::read(&filepath).await?;
103            return Ok((filepath, bytes));
104        }
105        let download_url = fetch_download_url(dingtalk, &self.download_code).await?;
106        let bytes = dingtalk
107            .http_client
108            .get(download_url)
109            .send()
110            .await?
111            .bytes()
112            .await?;
113        tokio::fs::write(&filepath, bytes.as_ref()).await?;
114        info!(
115            "Downloaded {} video to {}",
116            self.download_code,
117            filepath.display()
118        );
119        Ok((filepath, bytes.to_vec()))
120    }
121}
122
123#[async_trait]
124impl DingtalkResource for PayloadFile {
125    type T = Vec<u8>;
126
127    async fn fetch(
128        &self,
129        dingtalk: &DingTalkStream,
130        save_to_dir: PathBuf,
131    ) -> crate::Result<(PathBuf, Self::T)> {
132        if !save_to_dir.exists() {
133            tokio::fs::create_dir_all(&save_to_dir).await?;
134        }
135        if save_to_dir.is_file() {
136            return Err(anyhow!("save_to_dir is a file"));
137        }
138        let filepath = save_to_dir.join(format!(
139            "{}_{}",
140            format!("{:x}", md5::compute(&self.download_code)),
141            self.file_name
142        ));
143        if filepath.exists() {
144            let bytes = tokio::fs::read(&filepath).await?;
145            return Ok((filepath, bytes));
146        }
147        let download_url = fetch_download_url(dingtalk, &self.download_code).await?;
148        let bytes = dingtalk
149            .http_client
150            .get(download_url)
151            .send()
152            .await?
153            .bytes()
154            .await?;
155        tokio::fs::write(&filepath, bytes.as_ref()).await?;
156        info!(
157            "Downloaded {} file to {}",
158            self.download_code,
159            filepath.display()
160        );
161        Ok((filepath, bytes.to_vec()))
162    }
163}
164
165async fn fetch_download_url(dingtalk: &DingTalkStream, download_code: &str) -> crate::Result<Url> {
166    let access_token = dingtalk.get_access_token().await?;
167    let response = dingtalk
168        .http_client
169        .post(crate::MESSAGE_FILES_DOWNLOAD_URL)
170        .header("x-acs-dingtalk-access-token", access_token.deref())
171        .header("Content-Type", "application/json")
172        .json(&json!({
173            "robotCode": &dingtalk.credential.client_id,
174            "downloadCode": download_code,
175        }))
176        .send()
177        .await?;
178    let code = response.status();
179    if response.status().is_success() {
180        let json = response.json::<serde_json::Value>().await?;
181        let download_url = json.get("downloadUrl").and_then(|it| it.as_str());
182        info!(
183            "Get download url by download_code: {download_code}, download_url: {}",
184            download_url.unwrap_or("None")
185        );
186        Ok(download_url
187            .ok_or(anyhow!("download_url is not found"))?
188            .try_into()?)
189    } else {
190        Err(anyhow!(
191            "Failed to download file with unexpected http-code: {}, download_code: {}",
192            code,
193            download_code
194        ))
195    }
196}