#!/usr/bin/env python3 import sys import os import argparse from selenium import webdriver from selenium.webdriver.common.by import By import xml.etree.ElementTree as ET class ScreenshotGenerator: def __init__(self, debug_port: int, debug: bool = False): """Initialize the screenshot generator with optional debug mode""" self.debug = debug options = webdriver.ChromeOptions() options.add_experimental_option("debuggerAddress", f"127.0.0.1:{debug_port}") self.driver = webdriver.Chrome(options=options) # Load JavaScript files js_dir = os.path.join(os.path.dirname(__file__), "..", "js") self.js_code = "" for filename in ["viewport.js", "mouse.js", "dom_analyzer.js"]: with open(os.path.join(js_dir, filename)) as f: self.js_code += f.read() + "\n" def log(self, msg: str): """Log message if verbose mode is enabled""" if self.debug: print(f"[Debug] {msg}") def get_visible_dom(self) -> str: """Get only the visible DOM elements with enhanced visibility detection""" if self.debug: self.log("Starting DOM analysis") print("Debug mode enabled, gathering DOM information...") try: script = """ // Debug helper const debug = { log: [], add: function(msg) { this.log.push(msg); } }; function isVisible(element) { debug.add(`Checking visibility for ${element.tagName}`); if (!element || element.nodeType !== Node.ELEMENT_NODE) { debug.add('Not an element node'); return false; } const style = window.getComputedStyle(element); const rect = element.getBoundingClientRect(); debug.add(`Element: ${element.tagName}, display: ${style.display}, visibility: ${style.visibility}`); debug.add(`Dimensions: ${rect.width}x${rect.height}`); return ( style.display !== 'none' && style.visibility !== 'hidden' && rect.width > 0 && rect.height > 0 ); } function getElementInfo(element) { const info = { tag: element.tagName.toLowerCase(), attributes: {}, text: element.textContent.trim() }; ['id', 'class', 'href', 'src'].forEach(attr => { if (element.hasAttribute(attr)) { info.attributes[attr] = element.getAttribute(attr); } }); return info; } function traverseDOM(element, depth = 0) { if (!element) { debug.add('Null element encountered'); return ''; } let output = ''; if (element.nodeType === Node.ELEMENT_NODE && isVisible(element)) { const info = getElementInfo(element); const indent = ' '.repeat(depth); let elementStr = `${indent}${info.tag}`; Object.entries(info.attributes).forEach(([key, value]) => { elementStr += ` ${key}="${value}"`; }); if (info.text && !element.children.length) { elementStr += ` [text: ${info.text}]`; } output += elementStr + '\\n'; Array.from(element.children).forEach(child => { output += traverseDOM(child, depth + 1); }); } return output; } const dom = traverseDOM(document.documentElement); debug.add(`Total output length: ${dom.length}`); return { result: dom, debug: debug.log, stats: { url: window.location.href, title: document.title, elementCount: document.getElementsByTagName('*').length, bodyChildCount: document.body ? document.body.children.length : 0 } }; """ result = self.driver.execute_script(script) if isinstance(result, dict): if self.debug: print("\n=== Page Statistics ===") stats = result.get('stats', {}) print(f"Title: {stats.get('title')}") print(f"Element count: {stats.get('elementCount')}") print(f"Body children: {stats.get('bodyChildCount')}") print("\n=== Debug Log ===") for log_entry in result.get('debug', []): print(f"Browser: {log_entry}") return result.get('result', '') return result except Exception as e: if self.debug: print(f"Error in get_visible_dom: {str(e)}") return "" def get_full_dom(self) -> str: """Get the full DOM structure""" return self.driver.execute_script(""" function traverseDOM(node, indent = '') { if (!node) return ''; let result = ''; // Add node information if (node.nodeType === Node.ELEMENT_NODE) { result += indent + node.tagName.toLowerCase(); if (node.id) result += ` id="${node.id}"`; if (node.className) result += ` class="${node.className}"`; result += '\\n'; // Process children for (let child of node.children) { result += traverseDOM(child, indent + ' '); } } return result; } return traverseDOM(document.documentElement); """) def generate_xml(self) -> str: page_data = self.driver.execute_script(self.js_code + "; return analyzePageContent();") # Create viewport element with attributes viewport = ET.Element("viewport") for key, value in page_data['viewport'].items(): viewport.set(key, str(value)) def process_element(element_data, parent): # Add fixed position indication if present if element_data.get('fixed'): if element_data.get('tag') == 'nav': nav = ET.SubElement(parent, "navigation") nav.set("fixed", "true") nav.set("x", str(element_data['x'])) nav.set("y", str(element_data['y'])) nav.set("width", str(element_data['width'])) nav.set("height", str(element_data['height'])) current_parent = nav else: fixed_container = ET.SubElement(parent, "fixed") fixed_container.set("x", str(element_data['x'])) fixed_container.set("y", str(element_data['y'])) fixed_container.set("width", str(element_data['width'])) fixed_container.set("height", str(element_data['height'])) current_parent = fixed_container else: current_parent = parent # Handle text content if element_data.get('text'): if len(current_parent) == 0: current_parent.text = (current_parent.text or '') + element_data['text'] + '\n' else: current_parent[-1].tail = (current_parent[-1].tail or '') + element_data['text'] + '\n' # Handle interactive elements if element_data.get('interactions'): interactive = ET.SubElement(current_parent, "interactive") interactive.set("x", str(element_data['x'])) interactive.set("y", str(element_data['y'])) interactive.set("width", str(element_data['width'])) interactive.set("height", str(element_data['height'])) interactive.set("interactions", ','.join(element_data['interactions'])) if element_data.get('id'): interactive.set("id", element_data['id']) elif element_data.get('suggestedClick'): interactive.set("suggested_click", f"{element_data['suggestedClick']['x']},{element_data['suggestedClick']['y']}") if element_data.get('placeholder'): interactive.set("placeholder", element_data['placeholder']) if element_data.get('buttonText'): interactive.text = element_data['buttonText'] # Handle images if element_data.get('tag') == 'img': img = ET.SubElement(current_parent, "img") img.set("x", str(element_data['x'])) img.set("y", str(element_data['y'])) img.set("width", str(element_data['width'])) img.set("height", str(element_data['height'])) img.set("src", element_data.get('src', '')) img.set("alt", element_data.get('alt', '')) if element_data.get('id'): img.set("id", element_data['id']) # Process children for child in element_data.get('children', []): process_element(child, current_parent) # Process all elements process_element(page_data['content'], viewport) return self.format_xml(ET.tostring(viewport, encoding='unicode')) def format_xml(self, xml_str: str) -> str: """Format XML with proper indentation""" root = ET.fromstring(xml_str) ET.indent(root, space=" ", level=0) return ET.tostring(root, encoding='unicode', method='xml') def main(): parser = argparse.ArgumentParser(description='Generate page screenshot in XML format') parser.add_argument('debug_port', type=int, help='Chrome debug port') parser.add_argument('--debug', action='store_true', help='Enable debug output') parser.add_argument('--dom', action='store_true', help='Print full DOM structure') parser.add_argument('--visible-dom', action='store_true', help='Print visible DOM structure') args = parser.parse_args() try: generator = ScreenshotGenerator(args.debug_port, args.debug) # Basic DOM check when debug is enabled if args.debug: print("\nPerforming basic DOM check...") basic_check = generator.driver.execute_script(""" return { title: document.title, bodyExists: !!document.body, elementCount: document.getElementsByTagName('*').length, firstElement: document.body ? document.body.firstElementChild.tagName : 'none' } """) print(f"Page title: {basic_check['title']}") print(f"Body exists: {basic_check['bodyExists']}") print(f"Total elements: {basic_check['elementCount']}") print(f"First element: {basic_check['firstElement']}") if args.dom: print("=== Full DOM Structure ===") try: dom_output = generator.get_full_dom() if dom_output.strip(): print(dom_output) else: print("No DOM content found.") except Exception as e: print(f"Error retrieving DOM structure: {str(e)}") elif args.visible_dom: print("=== Visible DOM Structure ===") try: visible_dom = generator.get_visible_dom() if visible_dom.strip(): print(visible_dom) else: print("No visible DOM content found.") except Exception as e: print(f"Error retrieving visible DOM structure: {str(e)}") else: xml_output = generator.generate_xml() print(xml_output) except Exception as e: print(f"Error: {e}", file=sys.stderr) sys.exit(1) if __name__ == "__main__": main()