headless_engine/dom/
screenshot.rs

1use crate::dom::interactive::InteractiveElement;
2use crate::dom::SearchResults;
3use anyhow::Result;
4use base64::Engine;
5use scraper::{Html, Selector};
6use serde::{Deserialize, Serialize};
7use std::path::PathBuf;
8
9use std::sync::atomic::{AtomicUsize, Ordering};
10
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct ScreenshotResult {
13    pub width: u32,
14    pub height: u32,
15    pub svg: String,
16    pub layout_wireframe: String,
17    pub element_count: usize,
18    #[serde(skip_serializing_if = "Vec::is_empty")]
19    pub png_bytes: Vec<u8>,
20    pub png_base64: String,
21}
22
23pub struct RealBrowserScreenshot;
24
25struct TempFileCleanup(PathBuf);
26impl Drop for TempFileCleanup {
27    fn drop(&mut self) {
28        let _ = std::fs::remove_file(&self.0);
29    }
30}
31
32impl RealBrowserScreenshot {
33    pub fn find_browser_binary() -> Option<PathBuf> {
34        #[cfg(target_os = "windows")]
35        {
36            let candidates = [
37                r"C:\Program Files\Google\Chrome\Application\chrome.exe",
38                r"C:\Program Files (x86)\Google\Chrome\Application\chrome.exe",
39                r"C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe",
40                r"C:\Program Files\Microsoft\Edge\Application\msedge.exe",
41            ];
42            for path_str in &candidates {
43                let p = std::path::Path::new(path_str);
44                if p.exists() {
45                    return Some(p.to_path_buf());
46                }
47            }
48            if let Ok(local_app_data) = std::env::var("LOCALAPPDATA") {
49                let chrome_local = std::path::Path::new(&local_app_data)
50                    .join(r"Google\Chrome\Application\chrome.exe");
51                if chrome_local.exists() {
52                    return Some(chrome_local);
53                }
54                let edge_local = std::path::Path::new(&local_app_data)
55                    .join(r"Microsoft\Edge\Application\msedge.exe");
56                if edge_local.exists() {
57                    return Some(edge_local);
58                }
59            }
60            // Dynamic fallback
61            if let Ok(output) = std::process::Command::new("where").arg("chrome").output() {
62                if output.status.success() {
63                    let path_str = String::from_utf8_lossy(&output.stdout)
64                        .lines()
65                        .next()
66                        .unwrap_or("")
67                        .trim()
68                        .to_string();
69                    if !path_str.is_empty() {
70                        return Some(PathBuf::from(path_str));
71                    }
72                }
73            }
74            if let Ok(output) = std::process::Command::new("where").arg("msedge").output() {
75                if output.status.success() {
76                    let path_str = String::from_utf8_lossy(&output.stdout)
77                        .lines()
78                        .next()
79                        .unwrap_or("")
80                        .trim()
81                        .to_string();
82                    if !path_str.is_empty() {
83                        return Some(PathBuf::from(path_str));
84                    }
85                }
86            }
87        }
88
89        #[cfg(target_os = "macos")]
90        {
91            let candidates = [
92                "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome",
93                "/Applications/Microsoft Edge.app/Contents/MacOS/Microsoft Edge",
94                "/Applications/Chromium.app/Contents/MacOS/Chromium",
95            ];
96            for path_str in &candidates {
97                let p = std::path::Path::new(path_str);
98                if p.exists() {
99                    return Some(p.to_path_buf());
100                }
101            }
102            // Dynamic fallback
103            if let Ok(output) = std::process::Command::new("which")
104                .arg("google-chrome")
105                .output()
106            {
107                if output.status.success() {
108                    let path_str = String::from_utf8_lossy(&output.stdout).trim().to_string();
109                    if !path_str.is_empty() {
110                        return Some(PathBuf::from(path_str));
111                    }
112                }
113            }
114        }
115
116        #[cfg(target_os = "linux")]
117        {
118            let candidates = [
119                "google-chrome",
120                "google-chrome-stable",
121                "chromium",
122                "chromium-browser",
123            ];
124            for bin in &candidates {
125                if let Ok(output) = std::process::Command::new("which").arg(bin).output() {
126                    if output.status.success() {
127                        let path_str = String::from_utf8_lossy(&output.stdout).trim().to_string();
128                        if !path_str.is_empty() {
129                            return Some(PathBuf::from(path_str));
130                        }
131                    }
132                }
133            }
134        }
135
136        None
137    }
138
139    fn generate_temp_png_path() -> PathBuf {
140        let temp_dir = std::env::temp_dir();
141        let pid = std::process::id();
142        let time = std::time::SystemTime::now()
143            .duration_since(std::time::UNIX_EPOCH)
144            .map(|d| d.as_millis())
145            .unwrap_or(0);
146        static COUNTER: AtomicUsize = AtomicUsize::new(0);
147        let count = COUNTER.fetch_add(1, Ordering::Relaxed);
148        temp_dir.join(format!("headless_shot_{}_{}_{}.png", pid, time, count))
149    }
150
151    fn is_valid_url(url: &str) -> bool {
152        url.starts_with("http://") || url.starts_with("https://") || url.starts_with("file://")
153    }
154
155    pub async fn capture_real_screenshot_async(
156        url: &str,
157        html_str: &str,
158        width: u32,
159        height: u32,
160    ) -> Option<ScreenshotResult> {
161        if !Self::is_valid_url(url) {
162            return None;
163        }
164
165        let browser_bin = Self::find_browser_binary()?;
166
167        let temp_png = Self::generate_temp_png_path();
168        let mut temp_html = temp_png.clone();
169        temp_html.set_extension("html");
170
171        let mut injected_html = html_str.to_string();
172        let base_tag = format!("<base href=\"{}\">", url);
173        if let Some(idx) = injected_html.find("<head>") {
174            injected_html.insert_str(idx + 6, &base_tag);
175        } else if let Some(idx) = injected_html.find("<head ") {
176            if let Some(close_idx) = injected_html[idx..].find('>') {
177                injected_html.insert_str(idx + close_idx + 1, &base_tag);
178            }
179        } else {
180            injected_html.insert_str(0, &base_tag);
181        }
182
183        let _ = std::fs::write(&temp_html, injected_html);
184
185        let _cleanup_png = TempFileCleanup(temp_png.clone());
186        let _cleanup_html = TempFileCleanup(temp_html.clone());
187
188        let temp_png_str = temp_png.to_string_lossy().to_string();
189        let temp_html_str = format!(
190            "file:///{}",
191            temp_html.to_string_lossy().to_string().replace('\\', "/")
192        );
193
194        let screenshot_arg = format!("--screenshot={}", temp_png_str);
195        let window_size_arg = format!("--window-size={},{}", width, height);
196
197        let target_arg = if url.starts_with("http://") || url.starts_with("https://") {
198            url.to_string()
199        } else {
200            temp_html_str
201        };
202
203        let profile_dir = std::env::temp_dir().join("headless_engine_stealth_profile");
204        let profile_arg = format!("--user-data-dir={}", profile_dir.to_string_lossy());
205
206        let mut cmd = tokio::process::Command::new(&browser_bin);
207        cmd.arg("--headless=new")
208            .arg("--disable-gpu")
209            .arg("--no-sandbox")
210            .arg(&profile_arg)
211            .arg("--hide-scrollbars")
212            .arg("--disable-blink-features=AutomationControlled")
213            .arg("--user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36")
214            .arg("--lang=en-US,en")
215            .arg("--virtual-time-budget=8000")
216            .arg(&window_size_arg)
217            .arg(&screenshot_arg)
218            .arg(&target_arg);
219
220        #[cfg(target_os = "windows")]
221        cmd.creation_flags(0x08000000); // CREATE_NO_WINDOW
222
223        if let Ok(child) = cmd.spawn() {
224            let _ =
225                tokio::time::timeout(std::time::Duration::from_secs(12), child.wait_with_output())
226                    .await;
227        }
228
229        if temp_png.exists() {
230            if let Ok(png_bytes) = std::fs::read(&temp_png) {
231                let png_base64 = format!(
232                    "data:image/png;base64,{}",
233                    base64::engine::general_purpose::STANDARD.encode(&png_bytes)
234                );
235                return Some(ScreenshotResult {
236                    width,
237                    height,
238                    svg: String::new(),
239                    layout_wireframe: String::new(),
240                    element_count: 1,
241                    png_bytes,
242                    png_base64,
243                });
244            }
245        }
246        None
247    }
248
249    pub fn capture_real_screenshot_sync(
250        url: &str,
251        html_str: &str,
252        width: u32,
253        height: u32,
254    ) -> Option<ScreenshotResult> {
255        if !Self::is_valid_url(url) {
256            return None;
257        }
258
259        let browser_bin = Self::find_browser_binary()?;
260
261        let temp_png = Self::generate_temp_png_path();
262        let mut temp_html = temp_png.clone();
263        temp_html.set_extension("html");
264
265        let mut injected_html = html_str.to_string();
266        let base_tag = format!("<base href=\"{}\">", url);
267        if let Some(idx) = injected_html.find("<head>") {
268            injected_html.insert_str(idx + 6, &base_tag);
269        } else if let Some(idx) = injected_html.find("<head ") {
270            if let Some(close_idx) = injected_html[idx..].find('>') {
271                injected_html.insert_str(idx + close_idx + 1, &base_tag);
272            }
273        } else {
274            injected_html.insert_str(0, &base_tag);
275        }
276
277        let _ = std::fs::write(&temp_html, injected_html);
278
279        let _cleanup_png = TempFileCleanup(temp_png.clone());
280        let _cleanup_html = TempFileCleanup(temp_html.clone());
281
282        let temp_png_str = temp_png.to_string_lossy().to_string();
283        let temp_html_str = format!(
284            "file:///{}",
285            temp_html.to_string_lossy().to_string().replace('\\', "/")
286        );
287
288        let screenshot_arg = format!("--screenshot={}", temp_png_str);
289        let window_size_arg = format!("--window-size={},{}", width, height);
290
291        let target_arg = if url.starts_with("http://") || url.starts_with("https://") {
292            url.to_string()
293        } else {
294            temp_html_str
295        };
296
297        let profile_dir = std::env::temp_dir().join("headless_engine_stealth_profile");
298        let profile_arg = format!("--user-data-dir={}", profile_dir.to_string_lossy());
299
300        let mut cmd = std::process::Command::new(&browser_bin);
301        cmd.arg("--headless=new")
302            .arg("--disable-gpu")
303            .arg("--no-sandbox")
304            .arg(&profile_arg)
305            .arg("--hide-scrollbars")
306            .arg("--disable-blink-features=AutomationControlled")
307            .arg("--user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36")
308            .arg("--lang=en-US,en")
309            .arg("--virtual-time-budget=8000")
310            .arg(&window_size_arg)
311            .arg(&screenshot_arg)
312            .arg(&target_arg);
313
314        #[cfg(target_os = "windows")]
315        {
316            use std::os::windows::process::CommandExt;
317            cmd.creation_flags(0x08000000); // CREATE_NO_WINDOW
318        }
319
320        if let Ok(mut child) = cmd.spawn() {
321            let start = std::time::Instant::now();
322            let timeout = std::time::Duration::from_secs(12);
323            loop {
324                if let Ok(Some(_)) = child.try_wait() {
325                    break;
326                }
327                if start.elapsed() > timeout {
328                    let _ = child.kill();
329                    let _ = child.wait();
330                    break;
331                }
332                std::thread::sleep(std::time::Duration::from_millis(50));
333            }
334        }
335
336        if temp_png.exists() {
337            if let Ok(png_bytes) = std::fs::read(&temp_png) {
338                let png_base64 = format!(
339                    "data:image/png;base64,{}",
340                    base64::engine::general_purpose::STANDARD.encode(&png_bytes)
341                );
342                return Some(ScreenshotResult {
343                    width,
344                    height,
345                    svg: String::new(),
346                    layout_wireframe: String::new(),
347                    element_count: 1,
348                    png_bytes,
349                    png_base64,
350                });
351            }
352        }
353        None
354    }
355}
356
357pub struct PageRenderer;
358
359impl PageRenderer {
360    pub async fn render_async(
361        url: &str,
362        title: &str,
363        html_str: &str,
364        interactive: &[InteractiveElement],
365        _search_results: Option<&SearchResults>,
366    ) -> ScreenshotResult {
367        let width = 1280;
368        let height = 720;
369
370        // Try one-shot ultra-light real browser screenshot first
371        if let Some(real_shot) =
372            RealBrowserScreenshot::capture_real_screenshot_async(url, html_str, width, height).await
373        {
374            return real_shot;
375        }
376
377        // Fallback to pure-Rust layout & rasterization engine
378        crate::render::HtmlRenderer::render_html_to_screenshot(
379            url,
380            title,
381            html_str,
382            interactive,
383            width,
384            height,
385        )
386        .await
387        .unwrap_or_else(|_| {
388            Self::render_general_page(url, title, html_str, interactive, width, height)
389        })
390    }
391
392    pub fn render(
393        url: &str,
394        title: &str,
395        html_str: &str,
396        interactive: &[InteractiveElement],
397        _search_results: Option<&SearchResults>,
398    ) -> ScreenshotResult {
399        let width = 1280;
400        let height = 720;
401
402        if let Some(real_shot) =
403            RealBrowserScreenshot::capture_real_screenshot_sync(url, html_str, width, height)
404        {
405            return real_shot;
406        }
407
408        Self::render_general_page(url, title, html_str, interactive, width, height)
409    }
410
411    fn render_general_page(
412        url: &str,
413        title: &str,
414        html_str: &str,
415        interactive: &[InteractiveElement],
416        width: u32,
417        height: u32,
418    ) -> ScreenshotResult {
419        let mut y_offset = 120;
420        let document = Html::parse_document(html_str);
421
422        let mut visual_blocks = Vec::new();
423        if let Ok(sel) = Selector::parse("h1, h2, h3, p, button, a[href], input, li") {
424            for el in document.select(&sel) {
425                let tag = el.value().name();
426                let text = el.text().collect::<Vec<_>>().join(" ").trim().to_string();
427                if text.is_empty() || text.len() < 2 {
428                    continue;
429                }
430
431                let linked_index = interactive.iter().find(|i| i.text == text).map(|i| i.index);
432                visual_blocks.push((tag.to_string(), text, linked_index));
433                if visual_blocks.len() >= 50 {
434                    break;
435                }
436            }
437        }
438
439        let title_escaped = Self::xml_escape(&title.chars().take(30).collect::<String>());
440        let url_escaped = Self::xml_escape(url);
441
442        let mut svg = format!(
443            "<svg xmlns=\"http://www.w3.org/2000/svg\" viewBox=\"0 0 {width} {height}\" width=\"{width}\" height=\"{height}\" style=\"background:#0f172a; font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;\">\n\
444  <rect width=\"{width}\" height=\"80\" fill=\"#1e293b\" />\n\
445  <circle cx=\"25\" cy=\"40\" r=\"7\" fill=\"#ef4444\" />\n\
446  <circle cx=\"45\" cy=\"40\" r=\"7\" fill=\"#f59e0b\" />\n\
447  <circle cx=\"65\" cy=\"40\" r=\"7\" fill=\"#10b981\" />\n\
448  <rect x=\"90\" y=\"24\" width=\"900\" height=\"32\" rx=\"6\" fill=\"#0f172a\" stroke=\"#334155\" stroke-width=\"1\" />\n\
449  <text x=\"105\" y=\"45\" fill=\"#94a3b8\" font-size=\"13\">&#128274; {url_escaped}</text>\n\
450  <text x=\"1010\" y=\"45\" fill=\"#e2e8f0\" font-size=\"13\" font-weight=\"600\">{title_escaped}</text>\n\
451  <g transform=\"translate(60, 100)\">\n"
452        );
453
454        let mut wireframe_lines = Vec::new();
455        wireframe_lines.push(
456            "╔══════════════════════════════════════════════════════════════════════════════╗"
457                .to_string(),
458        );
459        wireframe_lines.push(format!("║ URL: {:<72} ║", url));
460        wireframe_lines.push(format!("║ TITLE: {:<70} ║", title));
461        wireframe_lines.push(
462            "╠══════════════════════════════════════════════════════════════════════════════╣"
463                .to_string(),
464        );
465
466        for (tag, text, index_opt) in &visual_blocks {
467            let truncated_text: String = text.chars().take(80).collect();
468            let escaped_text = Self::xml_escape(&truncated_text);
469
470            match tag.as_str() {
471                "h1" => {
472                    svg.push_str(&format!(
473                        "    <text x=\"0\" y=\"{}\" fill=\"#f8fafc\" font-size=\"24\" font-weight=\"bold\">{}</text>\n",
474                        y_offset, escaped_text
475                    ));
476                    wireframe_lines.push(format!("║ # {:<74} ║", truncated_text));
477                    y_offset += 40;
478                }
479                "h2" => {
480                    svg.push_str(&format!(
481                        "    <text x=\"0\" y=\"{}\" fill=\"#38bdf8\" font-size=\"18\" font-weight=\"bold\">{}</text>\n",
482                        y_offset, escaped_text
483                    ));
484                    wireframe_lines.push(format!("║ ## {:<73} ║", truncated_text));
485                    y_offset += 32;
486                }
487                _ => {
488                    let badge = index_opt.map(|i| format!("[{}] ", i)).unwrap_or_default();
489                    let color = if index_opt.is_some() {
490                        "#60a5fa"
491                    } else {
492                        "#cbd5e1"
493                    };
494                    svg.push_str(&format!(
495                        "    <text x=\"0\" y=\"{}\" fill=\"{}\" font-size=\"13\">{}{}</text>\n",
496                        y_offset, color, badge, escaped_text
497                    ));
498                    if index_opt.is_some() {
499                        wireframe_lines.push(format!("║ [LINK] {}{:<67} ║", badge, truncated_text));
500                    } else {
501                        wireframe_lines.push(format!("║ {:<76} ║", truncated_text));
502                    }
503                    y_offset += 22;
504                }
505            }
506
507            if y_offset > 1100 {
508                break;
509            }
510        }
511
512        svg.push_str("  </g>\n</svg>");
513        wireframe_lines.push(
514            "╚══════════════════════════════════════════════════════════════════════════════╝"
515                .to_string(),
516        );
517
518        let png_bytes = Self::render_png(&svg, width, height).unwrap_or_default();
519        let png_base64 = if !png_bytes.is_empty() {
520            format!(
521                "data:image/png;base64,{}",
522                base64::engine::general_purpose::STANDARD.encode(&png_bytes)
523            )
524        } else {
525            String::new()
526        };
527
528        ScreenshotResult {
529            width,
530            height,
531            svg,
532            layout_wireframe: wireframe_lines.join("\n"),
533            element_count: visual_blocks.len(),
534            png_bytes,
535            png_base64,
536        }
537    }
538
539    pub fn render_png(svg_str: &str, width: u32, height: u32) -> Result<Vec<u8>> {
540        let opt = resvg::usvg::Options {
541            font_family: "sans-serif".to_string(),
542            ..Default::default()
543        };
544        let tree = resvg::usvg::Tree::from_str(svg_str, &opt)?;
545        let mut pixmap = resvg::tiny_skia::Pixmap::new(width, height)
546            .ok_or_else(|| anyhow::anyhow!("Failed to allocate raster pixmap"))?;
547
548        resvg::render(
549            &tree,
550            resvg::tiny_skia::Transform::default(),
551            &mut pixmap.as_mut(),
552        );
553        let png_data = pixmap.encode_png()?;
554        Ok(png_data)
555    }
556
557    fn xml_escape(s: &str) -> String {
558        s.replace('&', "&amp;")
559            .replace('<', "&lt;")
560            .replace('>', "&gt;")
561            .replace('"', "&quot;")
562            .replace('\'', "&apos;")
563    }
564}
headless_engine/dom/screenshot.rs

headless_engine/dom/
screenshot.rs
