modde_sources/mega/
mod.rs

1//! Mega.nz download source, including client-side AES-128-CTR decryption of
2//! the encrypted payload.
3
4use std::collections::HashMap;
5use std::path::Path;
6
7use aes::Aes128;
8use anyhow::{Context, Result, bail};
9use base64::Engine;
10use base64::engine::general_purpose::URL_SAFE_NO_PAD;
11use ctr::Ctr128BE;
12use ctr::cipher::{KeyIvInit, StreamCipher};
13use futures::StreamExt;
14use reqwest::Client;
15use serde::Deserialize;
16use tokio::io::AsyncWriteExt;
17use tracing::debug;
18
19use modde_core::manifest::wabbajack::DownloadDirective;
20
21use crate::common::{ensure_parent, verify_and_wrap};
22use crate::error::{SourceError, SourceResult, status_error};
23use crate::traits::{DownloadHandle, DownloadSource, ProgressCallback, VerifiedFile};
24
25const MEGA_API_URL: &str = "https://g.api.mega.co.nz/cs";
26
27/// Mega.nz download source.
28///
29/// Handles Mega's client-side AES-128-CTR decryption protocol.
30pub struct MegaSource {
31    client: Client,
32}
33
34#[derive(Debug, Deserialize)]
35struct MegaFileResponse {
36    /// Download URL
37    g: String,
38    /// File size
39    s: u64,
40}
41
42/// Parse a Mega URL to extract the file handle and key.
43/// Supports both new format `/file/HANDLE#KEY` and old format `/#!HANDLE!KEY`.
44fn parse_mega_url(url: &str) -> Result<(String, String)> {
45    // New format: https://mega.nz/file/HANDLE#KEY
46    if let Some(rest) = url
47        .strip_prefix("https://mega.nz/file/")
48        .or_else(|| url.strip_prefix("http://mega.nz/file/"))
49    {
50        let parts: Vec<&str> = rest.splitn(2, '#').collect();
51        if parts.len() == 2 {
52            return Ok((parts[0].to_string(), parts[1].to_string()));
53        }
54    }
55
56    // Old format: https://mega.nz/#!HANDLE!KEY
57    if let Some(rest) = url.find("#!") {
58        let fragment = &url[rest + 2..];
59        let parts: Vec<&str> = fragment.splitn(2, '!').collect();
60        if parts.len() == 2 {
61            return Ok((parts[0].to_string(), parts[1].to_string()));
62        }
63    }
64
65    bail!("invalid Mega URL format: {url}")
66}
67
68/// Decode a Mega key from base64url, XOR first/second halves for AES-128 key, extract IV.
69fn decode_mega_key(key_b64: &str) -> Result<([u8; 16], [u8; 16])> {
70    let key_bytes = URL_SAFE_NO_PAD
71        .decode(key_b64)
72        .context("failed to decode Mega key from base64url")?;
73
74    if key_bytes.len() != 32 {
75        bail!("expected 32-byte Mega key, got {} bytes", key_bytes.len());
76    }
77
78    // XOR first 16 bytes with second 16 bytes to get AES key
79    let mut aes_key = [0u8; 16];
80    for i in 0..16 {
81        aes_key[i] = key_bytes[i] ^ key_bytes[i + 16];
82    }
83
84    // IV is bytes 16..24, zero-padded to 16 bytes (counter starts at 0)
85    let mut iv = [0u8; 16];
86    iv[..8].copy_from_slice(&key_bytes[16..24]);
87    // bytes 8..16 of IV are zero (counter)
88
89    Ok((aes_key, iv))
90}
91
92impl MegaSource {
93    /// Create a source that downloads over the given HTTP `client`.
94    #[must_use]
95    pub fn new(client: Client) -> Self {
96        Self { client }
97    }
98}
99
100impl DownloadSource for MegaSource {
101    fn can_handle(&self, directive: &DownloadDirective) -> bool {
102        matches!(directive, DownloadDirective::Mega { .. })
103    }
104
105    async fn resolve(&self, directive: &DownloadDirective) -> SourceResult<DownloadHandle> {
106        let DownloadDirective::Mega { url, hash } = directive else {
107            return Err(SourceError::other(anyhow::anyhow!("not a Mega directive")));
108        };
109
110        let (handle_id, key_b64) = parse_mega_url(url).map_err(SourceError::other)?;
111
112        // Call Mega API to get download URL
113        let api_url = format!("{MEGA_API_URL}?id=0");
114        let payload = serde_json::json!([{"a": "g", "g": 1, "p": handle_id}]);
115
116        let resp = status_error(self.client.post(&api_url).json(&payload).send().await?)?;
117
118        let body: Vec<MegaFileResponse> = resp.json().await?;
119        let file_info = body
120            .into_iter()
121            .next()
122            .ok_or_else(|| SourceError::other(anyhow::anyhow!("empty response from Mega API")))?;
123
124        debug!(download_url = %file_info.g, size = file_info.s, "resolved Mega download URL");
125
126        let mut headers = HashMap::new();
127        headers.insert("x-mega-key".to_string(), key_b64);
128
129        Ok(DownloadHandle {
130            url: file_info.g,
131            candidate_urls: Vec::new(),
132            headers,
133            expected_hash: *hash,
134            size_hint: Some(file_info.s),
135        })
136    }
137
138    async fn download_with_progress(
139        &self,
140        handle: DownloadHandle,
141        dest: &Path,
142        progress: ProgressCallback,
143    ) -> SourceResult<VerifiedFile> {
144        ensure_parent(dest).await?;
145
146        let key_b64 = handle
147            .headers
148            .get("x-mega-key")
149            .ok_or_else(|| {
150                SourceError::other(anyhow::anyhow!(
151                    "missing x-mega-key header in download handle"
152                ))
153            })?
154            .clone();
155
156        let (aes_key, iv) = decode_mega_key(&key_b64).map_err(SourceError::other)?;
157
158        let resp = status_error(self.client.get(&handle.url).send().await?)?;
159
160        let total = resp.content_length().or(handle.size_hint).unwrap_or(0);
161        let mut file = tokio::fs::File::create(dest).await?;
162        let mut downloaded: u64 = 0;
163
164        let mut cipher = Ctr128BE::<Aes128>::new(&aes_key.into(), &iv.into());
165
166        let mut stream = resp.bytes_stream();
167        while let Some(chunk) = stream.next().await {
168            let mut chunk = chunk?.to_vec();
169            cipher.apply_keystream(&mut chunk);
170            file.write_all(&chunk).await?;
171            downloaded += chunk.len() as u64;
172            progress(downloaded, total);
173        }
174
175        file.flush().await?;
176        debug!(bytes = downloaded, "Mega download complete");
177
178        verify_and_wrap(dest, handle.expected_hash).await
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185    use modde_core::GameId;
186
187    // ── parse_mega_url: new format /file/HANDLE#KEY ──────────────────────
188
189    #[test]
190    fn parse_new_format_https() {
191        let (handle, key) = parse_mega_url("https://mega.nz/file/ABC123#some_key_base64").unwrap();
192        assert_eq!(handle, "ABC123");
193        assert_eq!(key, "some_key_base64");
194    }
195
196    #[test]
197    fn parse_new_format_http() {
198        let (handle, key) = parse_mega_url("http://mega.nz/file/XYZ789#another_key").unwrap();
199        assert_eq!(handle, "XYZ789");
200        assert_eq!(key, "another_key");
201    }
202
203    #[test]
204    fn parse_new_format_long_handle_and_key() {
205        let (handle, key) = parse_mega_url(
206            "https://mega.nz/file/AbCdEfGhIjKlMnOp#AAAAAAAAAAAABBBBBBBBBBBBCCCCCCCCCCCCDDDDDDDDDDDD",
207        )
208        .unwrap();
209        assert_eq!(handle, "AbCdEfGhIjKlMnOp");
210        assert_eq!(key, "AAAAAAAAAAAABBBBBBBBBBBBCCCCCCCCCCCCDDDDDDDDDDDD");
211    }
212
213    #[test]
214    fn parse_new_format_key_with_special_base64url_chars() {
215        // base64url uses - and _ instead of + and /
216        let (handle, key) =
217            parse_mega_url("https://mega.nz/file/HANDLE#a-b_c-d_e-f_g-h_i-j_k").unwrap();
218        assert_eq!(handle, "HANDLE");
219        assert_eq!(key, "a-b_c-d_e-f_g-h_i-j_k");
220    }
221
222    // ── parse_mega_url: old format /#!HANDLE!KEY ─────────────────────────
223
224    #[test]
225    fn parse_old_format_https() {
226        let (handle, key) = parse_mega_url("https://mega.nz/#!ABC123!some_key_base64").unwrap();
227        assert_eq!(handle, "ABC123");
228        assert_eq!(key, "some_key_base64");
229    }
230
231    #[test]
232    fn parse_old_format_http() {
233        let (handle, key) = parse_mega_url("http://mega.nz/#!OldHandle!OldKey123").unwrap();
234        assert_eq!(handle, "OldHandle");
235        assert_eq!(key, "OldKey123");
236    }
237
238    #[test]
239    fn parse_old_format_with_extra_prefix() {
240        // The old-format parser uses find("#!"), so it works even with odd prefixes
241        let (handle, key) = parse_mega_url("https://mega.co.nz/#!HANDLE!KEY").unwrap();
242        assert_eq!(handle, "HANDLE");
243        assert_eq!(key, "KEY");
244    }
245
246    // ── parse_mega_url: invalid URLs ─────────────────────────────────────
247
248    #[test]
249    fn parse_url_no_hash_new_format() {
250        // Missing # separator in new format
251        assert!(parse_mega_url("https://mega.nz/file/ABCnohash").is_err());
252    }
253
254    #[test]
255    fn parse_url_random_url() {
256        assert!(parse_mega_url("https://example.com/file").is_err());
257    }
258
259    #[test]
260    fn parse_url_empty_string() {
261        assert!(parse_mega_url("").is_err());
262    }
263
264    #[test]
265    fn parse_url_only_domain() {
266        assert!(parse_mega_url("https://mega.nz").is_err());
267    }
268
269    #[test]
270    fn parse_url_no_key_after_hash() {
271        // "splitn(2, '#')" produces ["HANDLE", ""] – length 2 but empty key
272        // The function does not reject empty keys, so this succeeds with empty key
273        let result = parse_mega_url("https://mega.nz/file/HANDLE#");
274        // Regardless of success/failure, document behaviour
275        if let Ok((_handle, key)) = &result {
276            assert!(key.is_empty());
277        }
278    }
279
280    #[test]
281    fn parse_url_with_query_params_new_format() {
282        // Query params end up as part of the key (since we only split on #)
283        let (handle, key) = parse_mega_url("https://mega.nz/file/HANDLE#KEY?foo=bar").unwrap();
284        assert_eq!(handle, "HANDLE");
285        assert_eq!(key, "KEY?foo=bar");
286    }
287
288    #[test]
289    fn parse_url_with_extra_path_segments() {
290        // "/file/" prefix is stripped, then everything up to # is handle
291        let (handle, key) = parse_mega_url("https://mega.nz/file/HANDLE/extra#KEY").unwrap();
292        assert_eq!(handle, "HANDLE/extra");
293        assert_eq!(key, "KEY");
294    }
295
296    // ── decode_mega_key: valid 32-byte keys ──────────────────────────────
297
298    #[test]
299    fn decode_key_valid_32_bytes() {
300        // 32 zero bytes -> base64url = AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=
301        // URL_SAFE_NO_PAD expects no padding, so use the unpadded form
302        let key_bytes = [0u8; 32];
303        let key_b64 = URL_SAFE_NO_PAD.encode(key_bytes);
304        let (aes_key, iv) = decode_mega_key(&key_b64).unwrap();
305        // 0 XOR 0 = 0 for all bytes
306        assert_eq!(aes_key, [0u8; 16]);
307        // IV = bytes 16..24 of original (all zero), zero-padded
308        assert_eq!(iv, [0u8; 16]);
309    }
310
311    #[test]
312    fn decode_key_xor_logic() {
313        // Construct a 32-byte key where first half = [1..=16], second half = [17..=32]
314        let mut key_bytes = [0u8; 32];
315        for i in 0..16 {
316            key_bytes[i] = (i + 1) as u8;
317            key_bytes[i + 16] = (i + 17) as u8;
318        }
319        let key_b64 = URL_SAFE_NO_PAD.encode(key_bytes);
320        let (aes_key, iv) = decode_mega_key(&key_b64).unwrap();
321
322        // Verify XOR: aes_key[i] = key_bytes[i] ^ key_bytes[i+16]
323        for i in 0..16 {
324            assert_eq!(
325                aes_key[i],
326                key_bytes[i] ^ key_bytes[i + 16],
327                "XOR mismatch at index {i}"
328            );
329        }
330
331        // Verify IV: bytes 16..24 of original, zero-padded to 16
332        let mut expected_iv = [0u8; 16];
333        expected_iv[..8].copy_from_slice(&key_bytes[16..24]);
334        assert_eq!(iv, expected_iv);
335    }
336
337    #[test]
338    fn decode_key_xor_inverse() {
339        // If first half == second half, XOR yields all zeros
340        let mut key_bytes = [0u8; 32];
341        for byte in key_bytes.iter_mut().take(16) {
342            *byte = 0xAB;
343        }
344        for byte in key_bytes.iter_mut().skip(16) {
345            *byte = 0xAB;
346        }
347        let key_b64 = URL_SAFE_NO_PAD.encode(key_bytes);
348        let (aes_key, _iv) = decode_mega_key(&key_b64).unwrap();
349        assert_eq!(aes_key, [0u8; 16]);
350    }
351
352    #[test]
353    fn decode_key_xor_all_ones() {
354        // first half = 0xFF, second half = 0x00 -> XOR = 0xFF
355        let mut key_bytes = [0u8; 32];
356        for byte in key_bytes.iter_mut().take(16) {
357            *byte = 0xFF;
358        }
359        let key_b64 = URL_SAFE_NO_PAD.encode(key_bytes);
360        let (aes_key, _iv) = decode_mega_key(&key_b64).unwrap();
361        assert_eq!(aes_key, [0xFF; 16]);
362    }
363
364    // ── decode_mega_key: IV extraction correctness ───────────────────────
365
366    #[test]
367    fn decode_key_iv_extraction() {
368        let mut key_bytes = [0u8; 32];
369        // Set bytes 16..24 to distinct values
370        for i in 0..8 {
371            key_bytes[16 + i] = (0x10 + i) as u8;
372        }
373        // Set bytes 24..32 to something else (should NOT appear in IV)
374        for i in 0..8 {
375            key_bytes[24 + i] = 0xFF;
376        }
377        let key_b64 = URL_SAFE_NO_PAD.encode(key_bytes);
378        let (_aes_key, iv) = decode_mega_key(&key_b64).unwrap();
379
380        // First 8 bytes of IV = bytes 16..24 of original
381        for (i, byte) in iv.iter().enumerate().take(8) {
382            assert_eq!(*byte, (0x10 + i) as u8, "IV byte {i} mismatch");
383        }
384        // Last 8 bytes of IV must be zero (counter)
385        for (i, byte) in iv.iter().enumerate().skip(8) {
386            assert_eq!(*byte, 0, "IV counter byte {i} should be zero");
387        }
388    }
389
390    // ── decode_mega_key: invalid inputs ──────────────────────────────────
391
392    #[test]
393    fn decode_key_too_short() {
394        let short = URL_SAFE_NO_PAD.encode([0u8; 16]);
395        let err = decode_mega_key(&short).unwrap_err();
396        assert!(
397            err.to_string().contains("expected 32-byte"),
398            "unexpected error: {err}"
399        );
400    }
401
402    #[test]
403    fn decode_key_too_long() {
404        let long = URL_SAFE_NO_PAD.encode([0u8; 48]);
405        let err = decode_mega_key(&long).unwrap_err();
406        assert!(
407            err.to_string().contains("expected 32-byte"),
408            "unexpected error: {err}"
409        );
410    }
411
412    #[test]
413    fn decode_key_empty() {
414        let err = decode_mega_key("").unwrap_err();
415        assert!(
416            err.to_string().contains("expected 32-byte"),
417            "unexpected error: {err}"
418        );
419    }
420
421    #[test]
422    fn decode_key_invalid_base64() {
423        let err = decode_mega_key("!!!not-valid-base64!!!").unwrap_err();
424        assert!(
425            err.to_string().contains("base64"),
426            "unexpected error: {err}"
427        );
428    }
429
430    #[test]
431    fn decode_key_one_byte() {
432        let one = URL_SAFE_NO_PAD.encode([0x42u8; 1]);
433        let err = decode_mega_key(&one).unwrap_err();
434        assert!(err.to_string().contains("expected 32-byte"));
435    }
436
437    #[test]
438    fn decode_key_31_bytes() {
439        let data = URL_SAFE_NO_PAD.encode([0u8; 31]);
440        assert!(decode_mega_key(&data).is_err());
441    }
442
443    #[test]
444    fn decode_key_33_bytes() {
445        let data = URL_SAFE_NO_PAD.encode([0u8; 33]);
446        assert!(decode_mega_key(&data).is_err());
447    }
448
449    // ── decode_mega_key: very long key string ────────────────────────────
450
451    #[test]
452    fn decode_key_very_long_base64() {
453        // 256 bytes is way too long
454        let long = URL_SAFE_NO_PAD.encode([0xABu8; 256]);
455        assert!(decode_mega_key(&long).is_err());
456    }
457
458    // ── can_handle ───────────────────────────────────────────────────────
459
460    #[test]
461    fn can_handle_mega_directive() {
462        let source = MegaSource::new(Client::new());
463        let directive = DownloadDirective::Mega {
464            url: "https://mega.nz/file/ABC#KEY".to_string(),
465            hash: 0,
466        };
467        assert!(source.can_handle(&directive));
468    }
469
470    #[test]
471    fn can_handle_rejects_nexus() {
472        let source = MegaSource::new(Client::new());
473        let directive = DownloadDirective::Nexus {
474            game_id: GameId::from("skyrim"),
475            mod_id: 1.into(),
476            file_id: 1.into(),
477            hash: 0,
478        };
479        assert!(!source.can_handle(&directive));
480    }
481
482    #[test]
483    fn can_handle_rejects_google_drive() {
484        let source = MegaSource::new(Client::new());
485        let directive = DownloadDirective::GoogleDrive {
486            id: "some-id".to_string(),
487            hash: 0,
488        };
489        assert!(!source.can_handle(&directive));
490    }
491
492    #[test]
493    fn can_handle_rejects_github() {
494        let source = MegaSource::new(Client::new());
495        let directive = DownloadDirective::GitHub {
496            user: "user".to_string(),
497            repo: "repo".to_string(),
498            tag: "v1".to_string(),
499            asset: "file.zip".to_string(),
500            hash: 0,
501        };
502        assert!(!source.can_handle(&directive));
503    }
504
505    #[test]
506    fn can_handle_rejects_direct_url() {
507        let source = MegaSource::new(Client::new());
508        let directive = DownloadDirective::DirectURL {
509            url: "https://example.com/file.zip".to_string(),
510            headers: HashMap::new(),
511            mirror_resolver: None,
512            hash: 0,
513        };
514        assert!(!source.can_handle(&directive));
515    }
516}
modde_sources/mega/mod.rs

modde_sources/mega/
mod.rs
