hottub/src/proxies/noodlemagazine.rs

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

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

#[derive(Debug, Clone)]
pub struct NoodlemagazineProxy {}

impl NoodlemagazineProxy {
    pub fn new() -> Self {
        NoodlemagazineProxy {}
    }

    fn extract_playlist(text: &str) -> Option<&str> {
        text.split("window.playlist = ").nth(1)?.split(';').next()
    }

    fn source_score(source: &Value) -> (u8, u32) {
        let file = source["file"].as_str().unwrap_or_default();
        let label = source["label"].as_str().unwrap_or_default();
        let is_hls = u8::from(file.contains(".m3u8"));
        let quality = label
            .chars()
            .filter(|c| c.is_ascii_digit())
            .collect::<String>()
            .parse::<u32>()
            .unwrap_or(0);

        (is_hls, quality)
    }

    fn select_best_source(playlist: &str) -> Option<String> {
        let json: Value = serde_json::from_str(playlist).ok()?;
        let sources = json["sources"].as_array()?;

        sources
            .iter()
            .filter(|source| {
                source["file"]
                    .as_str()
                    .map(|file| !file.is_empty())
                    .unwrap_or(false)
            })
            .max_by_key(|source| Self::source_score(source))
            .and_then(|source| source["file"].as_str())
            .map(str::to_string)
    }

    fn normalize_video_page_url(url: &str) -> String {
        if url.starts_with("http://") || url.starts_with("https://") {
            url.to_string()
        } else {
            format!("https://{}", url.trim_start_matches('/'))
        }
    }

    fn is_hls_url(url: &str) -> bool {
        Url::parse(url)
            .ok()
            .map(|parsed| parsed.path().ends_with(".m3u8"))
            .unwrap_or(false)
    }

    fn absolutize_uri(base_url: &Url, value: &str) -> String {
        if value.is_empty() {
            return String::new();
        }

        if value.starts_with('#')
            || value.starts_with("data:")
            || value.starts_with("http://")
            || value.starts_with("https://")
        {
            return value.to_string();
        }

        base_url
            .join(value)
            .map(|url| url.to_string())
            .unwrap_or_else(|_| value.to_string())
    }

    fn rewrite_manifest_line(base_url: &Url, line: &str) -> String {
        if line.trim().is_empty() {
            return line.to_string();
        }

        if !line.starts_with('#') {
            return Self::absolutize_uri(base_url, line);
        }

        let Some(uri_start) = line.find("URI=\"") else {
            return line.to_string();
        };
        let value_start = uri_start + 5;
        let Some(relative_end) = line[value_start..].find('"') else {
            return line.to_string();
        };
        let value_end = value_start + relative_end;
        let value = &line[value_start..value_end];
        let rewritten = Self::absolutize_uri(base_url, value);

        format!(
            "{}{}{}",
            &line[..value_start],
            rewritten,
            &line[value_end..]
        )
    }

    fn rewrite_manifest(manifest_url: &str, body: &str) -> Option<String> {
        let base_url = Url::parse(manifest_url).ok()?;

        Some(
            body.lines()
                .map(|line| Self::rewrite_manifest_line(&base_url, line))
                .collect::<Vec<_>>()
                .join("\n"),
        )
    }

    async fn resolve_source_url(
        &self,
        url: String,
        requester: web::types::State<Requester>,
    ) -> Option<(String, String)> {
        let mut requester = requester.get_ref().clone();
        let url = Self::normalize_video_page_url(&url);
        let text = requester
            .get(&url, Some(Version::HTTP_2))
            .await
            .unwrap_or_default();
        if text.is_empty() {
            return None;
        }

        let Some(playlist) = Self::extract_playlist(&text) else {
            return None;
        };

        Self::select_best_source(playlist).map(|source_url| (url, source_url))
    }
}

pub async fn serve_media(
    req: HttpRequest,
    requester: web::types::State<Requester>,
) -> Result<impl web::Responder, web::Error> {
    let endpoint = req.match_info().query("endpoint").to_string();
    let proxy = NoodlemagazineProxy::new();
    let Some((video_page_url, source_url)) =
        proxy.resolve_source_url(endpoint, requester.clone()).await
    else {
        return Ok(web::HttpResponse::BadGateway().finish());
    };

    if !NoodlemagazineProxy::is_hls_url(&source_url) {
        return Ok(web::HttpResponse::Found()
            .header("Location", source_url)
            .finish());
    }

    let mut upstream_requester = requester.get_ref().clone();
    let upstream = match upstream_requester
        .get_raw_with_headers(&source_url, vec![("Referer".to_string(), video_page_url)])
        .await
    {
        Ok(response) => response,
        Err(_) => return Ok(web::HttpResponse::BadGateway().finish()),
    };

    let manifest_body = upstream.text().await.map_err(error::ErrorBadGateway)?;
    let rewritten_manifest =
        match NoodlemagazineProxy::rewrite_manifest(&source_url, &manifest_body) {
            Some(body) => body,
            None => return Ok(web::HttpResponse::BadGateway().finish()),
        };

    Ok(web::HttpResponse::Ok()
        .header(CONTENT_TYPE, "application/vnd.apple.mpegurl")
        .body(rewritten_manifest))
}

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

    #[test]
    fn extracts_playlist_from_page() {
        let html = r#"
            <script>
                window.playlist = {"sources":[{"file":"https://cdn.example/360.mp4","label":"360p"}]};
            </script>
        "#;

        assert_eq!(
            NoodlemagazineProxy::extract_playlist(html),
            Some(r#"{"sources":[{"file":"https://cdn.example/360.mp4","label":"360p"}]}"#)
        );
    }

    #[test]
    fn prefers_hls_then_highest_quality() {
        let playlist = r#"{
            "sources": [
                {"file":"https://cdn.example/360.mp4","label":"360p"},
                {"file":"https://cdn.example/720.mp4","label":"720p"},
                {"file":"https://cdn.example/master.m3u8","label":"1080p"}
            ]
        }"#;

        assert_eq!(
            NoodlemagazineProxy::select_best_source(playlist).as_deref(),
            Some("https://cdn.example/master.m3u8")
        );
    }

    #[test]
    fn rewrites_manifest_to_direct_absolute_urls() {
        let manifest = "#EXTM3U\n#EXT-X-STREAM-INF:BANDWIDTH=1\nlow/index.m3u8\n#EXT-X-KEY:METHOD=AES-128,URI=\"keys/key.bin\"\nsegment0.ts";

        let rewritten =
            NoodlemagazineProxy::rewrite_manifest("https://cdn.example/hls/master.m3u8", manifest)
                .unwrap();

        assert_eq!(
            rewritten,
            "#EXTM3U\n#EXT-X-STREAM-INF:BANDWIDTH=1\nhttps://cdn.example/hls/low/index.m3u8\n#EXT-X-KEY:METHOD=AES-128,URI=\"https://cdn.example/hls/keys/key.bin\"\nhttps://cdn.example/hls/segment0.ts"
        );
    }
}