llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update to itb_screenshot and friends

Browse Source
This commit is contained in:
geens
2024-12-27 17:07:19 +01:00
parent fe72969b4f
commit e0bc7a7159
15 changed files with 6739 additions and 552 deletions
Download Patch File Download Diff File Expand all files Collapse all files

328
tools/itb/bin/itb_screenshot
View File

@@ -1,12 +1,15 @@
#!/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):
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)
@@ -18,69 +21,294 @@ class ScreenshotGenerator:
with open(os.path.join(js_dir, filename)) as f:
self.js_code += f.read() + "\n"
def get_page_content(self) -> dict:
return self.driver.execute_script(self.js_code + "; return analyzePageContent();")
def log(self, msg: str):
"""Log message if verbose mode is enabled"""
if self.debug:
print(f"[Debug] {msg}")
def format_xml(self, xml_str: str) -> str:
root = ET.fromstring(xml_str)
ET.indent(ET.ElementTree(root), space=" ", level=0)
return ET.tostring(root, encoding='unicode', method='xml')
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.get_page_content()
page_data = self.driver.execute_script(self.js_code + "; return analyzePageContent();")
# Create viewport element with attributes
viewport = ET.Element("viewport")
for k, v in page_data['viewport'].items():
viewport.set(k, str(v))
for key, value in page_data['viewport'].items():
viewport.set(key, str(value))
def process_node(node, parent):
if node is None:
return
if node['text']:
if len(parent) == 0:
parent.text = (parent.text or '') + node['text'] + '\n'
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:
parent[-1].tail = (parent[-1].tail or '') + node['text'] + '\n'
if node['interactions']:
interactive = ET.SubElement(parent, "interactive")
interactive.set("x", str(node['location']['x']))
interactive.set("y", str(node['location']['y']))
interactive.set("width", str(node['size']['width']))
interactive.set("height", str(node['size']['height']))
interactive.set("interactions", ",".join(
i for i in node['interactions'] if not i.startswith('placeholder:')
))
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']))
placeholder = next(
(i.split(':', 1)[1] for i in node['interactions']
if i.startswith('placeholder:')),
None
)
if placeholder:
interactive.set('placeholder', placeholder)
if node['text']:
interactive.text = node['text']
for child in node['children']:
process_node(child, parent)
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)
process_node(page_data['content'], viewport)
return self.format_xml(ET.tostring(viewport, encoding='unicode'))
def main(debug_port: int):
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(debug_port)
xml_output = generator.generate_xml()
print(xml_output)
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__":
if len(sys.argv) != 2:
print("Usage: itb_screenshot <debug_port>")
sys.exit(1)
main(int(sys.argv[1]))
main()