1use std::collections::HashMap;
25use std::collections::hash_map::DefaultHasher;
26use std::hash::{Hash, Hasher};
27
28use serde_json::json;
29
30const PYX_TEST_TOKEN: &str = "pyx-test-token";
31
32pub fn pyx_test_token() -> &'static str {
33 PYX_TEST_TOKEN
34}
35
36struct PackageEntry {
38 filename: &'static str,
39 url: &'static str,
40 sha256: &'static str,
41 requires_python: Option<&'static str>,
42 size: u64,
43 upload_time: &'static str,
44}
45
46fn package_database() -> HashMap<&'static str, Vec<PackageEntry>> {
48 let mut db = HashMap::new();
49
50 db.insert(
51 "iniconfig",
52 vec![
53 PackageEntry {
54 filename: "iniconfig-2.0.0-py3-none-any.whl",
55 url: "https://files.pythonhosted.org/packages/ef/a6/62565a6e1cf69e10f5727360368e451d4b7f58beeac6173dc9db836a5b46/iniconfig-2.0.0-py3-none-any.whl",
56 sha256: "b6a85871a79d2e3b22d2d1b94ac2824226a63c6b741c88f7ae975f18b6778374",
57 requires_python: Some(">=3.7"),
58 size: 5892,
59 upload_time: "2023-01-07T11:08:09.864Z",
60 },
61 PackageEntry {
62 filename: "iniconfig-2.0.0.tar.gz",
63 url: "https://files.pythonhosted.org/packages/d7/4b/cbd8e699e64a6f16ca3a8220661b5f83792b3017d0f79807cb8708d33913/iniconfig-2.0.0.tar.gz",
64 sha256: "2d91e135bf72d31a410b17c16da610a82cb55f6b0477d1a902134b24a455b8b3",
65 requires_python: Some(">=3.7"),
66 size: 4646,
67 upload_time: "2023-01-07T11:08:11.254Z",
68 },
69 ],
70 );
71
72 db.insert(
73 "anyio",
74 vec![
75 PackageEntry {
76 filename: "anyio-4.3.0-py3-none-any.whl",
77 url: "https://files.pythonhosted.org/packages/14/fd/2f20c40b45e4fb4324834aea24bd4afdf1143390242c0b33774da0e2e34f/anyio-4.3.0-py3-none-any.whl",
78 sha256: "048e05d0f6caeed70d731f3db756d35dcc1f35747c8c403364a8332c630441b8",
79 requires_python: Some(">=3.8"),
80 size: 85_584,
81 upload_time: "2024-02-19T08:36:26.842Z",
82 },
83 PackageEntry {
84 filename: "anyio-4.3.0.tar.gz",
85 url: "https://files.pythonhosted.org/packages/db/4d/3970183622f0330d3c23d9b8a5f52e365e50381fd484d08e3285104333d3/anyio-4.3.0.tar.gz",
86 sha256: "f75253795a87df48568485fd18cdd2a3fa5c4f7c5be8e5e36637733fce06fed6",
87 requires_python: Some(">=3.8"),
88 size: 159_642,
89 upload_time: "2024-02-19T08:36:28.641Z",
90 },
91 ],
92 );
93
94 db.insert(
95 "sniffio",
96 vec![
97 PackageEntry {
98 filename: "sniffio-1.3.1-py3-none-any.whl",
99 url: "https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl",
100 sha256: "2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2",
101 requires_python: Some(">=3.7"),
102 size: 10235,
103 upload_time: "2024-02-25T23:20:01.196Z",
104 },
105 PackageEntry {
106 filename: "sniffio-1.3.1.tar.gz",
107 url: "https://files.pythonhosted.org/packages/a2/87/a6771e1546d97e7e041b6ae58d80074f81b7d5121207425c964ddf5cfdbd/sniffio-1.3.1.tar.gz",
108 sha256: "f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc",
109 requires_python: Some(">=3.7"),
110 size: 20372,
111 upload_time: "2024-02-25T23:20:04.057Z",
112 },
113 ],
114 );
115
116 db.insert(
117 "idna",
118 vec![
119 PackageEntry {
120 filename: "idna-3.6-py3-none-any.whl",
121 url: "https://files.pythonhosted.org/packages/c2/e7/a82b05cf63a603df6e68d59ae6a68bf5064484a0718ea5033660af4b54a9/idna-3.6-py3-none-any.whl",
122 sha256: "c05567e9c24a6b9faaa835c4821bad0590fbb9d5779e7caa6e1cc4978e7eb24f",
123 requires_python: Some(">=3.5"),
124 size: 61567,
125 upload_time: "2023-11-25T15:40:52.604Z",
126 },
127 PackageEntry {
128 filename: "idna-3.6.tar.gz",
129 url: "https://files.pythonhosted.org/packages/bf/3f/ea4b9117521a1e9c50344b909be7886dd00a519552724809bb1f486986c2/idna-3.6.tar.gz",
130 sha256: "9ecdbbd083b06798ae1e86adcbfe8ab1479cf864e4ee30fe4e46a003d12491ca",
131 requires_python: Some(">=3.5"),
132 size: 175_426,
133 upload_time: "2023-11-25T15:40:54.902Z",
134 },
135 ],
136 );
137
138 db.insert(
139 "executable-application",
140 vec![
141 PackageEntry {
142 filename: "executable_application-0.3.0-py3-none-any.whl",
143 url: "https://files.pythonhosted.org/packages/32/97/8ab6fa1bbcb0a888f460c0a19c301f4cc4180573564ad7dd98b5ceca2ab6/executable_application-0.3.0-py3-none-any.whl",
144 sha256: "ca272aee7332e9d266663bc70037cd3ef1d74ffae40030eaf9ca46462dc8dcc6",
145 requires_python: Some(">=3.8"),
146 size: 1719,
147 upload_time: "2025-01-17T23:21:22.716Z",
148 },
149 PackageEntry {
150 filename: "executable_application-0.3.0.tar.gz",
151 url: "https://files.pythonhosted.org/packages/9a/36/e803315469274d62f2dab543e3916c0b5b65730074d295f7d48711aa9e36/executable_application-0.3.0.tar.gz",
152 sha256: "0ef8c5ddd28649503c6e4a9f55be17e5b3bd0685df7b83ff7c260b481025f261",
153 requires_python: Some(">=3.8"),
154 size: 914,
155 upload_time: "2025-01-17T23:21:24.559Z",
156 },
157 ],
158 );
159
160 db.insert(
161 "typing-extensions",
162 vec![
163 PackageEntry {
164 filename: "typing_extensions-4.10.0-py3-none-any.whl",
165 url: "https://files.pythonhosted.org/packages/f9/de/dc04a3ea60b22624b51c703a84bbe0184abcd1d0b9bc8074b5d6b7ab90bb/typing_extensions-4.10.0-py3-none-any.whl",
166 sha256: "69b1a937c3a517342112fb4c6df7e72fc39a38e7891a5730ed4985b5214b5475",
167 requires_python: Some(">=3.8"),
168 size: 33926,
169 upload_time: "2024-02-25T22:12:47.72Z",
170 },
171 PackageEntry {
172 filename: "typing_extensions-4.10.0.tar.gz",
173 url: "https://files.pythonhosted.org/packages/16/3a/0d26ce356c7465a19c9ea8814b960f8a36c3b0d07c323176620b7b483e44/typing_extensions-4.10.0.tar.gz",
174 sha256: "b0abd7c89e8fb96f98db18d86106ff1d90ab692004eb746cf6eda2682f91b3cb",
175 requires_python: Some(">=3.8"),
176 size: 77558,
177 upload_time: "2024-02-25T22:12:49.693Z",
178 },
179 ],
180 );
181
182 db
183}
184
185fn build_simple_api_response(
189 package_name: &str,
190 entries: &[PackageEntry],
191 file_url_prefix: &str,
192) -> serde_json::Value {
193 let files: Vec<serde_json::Value> = entries
194 .iter()
195 .map(|entry| {
196 let rewritten_url = entry.url.replace(
198 "https://files.pythonhosted.org/",
199 &format!("{file_url_prefix}/"),
200 );
201 let mut file_obj = json!({
202 "filename": entry.filename,
203 "url": rewritten_url,
204 "hashes": {
205 "sha256": entry.sha256
206 },
207 "size": entry.size,
208 "upload-time": entry.upload_time,
209 });
210 if let Some(rp) = entry.requires_python {
211 file_obj["requires-python"] = json!(rp);
212 }
213 file_obj
214 })
215 .collect();
216
217 json!({
218 "meta": { "api-version": "1.1" },
219 "name": package_name,
220 "files": files,
221 })
222}
223
224fn build_simple_api_response_with_project_status(
229 package_name: &str,
230 entries: &[PackageEntry],
231 file_url_prefix: &str,
232 status: &str,
233 reason: Option<&str>,
234) -> serde_json::Value {
235 let mut body = build_simple_api_response(package_name, entries, file_url_prefix);
236 let mut project_status = json!({ "status": status });
237 if let Some(reason) = reason {
238 project_status["reason"] = json!(reason);
239 }
240 body["project-status"] = project_status;
241 body
242}
243
244fn build_simple_api_response_without_upload_time(
246 package_name: &str,
247 entries: &[PackageEntry],
248 file_url_prefix: &str,
249) -> serde_json::Value {
250 let files: Vec<serde_json::Value> = entries
251 .iter()
252 .map(|entry| {
253 let rewritten_url = entry.url.replace(
254 "https://files.pythonhosted.org/",
255 &format!("{file_url_prefix}/"),
256 );
257 let mut file_obj = json!({
258 "filename": entry.filename,
259 "url": rewritten_url,
260 "hashes": {
261 "sha256": entry.sha256
262 },
263 "size": entry.size,
264 });
265 if let Some(rp) = entry.requires_python {
266 file_obj["requires-python"] = json!(rp);
267 }
268 file_obj
269 })
270 .collect();
271
272 json!({
273 "meta": { "api-version": "1.1" },
274 "name": package_name,
275 "files": files,
276 })
277}
278
279fn build_simple_api_response_relative(
283 package_name: &str,
284 entries: &[PackageEntry],
285) -> serde_json::Value {
286 let files: Vec<serde_json::Value> = entries
287 .iter()
288 .map(|entry| {
289 let rewritten_url = entry
290 .url
291 .replace("https://files.pythonhosted.org/", "../../../files/");
292 let mut file_obj = json!({
293 "filename": entry.filename,
294 "url": rewritten_url,
295 "hashes": {
296 "sha256": entry.sha256
297 },
298 "size": entry.size,
299 "upload-time": entry.upload_time,
300 });
301 if let Some(rp) = entry.requires_python {
302 file_obj["requires-python"] = json!(rp);
303 }
304 file_obj
305 })
306 .collect();
307
308 json!({
309 "meta": { "api-version": "1.1" },
310 "name": package_name,
311 "files": files,
312 })
313}
314
315pub struct PypiProxy {
317 server: wiremock::MockServer,
318}
319
320impl PypiProxy {
321 pub fn uri(&self) -> String {
323 self.server.uri()
324 }
325
326 pub fn host(&self) -> String {
328 let url = url::Url::parse(&self.server.uri()).expect("valid URL");
329 url.host_str().expect("has host").to_string()
330 }
331
332 pub fn host_port(&self) -> String {
334 self.server
335 .uri()
336 .strip_prefix("http://")
337 .expect("wiremock server URI should start with http://")
338 .to_string()
339 }
340
341 pub fn authenticated_url(&self, username: &str, password: &str, path: &str) -> String {
343 format!("http://{username}:{password}@{}{path}", self.host_port())
344 }
345
346 pub fn authenticated_uri(&self, username: &str, password: &str) -> String {
348 format!("http://{username}:{password}@{}", self.host_port())
349 }
350
351 pub fn username_url(&self, username: &str, path: &str) -> String {
353 format!("http://{username}@{}{path}", self.host_port())
354 }
355
356 pub fn url(&self, path: &str) -> String {
358 format!("{}{path}", self.uri())
359 }
360}
361
362pub async fn start() -> PypiProxy {
377 use wiremock::{Mock, MockServer, Request, ResponseTemplate};
378
379 let server = MockServer::start().await;
380 let db = package_database();
381 let server_uri = server.uri();
382
383 Mock::given(wiremock::matchers::any())
386 .respond_with(move |req: &Request| {
387 let path = req.url.path();
388
389 let auth = req
391 .headers
392 .get(&http::header::AUTHORIZATION)
393 .and_then(parse_basic_auth);
394 let bearer_auth = req
395 .headers
396 .get(&http::header::AUTHORIZATION)
397 .and_then(parse_bearer_auth);
398
399 if let Some(rest) = path.strip_prefix("/basic-auth/files/") {
401 if auth
402 .as_ref()
403 .is_some_and(|(u, p)| u == "public" && p == "heron")
404 {
405 let target = format!("https://files.pythonhosted.org/{rest}");
406 return ResponseTemplate::new(302).insert_header("Location", target);
407 }
408 return unauthorized_response();
409 }
410
411 if let Some(rest) = path.strip_prefix("/bearer-auth/files/") {
413 if bearer_auth
414 .as_ref()
415 .is_some_and(|token| token == PYX_TEST_TOKEN)
416 {
417 let target = format!("https://files.pythonhosted.org/{rest}");
418 return ResponseTemplate::new(302).insert_header("Location", target);
419 }
420 return unauthorized_response();
421 }
422
423 if let Some(rest) = path.strip_prefix("/basic-auth-heron/files/") {
425 if auth
426 .as_ref()
427 .is_some_and(|(u, p)| u == "public" && p == "heron")
428 {
429 let target = format!("https://files.pythonhosted.org/{rest}");
430 return ResponseTemplate::new(302).insert_header("Location", target);
431 }
432 return unauthorized_response();
433 }
434
435 if let Some(rest) = path.strip_prefix("/basic-auth-eagle/files/") {
437 if auth
438 .as_ref()
439 .is_some_and(|(u, p)| u == "public" && p == "eagle")
440 {
441 let target = format!("https://files.pythonhosted.org/{rest}");
442 return ResponseTemplate::new(302).insert_header("Location", target);
443 }
444 return unauthorized_response();
445 }
446
447 if let Some(rest) = path.strip_prefix("/files/") {
449 let target = format!("https://files.pythonhosted.org/{rest}");
450 return ResponseTemplate::new(302).insert_header("Location", target);
451 }
452
453 if let Some(pkg) = extract_package_name(path, "/basic-auth/relative/simple/") {
455 if auth
456 .as_ref()
457 .is_some_and(|(u, p)| u == "public" && p == "heron")
458 {
459 if let Some(entries) = db.get(pkg) {
460 let body = build_simple_api_response_relative(pkg, entries);
461 return simple_api_response(&body);
462 }
463 return ResponseTemplate::new(404);
464 }
465 return unauthorized_response();
466 }
467
468 if let Some(pkg) = extract_package_name(path, "/basic-auth/simple/") {
470 if auth
471 .as_ref()
472 .is_some_and(|(u, p)| u == "public" && p == "heron")
473 {
474 if let Some(entries) = db.get(pkg) {
475 let file_prefix = format!("{server_uri}/basic-auth/files");
476 let body = build_simple_api_response(pkg, entries, &file_prefix);
477 return simple_api_response(&body);
478 }
479 return ResponseTemplate::new(404);
480 }
481 return unauthorized_response();
482 }
483
484 if let Some(pkg) = extract_package_name(path, "/bearer-auth/simple/") {
486 if bearer_auth
487 .as_ref()
488 .is_some_and(|token| token == PYX_TEST_TOKEN)
489 {
490 if let Some(entries) = db.get(pkg) {
491 let file_prefix = format!("{server_uri}/bearer-auth/files");
492 let body = build_simple_api_response(pkg, entries, &file_prefix);
493 return simple_api_response(&body);
494 }
495 return ResponseTemplate::new(404);
496 }
497 return unauthorized_response();
498 }
499
500 if let Some(pkg) = extract_package_name(path, "/basic-auth-heron/simple/") {
502 if auth
503 .as_ref()
504 .is_some_and(|(u, p)| u == "public" && p == "heron")
505 {
506 if let Some(entries) = db.get(pkg) {
507 let file_prefix = format!("{server_uri}/basic-auth-heron/files");
508 let body = build_simple_api_response(pkg, entries, &file_prefix);
509 return simple_api_response(&body);
510 }
511 return ResponseTemplate::new(404);
512 }
513 return unauthorized_response();
514 }
515
516 if let Some(pkg) = extract_package_name(path, "/basic-auth-eagle/simple/") {
518 if auth
519 .as_ref()
520 .is_some_and(|(u, p)| u == "public" && p == "eagle")
521 {
522 if let Some(entries) = db.get(pkg) {
523 let file_prefix = format!("{server_uri}/basic-auth-eagle/files");
524 let body = build_simple_api_response(pkg, entries, &file_prefix);
525 return simple_api_response(&body);
526 }
527 return ResponseTemplate::new(404);
528 }
529 return unauthorized_response();
530 }
531
532 if let Some(pkg) = extract_package_name(path, "/relative/simple/") {
534 if let Some(entries) = db.get(pkg) {
535 let body = build_simple_api_response_relative(pkg, entries);
536 return simple_api_response(&body);
537 }
538 return ResponseTemplate::new(404);
539 }
540
541 if let Some(pkg) = extract_package_name(path, "/no-upload-time/simple/") {
543 if let Some(entries) = db.get(pkg) {
544 let file_prefix = "https://files.pythonhosted.org";
545 let body =
546 build_simple_api_response_without_upload_time(pkg, entries, file_prefix);
547 return simple_api_response(&body);
548 }
549 return ResponseTemplate::new(404);
550 }
551
552 if let Some(pkg) = extract_package_name(path, "/simple/") {
556 if let Some(entries) = db.get(pkg) {
557 let file_prefix = "https://files.pythonhosted.org";
558 let body = build_simple_api_response(pkg, entries, file_prefix);
559 return simple_api_response(&body);
560 }
561 return ResponseTemplate::new(404);
562 }
563
564 if let Some((status, reason, kind, suffix)) = parse_status_path(path) {
569 match kind {
570 StatusRouteKind::Simple => {
571 if let Some(pkg) = suffix.strip_suffix('/')
574 && !pkg.contains('/')
575 && let Some(entries) = db.get(pkg)
576 {
577 let file_prefix = format!(
578 "{server_uri}{prefix}/files",
579 prefix = status_route_prefix(status, reason),
580 );
581 let body = build_simple_api_response_with_project_status(
582 pkg,
583 entries,
584 &file_prefix,
585 status,
586 reason,
587 );
588 return simple_api_response(&body);
589 }
590 return ResponseTemplate::new(404);
591 }
592 StatusRouteKind::Files => {
593 let target = format!("https://files.pythonhosted.org/{suffix}");
594 return ResponseTemplate::new(302).insert_header("Location", target);
595 }
596 }
597 }
598
599 ResponseTemplate::new(404)
600 })
601 .mount(&server)
602 .await;
603
604 PypiProxy { server }
605}
606
607enum StatusRouteKind {
608 Simple,
609 Files,
610}
611
612fn parse_status_path(path: &str) -> Option<(&str, Option<&str>, StatusRouteKind, &str)> {
616 let after_status_prefix = path.strip_prefix("/status/")?;
617 let (status, rest) = after_status_prefix.split_once('/')?;
618 let (reason, rest) = if let Some(after_reason_prefix) = rest.strip_prefix("reason/") {
620 let (reason, rest) = after_reason_prefix.split_once('/')?;
621 (Some(reason), rest)
622 } else {
623 (None, rest)
624 };
625 if let Some(suffix) = rest.strip_prefix("simple/") {
628 Some((status, reason, StatusRouteKind::Simple, suffix))
629 } else if let Some(suffix) = rest.strip_prefix("files/") {
630 Some((status, reason, StatusRouteKind::Files, suffix))
631 } else {
632 None
633 }
634}
635
636fn status_route_prefix(status: &str, reason: Option<&str>) -> String {
640 match reason {
641 Some(reason) => format!("/status/{status}/reason/{reason}"),
642 None => format!("/status/{status}"),
643 }
644}
645
646fn extract_package_name<'a>(path: &'a str, prefix: &str) -> Option<&'a str> {
648 let rest = path.strip_prefix(prefix)?;
649 let pkg = rest.strip_suffix('/')?;
650 if pkg.contains('/') {
652 return None;
653 }
654 Some(pkg)
655}
656
657fn parse_basic_auth(value: &wiremock::http::HeaderValue) -> Option<(String, String)> {
659 use base64::Engine;
660
661 let s = value.as_bytes();
662 let s = std::str::from_utf8(s).ok()?;
663 let encoded = s.strip_prefix("Basic ")?;
664 let decoded = base64::engine::general_purpose::STANDARD
665 .decode(encoded)
666 .ok()?;
667 let decoded = String::from_utf8(decoded).ok()?;
668 let (user, pass) = decoded.split_once(':')?;
669 Some((user.to_string(), pass.to_string()))
670}
671
672fn parse_bearer_auth(value: &wiremock::http::HeaderValue) -> Option<String> {
674 let s = value.as_bytes();
675 let s = std::str::from_utf8(s).ok()?;
676 let token = s.strip_prefix("Bearer ")?;
677 Some(token.to_string())
678}
679
680fn unauthorized_response() -> wiremock::ResponseTemplate {
681 wiremock::ResponseTemplate::new(401)
682 .insert_header("WWW-Authenticate", r#"Basic realm="authenticated""#)
683}
684
685fn simple_api_response(body: &serde_json::Value) -> wiremock::ResponseTemplate {
686 let body_str = body.to_string();
688 let mut hasher = DefaultHasher::new();
689 body_str.hash(&mut hasher);
690 let etag = format!("\"{}\"", hasher.finish());
691
692 wiremock::ResponseTemplate::new(200)
695 .insert_header("Cache-Control", "max-age=600, public")
696 .insert_header("ETag", etag)
697 .set_body_raw(body_str, "application/vnd.pypi.simple.v1+json")
698}