modde_sources/gdrive/
mod.rs

1use std::collections::HashMap;
2use std::path::Path;
3
4use anyhow::{Context, Result};
5use reqwest::Client;
6use tracing::debug;
7
8use modde_core::manifest::wabbajack::DownloadDirective;
9
10use crate::common::{ensure_parent, stream_to_file, verify_and_wrap};
11use crate::traits::{DownloadHandle, DownloadSource, ProgressCallback, VerifiedFile};
12
13/// Google Drive download source.
14///
15/// Handles the virus scan warning page for large files.
16pub struct GoogleDriveSource {
17    client: Client,
18}
19
20impl GoogleDriveSource {
21    pub fn new(client: Client) -> Self {
22        Self { client }
23    }
24}
25
26impl DownloadSource for GoogleDriveSource {
27    fn can_handle(&self, directive: &DownloadDirective) -> bool {
28        matches!(directive, DownloadDirective::GoogleDrive { .. })
29    }
30
31    async fn resolve(&self, directive: &DownloadDirective) -> Result<DownloadHandle> {
32        let DownloadDirective::GoogleDrive { id, hash } = directive else {
33            anyhow::bail!("not a Google Drive directive");
34        };
35
36        let url = format!("https://drive.google.com/uc?id={id}&export=download");
37
38        Ok(DownloadHandle {
39            url,
40            headers: HashMap::new(),
41            expected_hash: *hash,
42            size_hint: None,
43        })
44    }
45
46    async fn download_with_progress(
47        &self,
48        handle: DownloadHandle,
49        dest: &Path,
50        progress: ProgressCallback,
51    ) -> Result<VerifiedFile> {
52        ensure_parent(dest).await?;
53
54        do_download(&self.client, &handle, dest, &progress)
55            .await
56            .context("Google Drive download failed")?;
57
58        verify_and_wrap(dest, handle.expected_hash).await
59    }
60}
61
62async fn do_download(
63    client: &Client,
64    handle: &DownloadHandle,
65    dest: &Path,
66    progress: &ProgressCallback,
67) -> Result<()> {
68    let resp = client.get(&handle.url).send().await?.error_for_status()?;
69    let content_type = resp
70        .headers()
71        .get("content-type")
72        .and_then(|v| v.to_str().ok())
73        .unwrap_or("")
74        .to_string();
75
76    // If Google returns HTML, it's the virus scan warning page
77    if content_type.contains("text/html") {
78        debug!("got virus scan warning page, extracting confirm token");
79        let body = resp.text().await?;
80        let confirm_token = extract_confirm_token(&body)
81            .ok_or_else(|| anyhow::anyhow!("failed to extract confirm token from virus scan page"))?;
82
83        let confirmed_url = format!("{}&confirm={confirm_token}", handle.url);
84        let resp = client.get(&confirmed_url).send().await?.error_for_status()?;
85        stream_to_file(resp, dest, handle.size_hint.unwrap_or(0), progress).await?;
86    } else {
87        stream_to_file(resp, dest, handle.size_hint.unwrap_or(0), progress).await?;
88    }
89    Ok(())
90}
91
92/// Extract the confirm token from Google Drive's virus scan warning HTML.
93fn extract_confirm_token(html: &str) -> Option<String> {
94    // Google uses a form with action containing confirm=TOKEN or an input named confirm
95    // Pattern 1: &confirm=TOKEN&
96    if let Some(pos) = html.find("confirm=") {
97        let rest = &html[pos + 8..];
98        let end = rest.find(|c: char| c == '&' || c == '"' || c == '\'' || c.is_whitespace())?;
99        let token = &rest[..end];
100        if !token.is_empty() {
101            return Some(token.to_string());
102        }
103    }
104    // Pattern 2: name="confirm" value="TOKEN"
105    if let Some(pos) = html.find("name=\"confirm\"") {
106        let rest = &html[pos..];
107        if let Some(val_pos) = rest.find("value=\"") {
108            let val_rest = &rest[val_pos + 7..];
109            let end = val_rest.find('"')?;
110            let token = &val_rest[..end];
111            if !token.is_empty() {
112                return Some(token.to_string());
113            }
114        }
115    }
116    // Pattern 3: id="uc-download-link" with href containing confirm=
117    if let Some(pos) = html.find("id=\"uc-download-link\"") {
118        let rest = &html[pos..];
119        if let Some(href_pos) = rest.find("confirm=") {
120            let val_rest = &rest[href_pos + 8..];
121            let end = val_rest.find(|c: char| c == '&' || c == '"' || c == '\'')?;
122            let token = &val_rest[..end];
123            if !token.is_empty() {
124                return Some(token.to_string());
125            }
126        }
127    }
128    None
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134    use modde_core::GameId;
135
136    // ── extract_confirm_token: Pattern 1 – &confirm=TOKEN& ───────────────
137
138    #[test]
139    fn confirm_token_pattern1_ampersand_delimited() {
140        let html = r#"<a href="https://drive.google.com/uc?id=ID&confirm=t&export=download">Download</a>"#;
141        assert_eq!(extract_confirm_token(html), Some("t".to_string()));
142    }
143
144    #[test]
145    fn confirm_token_pattern1_long_token() {
146        let html = r#"something confirm=AbCdEfGh1234&rest"#;
147        assert_eq!(
148            extract_confirm_token(html),
149            Some("AbCdEfGh1234".to_string())
150        );
151    }
152
153    #[test]
154    fn confirm_token_pattern1_quote_delimited() {
155        let html = r#"href="https://example.com?confirm=mytoken""#;
156        assert_eq!(extract_confirm_token(html), Some("mytoken".to_string()));
157    }
158
159    #[test]
160    fn confirm_token_pattern1_single_quote_delimited() {
161        let html = r#"href='https://example.com?confirm=tok123'"#;
162        assert_eq!(extract_confirm_token(html), Some("tok123".to_string()));
163    }
164
165    #[test]
166    fn confirm_token_pattern1_whitespace_delimited() {
167        let html = "url?confirm=TOKEN rest of text";
168        assert_eq!(extract_confirm_token(html), Some("TOKEN".to_string()));
169    }
170
171    // ── extract_confirm_token: Pattern 2 – name="confirm" value="TOKEN" ──
172
173    #[test]
174    fn confirm_token_pattern2_input_field() {
175        let html =
176            r#"<input type="hidden" name="confirm" value="SecretVal"><input type="submit">"#;
177        assert_eq!(
178            extract_confirm_token(html),
179            Some("SecretVal".to_string())
180        );
181    }
182
183    #[test]
184    fn confirm_token_pattern2_with_extra_attrs() {
185        let html =
186            r#"<input class="foo" name="confirm" id="bar" value="TOKEN42">"#;
187        assert_eq!(extract_confirm_token(html), Some("TOKEN42".to_string()));
188    }
189
190    // ── extract_confirm_token: Pattern 3 – id="uc-download-link" ─────────
191
192    #[test]
193    fn confirm_token_pattern3_uc_download_link() {
194        let html = r#"<a id="uc-download-link" href="/uc?export=download&confirm=XyZ123&id=abc">Download anyway</a>"#;
195        assert_eq!(extract_confirm_token(html), Some("XyZ123".to_string()));
196    }
197
198    #[test]
199    fn confirm_token_pattern3_uc_download_link_quote_end() {
200        let html =
201            r#"<a id="uc-download-link" href="/uc?export=download&confirm=TOK">"#;
202        assert_eq!(extract_confirm_token(html), Some("TOK".to_string()));
203    }
204
205    // ── extract_confirm_token: non-matching / edge cases ─────────────────
206
207    #[test]
208    fn confirm_token_no_match_random_html() {
209        let html = "<html><body><p>Hello world</p></body></html>";
210        assert_eq!(extract_confirm_token(html), None);
211    }
212
213    #[test]
214    fn confirm_token_no_match_empty_string() {
215        assert_eq!(extract_confirm_token(""), None);
216    }
217
218    #[test]
219    fn confirm_token_no_match_similar_but_not_confirm() {
220        let html = r#"<input name="confirmed" value="nope">"#;
221        // "confirmed" does not match "confirm=" pattern for pattern 1
222        // but it does contain "confirm=" when the 'e' follows... let's check:
223        // "confirmed" contains substring "confirm" so pattern 1 ("confirm=") won't match
224        // because after "confirm" comes "ed" not "=".
225        // But wait, the html also doesn't have "confirm=" anywhere.
226        assert_eq!(extract_confirm_token(html), None);
227    }
228
229    #[test]
230    fn confirm_token_empty_token_returns_none() {
231        // Pattern 1: confirm= followed immediately by &
232        let html = "confirm=&rest";
233        // Token would be "", which is checked: `if !token.is_empty()`
234        assert_eq!(extract_confirm_token(html), None);
235    }
236
237    // ── can_handle for GoogleDriveSource ─────────────────────────────────
238
239    #[test]
240    fn can_handle_google_drive_directive() {
241        let source = GoogleDriveSource::new(Client::new());
242        let directive = DownloadDirective::GoogleDrive {
243            id: "1AbCdEfGh".to_string(),
244            hash: 42,
245        };
246        assert!(source.can_handle(&directive));
247    }
248
249    #[test]
250    fn can_handle_rejects_mega() {
251        let source = GoogleDriveSource::new(Client::new());
252        let directive = DownloadDirective::Mega {
253            url: "https://mega.nz/file/X#Y".to_string(),
254            hash: 0,
255        };
256        assert!(!source.can_handle(&directive));
257    }
258
259    #[test]
260    fn can_handle_rejects_nexus() {
261        let source = GoogleDriveSource::new(Client::new());
262        let directive = DownloadDirective::Nexus {
263            game_id: GameId::from("skyrim"),
264            mod_id: 1,
265            file_id: 1,
266            hash: 0,
267        };
268        assert!(!source.can_handle(&directive));
269    }
270
271    #[test]
272    fn can_handle_rejects_github() {
273        let source = GoogleDriveSource::new(Client::new());
274        let directive = DownloadDirective::GitHub {
275            user: "u".to_string(),
276            repo: "r".to_string(),
277            tag: "t".to_string(),
278            asset: "a".to_string(),
279            hash: 0,
280        };
281        assert!(!source.can_handle(&directive));
282    }
283
284    #[test]
285    fn can_handle_rejects_direct_url() {
286        let source = GoogleDriveSource::new(Client::new());
287        let directive = DownloadDirective::DirectURL {
288            url: "https://example.com/file".to_string(),
289            headers: std::collections::HashMap::new(),
290            hash: 0,
291        };
292        assert!(!source.can_handle(&directive));
293    }
294}
modde_sources/gdrive/mod.rs

modde_sources/gdrive/
mod.rs
