hottub/src/util/requester.rs

use serde::Serialize;
use std::env;
use wreq::Client;
use wreq::Proxy;
use wreq::Response;
use wreq::Version;
use wreq::header::HeaderValue;
use wreq::redirect::Policy;
use wreq_util::Emulation;

use crate::util::flaresolverr::FlareSolverrRequest;
use crate::util::flaresolverr::Flaresolverr;

// A Send + Sync error type for all async paths
type AnyErr = Box<dyn std::error::Error + Send + Sync + 'static>;

#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
pub struct Requester {
    #[serde(skip)]
    client: Client,
    proxy: bool,
    flaresolverr_session: Option<String>,
}

impl Requester {
    pub fn new() -> Self {
        let client = Client::builder()
            .cert_verification(false)
            .emulation(Emulation::Firefox136)
            .cookie_store(true)
            .redirect(Policy::default())
            .build()
            .expect("Failed to create HTTP client");

        Requester {
            client,
            proxy: false,
            flaresolverr_session: None,
        }
    }

    pub fn set_proxy(&mut self, proxy: bool) {
        if proxy {
            println!("Proxy enabled");
        }
        self.proxy = proxy;
    }

    // pub fn set_flaresolverr_session(&mut self, session: String) {
    //     self.flaresolverr_session = Some(session);
    // }

    // fn get_url_from_location_header(&self, prev_url: &str, location: &str) -> String {
    //     if location.starts_with("http://") || location.starts_with("https://") {
    //         location.to_string()
    //     } else if location.starts_with("//") {
    //         format!("{}{}", "https:", location)
    //     } else if location.starts_with('/') {
    //         let base_url = prev_url.split('/').take(3).collect::<Vec<&str>>().join("/");
    //         format!("{}{}", base_url, location)
    //     } else {
    //         format!("{}/{}", prev_url, location)
    //     }
    // }

    pub async fn get_raw(&mut self, url: &str) -> Result<Response, wreq::Error> {
        let client = Client::builder()
            .cert_verification(false)
            .emulation(Emulation::Firefox136)
            .cookie_store(true)
            .build()
            .expect("Failed to create HTTP client");

        let mut request = client.get(url).version(Version::HTTP_11);

        if self.proxy {
            if let Ok(proxy_url) = env::var("BURP_URL") {
                let proxy = Proxy::all(&proxy_url).unwrap();
                request = request.proxy(proxy);
            }
        }

        request.send().await
    }

    pub async fn post<S>(&mut self, url: &str, data: &S) -> Result<Response, wreq::Error>
    where
        S: Serialize + ?Sized,
    {
        let client = Client::builder()
            .cert_verification(false)
            .emulation(Emulation::Firefox136)
            .cookie_store(true)
            .build()
            .expect("Failed to create HTTP client");

        let mut request = client.post(url).version(Version::HTTP_11).json(data);

        if self.proxy {
            if let Ok(proxy_url) = env::var("BURP_URL") {
                let proxy = Proxy::all(&proxy_url).unwrap();
                request = request.proxy(proxy);
            }
        }

        request.send().await
    }

    pub async fn get(&mut self, url: &str) -> Result<String, AnyErr> {
        let mut request = self.client.get(url).version(Version::HTTP_11);

        if self.proxy {
            if let Ok(proxy_url) = env::var("BURP_URL") {
                let proxy = Proxy::all(&proxy_url).unwrap();
                request = request.proxy(proxy);
            }
        }

        let response = request.send().await?;

        if response.status().is_success() {
            return Ok(response.text().await?);
        }

        // If direct request failed, try FlareSolverr. Map its error to a Send+Sync error immediately,
        // so no non-Send error value lives across later `.await`s.
        let flare_url = env::var("FLARE_URL").expect("FLARE_URL not set");
        let mut flare = Flaresolverr::new(flare_url);
        if self.proxy && env::var("BURP_URL").is_ok() {
            flare.set_proxy(true);
        }

        let res = flare
            .solve(FlareSolverrRequest {
                cmd: "request.get".to_string(),
                url: url.to_string(),
                maxTimeout: 60000,
            })
            .await
            .map_err(|e| -> AnyErr { format!("Failed to solve FlareSolverr: {e}").into() })?;

        // Rebuild client and apply UA/cookies from FlareSolverr
        let cookie_origin = url.split('/').take(3).collect::<Vec<&str>>().join("/");

        self.client = Client::builder()
            .cert_verification(false)
            .emulation(Emulation::Firefox136)
            .cookie_store(true)
            .redirect(Policy::default())
            .build()
            .expect("Failed to create HTTP client");

        let useragent = res.solution.userAgent;
        self.client
            .update()
            .headers(|headers| {
                headers.insert("User-Agent", HeaderValue::from_str(&useragent).unwrap());
            })
            .apply()
            .unwrap();

        if let Ok(origin) = url::Url::parse(&cookie_origin) {
            for cookie in res.solution.cookies {
                let header =
                    HeaderValue::from_str(&format!("{}={}", cookie.name, cookie.value)).unwrap();
                self.client.set_cookie(&origin, header);
            }
        }

        // Retry the original URL with the updated client & (optional) proxy
        let mut request = self.client.get(url).version(Version::HTTP_11);
        if self.proxy {
            if let Ok(proxy_url) = env::var("BURP_URL") {
                let proxy = Proxy::all(&proxy_url).unwrap();
                request = request.proxy(proxy);
            }
        }

        let response = request.send().await?;
        if response.status().is_success() {
            return Ok(response.text().await?);
        }

        // Fall back to FlareSolverr-provided body
        Ok(res.solution.response)
    }
}