315 lines
13 KiB
Python
Executable File
315 lines
13 KiB
Python
Executable File
#!/root/venvs/itb/bin/python
|
|
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()
|