use crate::activity_log::{ActionType, ActivityEntry};
use crate::audit::AuditEntry;
use crate::config::basic_auth_header;
use crate::AppState;
use axum::{
extract::{Path, State},
http::{header, StatusCode},
response::{Html, IntoResponse, Response},
routing::get,
Router,
};
use std::sync::Arc;
use std::time::Duration;
pub fn routes() -> Router<Arc<AppState>> {
Router::new()
.route("/simple/", get(list_packages))
.route("/simple/{name}/", get(package_versions))
.route("/simple/{name}/{filename}", get(download_file))
}
async fn list_packages(State(state): State<Arc<AppState>>) -> impl IntoResponse {
let keys = state.storage.list("pypi/").await;
let mut packages = std::collections::HashSet::new();
for key in keys {
if let Some(pkg) = key.strip_prefix("pypi/").and_then(|k| k.split('/').next()) {
if !pkg.is_empty() {
packages.insert(pkg.to_string());
}
}
}
let mut html = String::from(
"<!DOCTYPE html>\n<html><head><title>Simple Index</title></head><body><h1>Simple Index</h1>\n",
);
let mut pkg_list: Vec<_> = packages.into_iter().collect();
pkg_list.sort();
for pkg in pkg_list {
html.push_str(&format!("<a href=\"/simple/{}/\">{}</a><br>\n", pkg, pkg));
}
html.push_str("</body></html>");
(StatusCode::OK, Html(html))
}
async fn package_versions(
State(state): State<Arc<AppState>>,
Path(name): Path<String>,
) -> Response {
let normalized = normalize_name(&name);
let prefix = format!("pypi/{}/", normalized);
let keys = state.storage.list(&prefix).await;
if !keys.is_empty() {
let mut html = format!(
"<!DOCTYPE html>\n<html><head><title>Links for {}</title></head><body><h1>Links for {}</h1>\n",
name, name
);
for key in &keys {
if let Some(filename) = key.strip_prefix(&prefix) {
if !filename.is_empty() {
html.push_str(&format!(
"<a href=\"/simple/{}/{}\">{}</a><br>\n",
normalized, filename, filename
));
}
}
}
html.push_str("</body></html>");
return (StatusCode::OK, Html(html)).into_response();
}
if let Some(proxy_url) = &state.config.pypi.proxy {
let url = format!("{}/{}/", proxy_url.trim_end_matches('/'), normalized);
if let Ok(html) = fetch_package_page(
&state.http_client,
&url,
state.config.pypi.proxy_timeout,
state.config.pypi.proxy_auth.as_deref(),
)
.await
{
let rewritten = rewrite_pypi_links(&html, &normalized);
return (StatusCode::OK, Html(rewritten)).into_response();
}
}
StatusCode::NOT_FOUND.into_response()
}
async fn download_file(
State(state): State<Arc<AppState>>,
Path((name, filename)): Path<(String, String)>,
) -> Response {
let normalized = normalize_name(&name);
let key = format!("pypi/{}/{}", normalized, filename);
if let Ok(data) = state.storage.get(&key).await {
state.metrics.record_download("pypi");
state.metrics.record_cache_hit();
state.activity.push(ActivityEntry::new(
ActionType::CacheHit,
format!("{}/{}", name, filename),
"pypi",
"CACHE",
));
state
.audit
.log(AuditEntry::new("cache_hit", "api", "", "pypi", ""));
let content_type = if filename.ends_with(".whl") {
"application/zip"
} else if filename.ends_with(".tar.gz") || filename.ends_with(".tgz") {
"application/gzip"
} else {
"application/octet-stream"
};
return (StatusCode::OK, [(header::CONTENT_TYPE, content_type)], data).into_response();
}
if let Some(proxy_url) = &state.config.pypi.proxy {
let page_url = format!("{}/{}/", proxy_url.trim_end_matches('/'), normalized);
if let Ok(html) = fetch_package_page(
&state.http_client,
&page_url,
state.config.pypi.proxy_timeout,
state.config.pypi.proxy_auth.as_deref(),
)
.await
{
if let Some(file_url) = find_file_url(&html, &filename) {
if let Ok(data) = fetch_file(
&state.http_client,
&file_url,
state.config.pypi.proxy_timeout,
state.config.pypi.proxy_auth.as_deref(),
)
.await
{
state.metrics.record_download("pypi");
state.metrics.record_cache_miss();
state.activity.push(ActivityEntry::new(
ActionType::ProxyFetch,
format!("{}/{}", name, filename),
"pypi",
"PROXY",
));
state
.audit
.log(AuditEntry::new("proxy_fetch", "api", "", "pypi", ""));
let storage = state.storage.clone();
let key_clone = key.clone();
let data_clone = data.clone();
tokio::spawn(async move {
let _ = storage.put(&key_clone, &data_clone).await;
});
state.repo_index.invalidate("pypi");
let content_type = if filename.ends_with(".whl") {
"application/zip"
} else if filename.ends_with(".tar.gz") || filename.ends_with(".tgz") {
"application/gzip"
} else {
"application/octet-stream"
};
return (StatusCode::OK, [(header::CONTENT_TYPE, content_type)], data)
.into_response();
}
}
}
}
StatusCode::NOT_FOUND.into_response()
}
fn normalize_name(name: &str) -> String {
name.to_lowercase().replace(['-', '_', '.'], "-")
}
async fn fetch_package_page(
client: &reqwest::Client,
url: &str,
timeout_secs: u64,
auth: Option<&str>,
) -> Result<String, ()> {
let mut request = client
.get(url)
.timeout(Duration::from_secs(timeout_secs))
.header("Accept", "text/html");
if let Some(credentials) = auth {
request = request.header("Authorization", basic_auth_header(credentials));
}
let response = request.send().await.map_err(|_| ())?;
if !response.status().is_success() {
return Err(());
}
response.text().await.map_err(|_| ())
}
async fn fetch_file(
client: &reqwest::Client,
url: &str,
timeout_secs: u64,
auth: Option<&str>,
) -> Result<Vec<u8>, ()> {
let mut request = client.get(url).timeout(Duration::from_secs(timeout_secs));
if let Some(credentials) = auth {
request = request.header("Authorization", basic_auth_header(credentials));
}
let response = request.send().await.map_err(|_| ())?;
if !response.status().is_success() {
return Err(());
}
response.bytes().await.map(|b| b.to_vec()).map_err(|_| ())
}
fn rewrite_pypi_links(html: &str, package_name: &str) -> String {
let mut result = String::with_capacity(html.len());
let mut remaining = html;
while let Some(href_start) = remaining.find("href=\"") {
result.push_str(&remaining[..href_start + 6]);
remaining = &remaining[href_start + 6..];
if let Some(href_end) = remaining.find('"') {
let url = &remaining[..href_end];
if let Some(filename) = extract_filename(url) {
result.push_str(&format!("/simple/{}/{}", package_name, filename));
} else {
result.push_str(url);
}
remaining = &remaining[href_end..];
}
}
result.push_str(remaining);
let result = remove_attribute(&result, "data-core-metadata");
remove_attribute(&result, "data-dist-info-metadata")
}
fn remove_attribute(html: &str, attr_name: &str) -> String {
let mut result = String::with_capacity(html.len());
let mut remaining = html;
let pattern = format!(" {}=\"", attr_name);
while let Some(attr_start) = remaining.find(&pattern) {
result.push_str(&remaining[..attr_start]);
remaining = &remaining[attr_start + pattern.len()..];
if let Some(attr_end) = remaining.find('"') {
remaining = &remaining[attr_end + 1..];
}
}
result.push_str(remaining);
result
}
fn extract_filename(url: &str) -> Option<&str> {
let url = url.split('#').next()?;
let filename = url.rsplit('/').next()?;
if filename.ends_with(".tar.gz")
|| filename.ends_with(".tgz")
|| filename.ends_with(".whl")
|| filename.ends_with(".zip")
|| filename.ends_with(".egg")
{
Some(filename)
} else {
None
}
}
fn find_file_url(html: &str, target_filename: &str) -> Option<String> {
let mut remaining = html;
while let Some(href_start) = remaining.find("href=\"") {
remaining = &remaining[href_start + 6..];
if let Some(href_end) = remaining.find('"') {
let url = &remaining[..href_end];
if let Some(filename) = extract_filename(url) {
if filename == target_filename {
return Some(url.split('#').next().unwrap_or(url).to_string());
}
}
remaining = &remaining[href_end..];
}
}
None
}