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        let http_request = crate::http::HttpRequest::get(url);
55        let response = self.http_client.execute(http_request).await?;
56
57        if response.status_code >= 300 {
58            return Err(anyhow!(
59                "Download failed with status: {}",
60                response.status_code
61            ));
62        }
63
64        // Decrypt in a blocking thread since it's CPU-intensive
65        tokio::task::spawn_blocking(move || {
66            DownloadUtils::decrypt_stream(&response.body[..], &media_key, app_info)
67        })
68        .await?
69    }
70
71    pub async fn download_to_file<W: Write + Seek + Send + Unpin>(
72        &self,
73        downloadable: &dyn Downloadable,
74        mut writer: W,
75    ) -> Result<()> {
76        let media_conn = self.refresh_media_conn(false).await?;
77        let core_media_conn = wacore::download::MediaConnection::from(&media_conn);
78        let requests = DownloadUtils::prepare_download_requests(downloadable, &core_media_conn)?;
79        let mut last_err: Option<anyhow::Error> = None;
80        for req in requests {
81            match self
82                .download_and_write(&req.url, &req.media_key, req.app_info, &mut writer)
83                .await
84            {
85                Ok(()) => return Ok(()),
86                Err(e) => {
87                    last_err = Some(e);
88                    continue;
89                }
90            }
91        }
92        match last_err {
93            Some(err) => Err(err),
94            None => Err(anyhow!("Failed to download from all available media hosts")),
95        }
96    }
97
98    async fn download_and_write<W: Write + Seek + Send + Unpin>(
99        &self,
100        url: &str,
101        media_key: &[u8],
102        media_type: MediaType,
103        writer: &mut W,
104    ) -> Result<()> {
105        let http_request = crate::http::HttpRequest::get(url);
106        let response = self.http_client.execute(http_request).await?;
107
108        if response.status_code >= 300 {
109            return Err(anyhow!(
110                "Download failed with status: {}",
111                response.status_code
112            ));
113        }
114
115        let media_key = media_key.to_vec();
116        let encrypted_bytes = response.body;
117
118        // Decrypt and verify in a blocking thread since it's CPU-intensive
119        let plaintext = tokio::task::spawn_blocking(move || {
120            DownloadUtils::verify_and_decrypt(&encrypted_bytes, &media_key, media_type)
121        })
122        .await??;
123
124        writer.seek(SeekFrom::Start(0))?;
125        writer.write_all(&plaintext)?;
126        Ok(())
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133    use std::io::Cursor;
134
135    #[test]
136    fn process_downloaded_media_ok() {
137        let data = b"Hello media test";
138        let enc = wacore::upload::encrypt_media(data, MediaType::Image)
139            .expect("encryption should succeed");
140        let mut cursor = Cursor::new(Vec::<u8>::new());
141        let plaintext = DownloadUtils::verify_and_decrypt(
142            &enc.data_to_upload,
143            &enc.media_key,
144            MediaType::Image,
145        )
146        .expect("decryption should succeed");
147        cursor.write_all(&plaintext).expect("write should succeed");
148        assert_eq!(cursor.into_inner(), data);
149    }
150
151    #[test]
152    fn process_downloaded_media_bad_mac() {
153        let data = b"Tamper";
154        let mut enc = wacore::upload::encrypt_media(data, MediaType::Image)
155            .expect("encryption should succeed");
156        let last = enc.data_to_upload.len() - 1;
157        enc.data_to_upload[last] ^= 0x01;
158
159        let err = DownloadUtils::verify_and_decrypt(
160            &enc.data_to_upload,
161            &enc.media_key,
162            MediaType::Image,
163        )
164        .unwrap_err();
165
166        assert!(err.to_string().to_lowercase().contains("invalid mac"));
167    }
168}
whatsapp_rust/download.rs

whatsapp_rust/
download.rs
