hottub/src/proxies/pimpbunnythumb.rs

use ntex::http::header::{CONTENT_LENGTH, CONTENT_TYPE};
use ntex::{
    http::Response,
    web::{self, HttpRequest, error},
};
use url::Url;
use wreq::Version;

use crate::util::requester::Requester;

const FIREFOX_USER_AGENT: &str =
    "Mozilla/5.0 (X11; Linux x86_64; rv:147.0) Gecko/20100101 Firefox/147.0";
const HTML_ACCEPT: &str =
    "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8";
const IMAGE_ACCEPT: &str = "image/avif,image/webp,image/png,image/svg+xml,image/*;q=0.8,*/*;q=0.5";

fn root_referer() -> &'static str {
    "https://pimpbunny.com/"
}

fn root_html_headers() -> Vec<(String, String)> {
    vec![
        ("Referer".to_string(), root_referer().to_string()),
        ("User-Agent".to_string(), FIREFOX_USER_AGENT.to_string()),
        ("Accept".to_string(), HTML_ACCEPT.to_string()),
        ("Accept-Language".to_string(), "en-US,en;q=0.9".to_string()),
    ]
}

fn image_headers(requester: &Requester, image_url: &str) -> Vec<(String, String)> {
    let mut headers = vec![
        ("Referer".to_string(), root_referer().to_string()),
        ("User-Agent".to_string(), FIREFOX_USER_AGENT.to_string()),
        ("Accept".to_string(), IMAGE_ACCEPT.to_string()),
        ("Accept-Language".to_string(), "en-US,en;q=0.9".to_string()),
    ];
    if let Some(cookie) = requester.cookie_header_for_url(image_url) {
        headers.push(("Cookie".to_string(), cookie));
    }
    headers
}

fn is_allowed_thumb_url(url: &str) -> bool {
    let Some(url) = Url::parse(url).ok() else {
        return false;
    };
    if url.scheme() != "https" {
        return false;
    }
    let Some(host) = url.host_str() else {
        return false;
    };

    matches!(host, "pimpbunny.com" | "www.pimpbunny.com")
        && url.path().starts_with("/contents/videos_screenshots/")
}

pub async fn get_image(
    req: HttpRequest,
    requester: web::types::State<Requester>,
) -> Result<impl web::Responder, web::Error> {
    let endpoint = req.match_info().query("endpoint").to_string();
    let image_url = if endpoint.starts_with("http://") || endpoint.starts_with("https://") {
        endpoint
    } else {
        format!("https://{}", endpoint.trim_start_matches('/'))
    };

    if !is_allowed_thumb_url(&image_url) {
        return Ok(web::HttpResponse::BadRequest().finish());
    }

    let mut requester = requester.get_ref().clone();
    let _ = requester
        .get_with_headers(root_referer(), root_html_headers(), Some(Version::HTTP_11))
        .await;

    let mut headers = image_headers(&requester, image_url.as_str());

    let mut upstream = requester
        .get_raw_with_headers(image_url.as_str(), headers.clone())
        .await
        .ok();

    let needs_warmup = upstream
        .as_ref()
        .map(|response| !response.status().is_success())
        .unwrap_or(true);

    if needs_warmup {
        let _ = requester
            .get_with_headers(root_referer(), root_html_headers(), Some(Version::HTTP_11))
            .await;
        headers = image_headers(&requester, image_url.as_str());

        upstream = requester
            .get_raw_with_headers(image_url.as_str(), headers.clone())
            .await
            .ok();
    }

    let needs_image_specific_warmup = upstream
        .as_ref()
        .map(|response| !response.status().is_success())
        .unwrap_or(true);

    if needs_image_specific_warmup {
        let _ = requester
            .get_with_headers(image_url.as_str(), headers.clone(), Some(Version::HTTP_11))
            .await;
        headers = image_headers(&requester, image_url.as_str());

        upstream = requester
            .get_raw_with_headers(image_url.as_str(), headers)
            .await
            .ok();
    }

    let Some(upstream) = upstream.filter(|response| response.status().is_success()) else {
        return Ok(web::HttpResponse::NotFound().finish());
    };

    let status = upstream.status();
    let headers = upstream.headers().clone();
    let bytes = upstream.bytes().await.map_err(error::ErrorBadGateway)?;

    let mut resp = Response::build(status);

    if let Some(ct) = headers.get(CONTENT_TYPE) {
        if let Ok(ct_str) = ct.to_str() {
            resp.set_header(CONTENT_TYPE, ct_str);
        }
    }
    if let Some(cl) = headers.get(CONTENT_LENGTH) {
        if let Ok(cl_str) = cl.to_str() {
            resp.set_header(CONTENT_LENGTH, cl_str);
        }
    }

    Ok(resp.body(bytes.to_vec()))
}

#[cfg(test)]
mod tests {
    use super::is_allowed_thumb_url;

    #[test]
    fn allows_expected_pimpbunny_thumb_paths() {
        assert!(is_allowed_thumb_url(
            "https://pimpbunny.com/contents/videos_screenshots/517000/517329/800x450/1.jpg"
        ));
        assert!(is_allowed_thumb_url(
            "https://www.pimpbunny.com/contents/videos_screenshots/1/2/800x450/3.webp"
        ));
    }

    #[test]
    fn rejects_non_thumb_or_non_pimpbunny_urls() {
        assert!(!is_allowed_thumb_url(
            "http://pimpbunny.com/contents/videos_screenshots/x.jpg"
        ));
        assert!(!is_allowed_thumb_url(
            "https://pimpbunny.com/videos/example-video/"
        ));
        assert!(!is_allowed_thumb_url(
            "https://example.com/contents/videos_screenshots/x.jpg"
        ));
    }
}