whatsapp_rust/
download.rs

1use crate::client::Client;
2use crate::mediaconn::MediaConn;
3use anyhow::{Result, anyhow};
4use std::io::{Seek, SeekFrom, Write};
5
6pub use wacore::download::{DownloadUtils, Downloadable, MediaType};
7
8impl From<&MediaConn> for wacore::download::MediaConnection {
9    fn from(conn: &MediaConn) -> Self {
10        wacore::download::MediaConnection {
11            hosts: conn
12                .hosts
13                .iter()
14                .map(|h| wacore::download::MediaHost {
15                    hostname: h.hostname.clone(),
16                })
17                .collect(),
18            auth: conn.auth.clone(),
19        }
20    }
21}
22
23impl Client {
24    pub async fn download(&self, downloadable: &dyn Downloadable) -> Result<Vec<u8>> {
25        let media_conn = self.refresh_media_conn(false).await?;
26
27        let core_media_conn = wacore::download::MediaConnection::from(&media_conn);
28        let requests = DownloadUtils::prepare_download_requests(downloadable, &core_media_conn)?;
29
30        for request in requests {
31            match self.download_and_decrypt_with_request(&request).await {
32                Ok(data) => return Ok(data),
33                Err(e) => {
34                    log::warn!(
35                        "Failed to download from URL {}: {:?}. Trying next host.",
36                        request.url,
37                        e
38                    );
39                    continue;
40                }
41            }
42        }
43
44        Err(anyhow!("Failed to download from all available media hosts"))
45    }
46
47    async fn download_and_decrypt_with_request(
48        &self,
49        request: &wacore::download::DownloadRequest,
50    ) -> Result<Vec<u8>> {
51        let url = request.url.clone();
52        let media_key = request.media_key.clone();
53        let app_info = request.app_info;
54        tokio::task::spawn_blocking(move || -> Result<Vec<u8>> {
55            let resp = ureq::get(&url).call()?;
56            let mut body = resp.into_body();
57            let reader = body.as_reader();
58            DownloadUtils::decrypt_stream(reader, &media_key, app_info)
59        })
60        .await?
61    }
62
63    pub async fn download_to_file<W: Write + Seek + Send + Unpin>(
64        &self,
65        downloadable: &dyn Downloadable,
66        mut writer: W,
67    ) -> Result<()> {
68        let media_conn = self.refresh_media_conn(false).await?;
69        let core_media_conn = wacore::download::MediaConnection::from(&media_conn);
70        let requests = DownloadUtils::prepare_download_requests(downloadable, &core_media_conn)?;
71        let mut last_err: Option<anyhow::Error> = None;
72        for req in requests {
73            match self
74                .download_and_write(&req.url, &req.media_key, req.app_info, &mut writer)
75                .await
76            {
77                Ok(()) => return Ok(()),
78                Err(e) => {
79                    last_err = Some(e);
80                    continue;
81                }
82            }
83        }
84        Err(last_err.unwrap_or_else(|| anyhow!("All media hosts failed")))
85    }
86
87    async fn download_and_write<W: Write + Seek + Send + Unpin>(
88        &self,
89        url: &str,
90        media_key: &[u8],
91        media_type: MediaType,
92        writer: &mut W,
93    ) -> Result<()> {
94        let url = url.to_string();
95        let media_key = media_key.to_vec();
96
97        let plaintext = tokio::task::spawn_blocking(move || -> Result<Vec<u8>> {
98            let mut resp = ureq::get(&url).call()?;
99            if resp.status().as_u16() >= 300 {
100                return Err(anyhow!("Download failed with status: {}", resp.status()));
101            }
102
103            let encrypted_bytes = resp.body_mut().read_to_vec()?;
104
105            DownloadUtils::verify_and_decrypt(&encrypted_bytes, &media_key, media_type)
106        })
107        .await??;
108
109        writer.seek(SeekFrom::Start(0))?;
110        writer.write_all(&plaintext)?;
111        Ok(())
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    use std::io::Cursor;
119
120    #[test]
121    fn process_downloaded_media_ok() {
122        let data = b"Hello media test";
123        let enc = wacore::upload::encrypt_media(data, MediaType::Image).unwrap();
124        let mut cursor = Cursor::new(Vec::<u8>::new());
125        let plaintext = DownloadUtils::verify_and_decrypt(
126            &enc.data_to_upload,
127            &enc.media_key,
128            MediaType::Image,
129        )
130        .unwrap();
131        cursor.write_all(&plaintext).unwrap();
132        assert_eq!(cursor.into_inner(), data);
133    }
134
135    #[test]
136    fn process_downloaded_media_bad_mac() {
137        let data = b"Tamper";
138        let mut enc = wacore::upload::encrypt_media(data, MediaType::Image).unwrap();
139        let last = enc.data_to_upload.len() - 1;
140        enc.data_to_upload[last] ^= 0x01;
141
142        let err = DownloadUtils::verify_and_decrypt(
143            &enc.data_to_upload,
144            &enc.media_key,
145            MediaType::Image,
146        )
147        .unwrap_err();
148
149        assert!(err.to_string().to_lowercase().contains("invalid mac"));
150    }
151}
whatsapp_rust/download.rs

whatsapp_rust/
download.rs
