main.rs

use chaser_oxide::{Browser, BrowserConfig, ChaserPage, ChaserProfile, Gpu};
use chaser_oxide::handler::viewport::Viewport;
use chaser_oxide::cdp::browser_protocol::network::{EnableParams, EventResponseReceived, GetResponseBodyParams};
use chaser_oxide::cdp::browser_protocol::emulation::{SetDeviceMetricsOverrideParams, SetGeolocationOverrideParams};
use chaser_oxide::cdp::browser_protocol::page::CaptureScreenshotParams;
use futures::StreamExt;
use serde::{Deserialize, Serialize};
use std::convert::Infallible;
use std::env;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::sync::Mutex;
use warp::Filter;

type SharedBrowser = Arc<Mutex<Browser>>;

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct AdvancedAction {
    r#type: String,
    selector: Option<String>,
    value: Option<String>,
    file_name: Option<String>,
    full: Option<bool>,
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct Payload {
    width: Option<u32>,
    height: Option<u32>,
    sniff_network: Option<bool>,
    actions: Vec<AdvancedAction>,
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct StandardResponse {
    logs: Vec<String>,
    result: serde_json::Value,
    error: Option<String>,
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct MemoryStats {
    total_system: String,
    free_system: String,
    allocated_by_api: String,
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct StatsResponse {
    uptime: String,
    engine_version: String,
    system_load_average: String,
    active_tabs: usize,
    memory: MemoryStats,
}

#[derive(Serialize, Clone)]
#[serde(rename_all = "camelCase")]
struct SniffedActivity {
    url: String,
    method: String,
    response_status: i32,
    response_body: String,
}

fn format_uptime(secs: u64) -> String {
    let hours = secs / 3600;
    let minutes = (secs % 3600) / 60;
    let seconds = secs % 60;
    format!("{}h {}m {}s", hours, minutes, seconds)
}

fn format_size(bytes: u64) -> String {
    if bytes >= 1073741824 {
        format!("{:.2} GB", bytes as f64 / 1073741824.0)
    } else {
        format!("{:.2} MB", bytes as f64 / 1048576.0)
    }
}

fn get_current_memory_bytes() -> u64 {
    if let Ok(statm) = std::fs::read_to_string("/proc/self/statm") {
        if let Some(first_val) = statm.split_whitespace().next() {
            if let Ok(pages) = first_val.parse::<u64>() {
                return pages * 4096;
            }
        }
    }
    0
}

fn get_system_memory_bytes() -> (u64, u64) {
    let mut total = 0;
    let mut free = 0;
    if let Ok(meminfo) = std::fs::read_to_string("/proc/meminfo") {
        for line in meminfo.lines() {
            if line.starts_with("MemTotal:") {
                total = line.split_whitespace().nth(1).unwrap_or("0").parse::<u64>().unwrap_or(0) * 1024;
            } else if line.starts_with("MemAvailable:") || line.starts_with("MemFree:") {
                free = line.split_whitespace().nth(1).unwrap_or("0").parse::<u64>().unwrap_or(0) * 1024;
                if line.starts_with("MemAvailable:") { break; }
            }
        }
    }
    (total, free)
}

fn get_system_load() -> String {
    if let Ok(loadavg) = std::fs::read_to_string("/proc/loadavg") {
        let parts: Vec<&str> = loadavg.split_whitespace().take(3).collect();
        if parts.len() == 3 {
            return format!("{} {} {}", parts[0], parts[1], parts[2]);
        }
    }
    "0.00 0.00 0.00".to_string()
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let start_time = std::time::Instant::now();
    let port_str = env::var("PORT").unwrap_or_else(|_| "7860".to_string());
    let port: u16 = port_str.parse().unwrap_or(7860);
    let chrome_bin = env::var("CHROME_PATH").unwrap_or_else(|_| "/usr/bin/google-chrome-stable".to_string());

    let config = BrowserConfig::builder()
        .chrome_executable(chrome_bin)
        .no_sandbox()
        .new_headless_mode()
        .arg("--disable-gpu")
        .arg("--disable-dev-shm-usage")
        .arg("--disable-gl-drawing-for-tests")
        .arg("--disable-software-rasterizer")
        .arg("--no-zygote")
        .arg("--single-process")
        .arg("--disable-blink-features=AutomationControlled")
        .arg("--user-data-dir=/tmp/isolated_chrome_profile")
        .arg("--disable-infobars")
        .arg("--password-store=basic")
        .arg("--use-mock-keychain")
        .viewport(Viewport {
            width: 1366,
            height: 768,
            ..Default::default()
        })
        .build()
        .map_err(|e| anyhow::anyhow!(e))?;

    let (browser, mut handler) = Browser::launch(config).await?;
    tokio::spawn(async move {
        while let Some(_) = handler.next().await {}
    });

    let shared_browser = Arc::new(Mutex::new(browser));
    
    let execute_browser_filter = shared_browser.clone();
    let execute_route = warp::post()
        .and(warp::path("execute"))
        .and(warp::body::json())
        .and(warp::any().map(move || execute_browser_filter.clone()))
        .and_then(handle_execute);

    let stats_browser_filter = shared_browser.clone();
    let stats_route = warp::get()
        .and(warp::path::end())
        .and(warp::any().map(move || stats_browser_filter.clone()))
        .and_then(move |b| handle_stats(start_time, b));

    let tmp_dir = env::temp_dir();
    let files_route = warp::path("tmp").and(warp::fs::dir(tmp_dir.clone()));

    let routes = execute_route.or(stats_route).or(files_route);
    let addr: SocketAddr = format!("0.0.0.0:{}", port).parse().unwrap();
    
    warp::serve(routes).run(addr).await;

    Ok(())
}

async fn handle_stats(start_time: std::time::Instant, browser_info: SharedBrowser) -> Result<impl warp::Reply, Infallible> {
    let (total_mem, free_mem) = get_system_memory_bytes();
    let api_mem = get_current_memory_bytes();
    let pkg_version = format!("rust-warp-v{}", env!("CARGO_PKG_VERSION"));
    
    let total_tabs = {
        let browser = browser_info.lock().await;
        if let Ok(pages) = browser.pages().await {
            pages.len()
        } else {
            1
        }
    };

    Ok(warp::reply::json(&StatsResponse {
        uptime: format_uptime(start_time.elapsed().as_secs()),
        engine_version: pkg_version,
        system_load_average: get_system_load(),
        active_tabs: total_tabs,
        memory: MemoryStats {
            total_system: format_size(total_mem),
            free_system: format_size(free_mem),
            allocated_by_api: format_size(api_mem),
        },
    }))
}

async fn handle_execute(payload: Payload, browser_info: SharedBrowser) -> Result<impl warp::Reply, Infallible> {
    let mut logs = Vec::new();
    let mut screenshots = Vec::new();
    let mut evals = Vec::new();
    let network_activities = Arc::new(Mutex::new(Vec::<SniffedActivity>::new()));
    
    let width = payload.width.unwrap_or(1366);
    let height = payload.height.unwrap_or(768);

    let first_url = payload.actions.iter()
        .find(|act| act.r#type == "navigate")
        .and_then(|act| act.value.clone())
        .unwrap_or_else(|| "about:blank".to_string());

    let browser = browser_info.lock().await;
    let page = browser.new_page(&first_url).await.unwrap();
    
    let _ = page.execute(SetDeviceMetricsOverrideParams::builder()
        .width(width as i32)
        .height(height as i32)
        .device_scale_factor(1.0)
        .mobile(false)
        .build()
        .unwrap()
    ).await;

    if payload.sniff_network.unwrap_or(false) {
        let _ = page.execute(EnableParams::default()).await;
        let mut response_rcvd_stream = page.event_listener::<EventResponseReceived>().await.unwrap();
        let net_logs_clone = network_activities.clone();
        let page_clone = page.clone();
        
        tokio::spawn(async move {
            while let Some(ev) = response_rcvd_stream.next().await {
                let req_id = ev.request_id.clone();
                let url = ev.response.url.clone();
                let status = ev.response.status;
                
                if url.contains("json") || url.contains("api") {
                    if let Ok(body_res) = page_clone.execute(GetResponseBodyParams { request_id: req_id }).await {
                        let mut lock = net_logs_clone.lock().await;
                        lock.push(SniffedActivity {
                            url,
                            method: "GET/POST".to_string(),
                            response_status: status as i32,
                            response_body: body_res.body.clone(),
                        });
                    }
                }
            }
        });
    }

    let chaser = ChaserPage::new(page.clone());

    let profile = ChaserProfile::windows()
        .chrome_version(131)
        .gpu(Gpu::NvidiaRTX3080)
        .memory_gb(16)
        .cpu_cores(8)
        .locale("en-US")
        .timezone("America/New_York")
        .screen(width, height)
        .build();

    chaser.apply_profile(&profile).await.unwrap();

    let os_tmp_dir = env::temp_dir();
    let mut turnstile_sitekey = String::new();
    let mut turnstile_token = String::new();

    for (i, act) in payload.actions.iter().enumerate() {
        let mut log_msg = format!("Step {}: [{}]", i + 1, act.r#type);
        if let Some(sel) = &act.selector {
            log_msg.push_str(&format!(" -> Target: {}", sel));
        }
        gather_logs(&mut logs, log_msg);
        
        let inner_page = chaser.raw_page();

        match act.r#type.as_str() {
            "navigate" => {
                if first_url != "about:blank" && i == 0 {
                    continue;
                }
                if let Some(url) = &act.value {
                    chaser.goto(url).await.unwrap();
                }
            }
            "wait" => {
                if let Some(sel) = &act.selector {
                    let _ = inner_page.find_element(sel).await.unwrap();
                }
            }
            "click" => {
                if let Some(sel) = &act.selector {
                    if let Ok(el) = inner_page.find_element(sel).await {
                        el.click().await.unwrap();
                    }
                }
            }
            "type" => {
                if let Some(sel) = &act.selector {
                    if let Some(val) = &act.value {
                        if let Ok(el) = inner_page.find_element(sel).await {
                            el.type_str(val).await.unwrap();
                        }
                    }
                }
            }
            "clear" => {
                if let Some(sel) = &act.selector {
                    let js_clear = format!("(() => {{ const el = document.querySelector('{}'); if(el) {{ el.value = ''; el.dispatchEvent(new Event('input', {{bubbles:true}})); }} }})()", sel);
                    let _: Option<serde_json::Value> = chaser.evaluate(&js_clear).await.unwrap();
                }
            }
            "hover" => {
                if let Some(sel) = &act.selector {
                    if let Ok(el) = inner_page.find_element(sel).await {
                        el.hover().await.unwrap();
                    }
                }
            }
            "select" => {
                if let Some(sel) = &act.selector {
                    if let Some(val) = &act.value {
                        let js_select = format!("(() => {{ const el = document.querySelector('{}'); if(el) {{ el.value = '{}'; el.dispatchEvent(new Event('change', {{bubbles:true}})); }} }})()", sel, val);
                        let _: Option<serde_json::Value> = chaser.evaluate(&js_select).await.unwrap();
                    }
                }
            }
            "key" => {
                if let Some(val) = &act.value {
                    let js_input = format!(
                        "(() => {{ const activeEl = document.activeElement; if (activeEl) {{ activeEl.value += '{}'; activeEl.dispatchEvent(new Event('input', {{ bubbles: true }})); activeEl.dispatchEvent(new Event('change', {{ bubbles: true }})); }} }})()",
                        val
                    );
                    let _: Option<serde_json::Value> = chaser.evaluate(&js_input).await.unwrap();
                }
            }
            "scroll" => {
                if let Some(sel) = &act.selector {
                    if let Ok(el) = inner_page.find_element(sel).await {
                        el.scroll_into_view().await.unwrap();
                    }
                } else if let Some(val) = &act.value {
                    if let Ok(val_int) = val.parse::<i32>() {
                        let js_scroll = format!("window.scrollTo(0, {});", val_int);
                        let _: Option<serde_json::Value> = chaser.evaluate(&js_scroll).await.unwrap();
                    }
                }
            }
            "setLocation" => {
                if let Some(val) = &act.value {
                    let coords: Vec<&str> = val.split(',').collect();
                    if coords.len() == 2 {
                        if let (Ok(lat), Ok(long)) = (coords[0].parse::<f64>(), coords[1].parse::<f64>()) {
                            let _ = inner_page.execute(SetGeolocationOverrideParams::builder()
                                .latitude(lat)
                                .longitude(long)
                                .accuracy(100.0)
                                .build()
                            ).await;
                        }
                    }
                }
            }
            "extractTurnstile" => {
                let harvest_script = r#"
                    (() => {
                        const iframe = document.querySelector('iframe[src*="challenges.cloudflare.com"]');
                        let sitekey = "";
                        if (iframe) {
                            const urlParams = new URLSearchParams(new URL(iframe.src).hash ? iframe.src : iframe.src.replace('#', '?'));
                            sitekey = urlParams.get('k') || "";
                        }
                        const input = document.querySelector('input[name="cf-turnstile-response"]') || 
                                      document.getElementById('cf-turnstile-response');
                        return JSON.stringify({
                            sitekey: sitekey,
                            token: input ? input.value : ""
                        });
                    })()
                "#;
                if let Ok(Some(res_json)) = chaser.evaluate(harvest_script).await {
                    if let Some(res_str) = res_json.as_str() {
                        if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(res_str) {
                            turnstile_sitekey = parsed["sitekey"].as_str().unwrap_or("").to_string();
                            turnstile_token = parsed["token"].as_str().unwrap_or("").to_string();
                        }
                    }
                }
            }
            "evaluate" => {
                if let Some(js_code) = &act.value {
                    if let Ok(Some(res)) = chaser.evaluate(js_code).await {
                        evals.push(res);
                    }
                }
            }
            "screenshot" => {
                let is_full_page = act.full.unwrap_or(false);
                let buf = inner_page.screenshot(
                    CaptureScreenshotParams::builder()
                        .from_surface(true)
                        .capture_beyond_viewport(is_full_page)
                        .build()
                ).await.unwrap();
                
                let mut name = act.file_name.clone().unwrap_or_else(|| format!("ss-{}.png", std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_nanos()));
                name = name.replace(" ", "_");
                let file_path = os_tmp_dir.join(&name);
                std::fs::write(file_path, buf).unwrap();
                screenshots.push(format!("/tmp/{}", name));
            }
            _ => {}
        }
    }

    let final_net_logs = network_activities.lock().await.clone();

    let res_val = serde_json::json!({
        "screenshots": screenshots,
        "evals": evals,
        "turnstile": {
            "sitekey": turnstile_sitekey,
            "token": turnstile_token
        },
        "networkLogs": final_net_logs
    });

    page.close().await.unwrap();

    Ok(warp::reply::json(&StandardResponse {
        logs,
        result: res_val,
        error: None,
    }))
}

fn gather_logs(logs: &mut Vec<String>, msg: String) {
    logs.push(msg);
}