Update to itb_screenshot and friends

This commit is contained in:
2024-12-27 17:07:19 +01:00
parent fe72969b4f
commit e0bc7a7159
15 changed files with 6739 additions and 552 deletions

View File

@@ -2,22 +2,17 @@
import sys
import os
import time
import random
import argparse
import json
from selenium import webdriver
from selenium.webdriver import ActionChains
from selenium.webdriver.common.actions.action_builder import ActionBuilder
from selenium.webdriver.common.actions.mouse_button import MouseButton
from selenium.webdriver.common.by import By
def debug(msg: str):
"""Print debug message to stderr"""
print(f"DEBUG: {msg}", file=sys.stderr)
class ClickExecutor:
def __init__(self, debug_port: int):
debug(f"Connecting to Chrome on debug port {debug_port}")
def __init__(self, debug_port: int, verbose: bool = False):
"""Initialize connection to Chrome instance"""
self.verbose = verbose
options = webdriver.ChromeOptions()
options.add_experimental_option("debuggerAddress", f"127.0.0.1:{debug_port}")
self.driver = webdriver.Chrome(options=options)
@@ -29,15 +24,15 @@ class ClickExecutor:
with open(os.path.join(js_dir, filename)) as f:
self.js_code += f.read() + "\n"
# Get initial window info
# Get window and viewport info
self.window_size = self.driver.get_window_size()
debug(f"Window size: {json.dumps(self.window_size)}")
self.viewport_size = self.get_viewport_size()
debug(f"Viewport size: {json.dumps(self.viewport_size)}")
self.scroll_position = self.get_scroll_position()
debug(f"Scroll position: {json.dumps(self.scroll_position)}")
def log(self, msg: str):
"""Log message if verbose mode is enabled"""
if self.verbose:
print(f"DEBUG: {msg}", file=sys.stderr)
def get_viewport_size(self):
"""Get viewport dimensions"""
@@ -58,37 +53,47 @@ class ClickExecutor:
max(0, x - self.viewport_size['width'] // 2),
max(0, y - self.viewport_size['height'] // 2)
)
# Wait for scroll to complete
time.sleep(0.5)
self.scroll_position = self.get_scroll_position()
debug(f"New scroll position: {json.dumps(self.scroll_position)}")
def get_element_at_position(self, x: int, y: int):
"""Get the element at the specified coordinates"""
return self.driver.execute_script(self.js_code + "; return getElementAtPosition(arguments[0], arguments[1]);", x, y)
return self.driver.execute_script(
self.js_code + "; return getElementAtPosition(arguments[0], arguments[1]);",
x, y
)
def execute_click(self, x: int, y: int, double: bool = False, right: bool = False):
"""Execute click operation with improved coordinate handling"""
def execute_click(self, x: int, y: int, double: bool = False, right: bool = False, element_id: str = None):
"""Execute click operation with improved handling"""
try:
debug(f"Attempting to click at coordinates ({x}, {y})")
debug(f"Click type: {'double' if double else 'right' if right else 'single'}")
self.log(f"Attempting {'double' if double else 'right' if right else 'single'} click at ({x}, {y})")
# First scroll to make target visible
self.scroll_to_target(x, y)
# Get the element at the target position
element = self.get_element_at_position(x, y)
if not element:
debug(f"No element found at coordinates ({x}, {y})")
return False
if element_id:
# Try to find element by ID first
try:
element = self.driver.find_element(By.ID, element_id)
self.log(f"Found element by ID: {element_id}")
except:
self.log(f"Failed to find element by ID: {element_id}")
return False
else:
# Scroll to make target visible
self.scroll_to_target(x, y)
# Try clicking the element directly
# Get element at position
element = self.get_element_at_position(x, y)
if not element:
self.log(f"No element found at coordinates ({x}, {y})")
return False
# Try clicking the element
try:
debug("Attempting to click element directly")
element = self.driver.execute_script("return arguments[0];", element)
actions = ActionChains(self.driver)
actions.move_to_element(element)
if element_id:
actions.move_to_element(element)
else:
actions.move_by_offset(x, y)
if right:
actions.context_click()
@@ -101,28 +106,29 @@ class ClickExecutor:
return True
except Exception as e:
debug(f"Direct element click failed: {str(e)}")
debug("Falling back to JavaScript click")
self.log(f"Direct click failed: {str(e)}")
self.log("Attempting JavaScript click fallback")
# Fallback to JavaScript click
try:
self.driver.execute_script("arguments[0].click();", element)
return True
except Exception as js_error:
debug(f"JavaScript click failed: {str(js_error)}")
self.log(f"JavaScript click failed: {str(js_error)}")
return False
except Exception as e:
debug(f"Click operation failed with error: {str(e)}")
self.log(f"Click operation failed: {str(e)}")
return False
def main():
parser = argparse.ArgumentParser(description='Execute mouse clicks in Chrome')
parser.add_argument('debug_port', type=int, help='Chrome debug port')
parser.add_argument('-x', type=int, required=True, help='X coordinate')
parser.add_argument('-y', type=int, required=True, help='Y coordinate')
parser.add_argument('-x', type=int, help='X coordinate')
parser.add_argument('-y', type=int, help='Y coordinate')
parser.add_argument('--double', action='store_true', help='Perform double click')
parser.add_argument('--right', action='store_true', help='Perform right click')
parser.add_argument('--id', help='Element ID to click')
parser.add_argument('-v', '--verbose', action='store_true', help='Enable debug output')
args = parser.parse_args()
@@ -130,15 +136,25 @@ def main():
print("Cannot specify both --double and --right", file=sys.stderr)
sys.exit(1)
if not args.id and (args.x is None or args.y is None):
print("Must specify either element ID or both x and y coordinates", file=sys.stderr)
sys.exit(1)
try:
debug("Initializing ClickExecutor")
executor = ClickExecutor(args.debug_port)
success = executor.execute_click(args.x, args.y, args.double, args.right)
executor = ClickExecutor(args.debug_port, args.verbose)
success = executor.execute_click(
x=args.x if args.x is not None else 0,
y=args.y if args.y is not None else 0,
double=args.double,
right=args.right,
element_id=args.id
)
sys.exit(0 if success else 1)
except Exception as e:
debug(f"Fatal error: {str(e)}")
if args.verbose:
print(f"Fatal error: {str(e)}", file=sys.stderr)
sys.exit(1)
if __name__ == "__main__":
main()
main()

View File

@@ -0,0 +1,168 @@
#!/usr/bin/env python3
import sys
import os
import time
import argparse
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.action_chains import ActionChains
class InputController:
def __init__(self, debug_port: int, verbose: bool = False):
"""Initialize connection to Chrome instance"""
self.verbose = verbose
options = webdriver.ChromeOptions()
options.add_experimental_option("debuggerAddress", f"127.0.0.1:{debug_port}")
self.driver = webdriver.Chrome(options=options)
def log(self, msg: str):
"""Log message if verbose mode is enabled"""
if self.verbose:
print(f"DEBUG: {msg}", file=sys.stderr)
def focus_element(self, element_id: str = None):
"""Focus on an input element"""
try:
if element_id:
element = self.driver.find_element(By.ID, element_id)
else:
# Get currently focused element
element = self.driver.switch_to.active_element
element.click()
return element
except Exception as e:
self.log(f"Failed to focus element: {str(e)}")
return None
def enter_text(self, text: str, element_id: str = None, append: bool = False):
"""Enter text into an input element"""
try:
element = self.focus_element(element_id)
if not element:
return False
if not append:
# Clear existing text first
if element.tag_name.lower() in ['input', 'textarea']:
element.clear()
else:
# For contenteditable elements
element.send_keys(Keys.CONTROL + 'a')
element.send_keys(Keys.DELETE)
# Type the text
element.send_keys(text)
return True
except Exception as e:
self.log(f"Failed to enter text: {str(e)}")
return False
def set_caret_position(self, position: int, element_id: str = None):
"""Set the text cursor position"""
try:
element = self.focus_element(element_id)
if not element:
return False
# Use JavaScript to set caret position
self.driver.execute_script(
"""
const el = arguments[0];
const pos = arguments[1];
if (el.setSelectionRange) {
el.setSelectionRange(pos, pos);
} else if (el.createTextRange) {
const range = el.createTextRange();
range.collapse(true);
range.moveEnd('character', pos);
range.moveStart('character', pos);
range.select();
}
""",
element, position
)
return True
except Exception as e:
self.log(f"Failed to set caret position: {str(e)}")
return False
def select_text(self, start: int, end: int, element_id: str = None):
"""Select text range in an input element"""
try:
element = self.focus_element(element_id)
if not element:
return False
# Use JavaScript to set selection
self.driver.execute_script(
"""
const el = arguments[0];
const start = arguments[1];
const end = arguments[2];
if (el.setSelectionRange) {
el.setSelectionRange(start, end);
} else if (el.createTextRange) {
const range = el.createTextRange();
range.collapse(true);
range.moveStart('character', start);
range.moveEnd('character', end);
range.select();
}
""",
element, start, end
)
return True
except Exception as e:
self.log(f"Failed to select text: {str(e)}")
return False
def main():
parser = argparse.ArgumentParser(description='Control text input in Chrome')
parser.add_argument('debug_port', type=int, help='Chrome debug port')
parser.add_argument('-i', '--id', help='Element ID')
parser.add_argument('-t', '--text', help='Text to input')
parser.add_argument('--append', action='store_true', help='Append text instead of replacing')
parser.add_argument('--caret', type=int, help='Set caret position')
parser.add_argument('--select-start', type=int, help='Text selection start')
parser.add_argument('--select-end', type=int, help='Text selection end')
parser.add_argument('-v', '--verbose', action='store_true', help='Enable debug output')
args = parser.parse_args()
# Validate args
if args.select_start is not None and args.select_end is None:
print("Must specify both --select-start and --select-end", file=sys.stderr)
sys.exit(1)
if args.select_end is not None and args.select_start is None:
print("Must specify both --select-start and --select-end", file=sys.stderr)
sys.exit(1)
try:
controller = InputController(args.debug_port, args.verbose)
success = False
if args.text is not None:
success = controller.enter_text(args.text, args.id, args.append)
elif args.caret is not None:
success = controller.set_caret_position(args.caret, args.id)
elif args.select_start is not None:
success = controller.select_text(args.select_start, args.select_end, args.id)
else:
success = controller.focus_element(args.id) is not None
sys.exit(0 if success else 1)
except Exception as e:
if args.verbose:
print(f"Fatal error: {str(e)}", file=sys.stderr)
sys.exit(1)
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,71 @@
#!/usr/bin/env python3
import sys
import os
import time
import argparse
from selenium import webdriver
class NavigationController:
def __init__(self, debug_port: int):
"""Initialize connection to Chrome instance"""
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 = ""
with open(os.path.join(js_dir, "viewport.js")) as f:
self.js_code = f.read()
def navigate(self, url: str) -> bool:
"""
Navigate to the specified URL
Returns True if navigation was successful
"""
try:
# Store initial URL for comparison
initial_url = self.driver.current_url
# Execute navigation
self.driver.get(url)
# Wait for navigation to complete
max_wait = 10 # seconds
start_time = time.time()
while time.time() - start_time < max_wait:
# Check if URL has changed
if self.driver.current_url != initial_url:
# Wait a bit more for page to stabilize
time.sleep(0.5)
return True
time.sleep(0.1)
# If we get here, navigation may have failed
print(f"Warning: Navigation timeout after {max_wait} seconds", file=sys.stderr)
return False
except Exception as e:
print(f"Error during navigation: {str(e)}", file=sys.stderr)
return False
def main():
parser = argparse.ArgumentParser(description='Navigate to URL in Chrome')
parser.add_argument('debug_port', type=int, help='Chrome debug port')
parser.add_argument('url', help='URL to navigate to')
args = parser.parse_args()
try:
controller = NavigationController(args.debug_port)
success = controller.navigate(args.url)
sys.exit(0 if success else 1)
except Exception as e:
print(f"Fatal error: {str(e)}", file=sys.stderr)
sys.exit(1)
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,162 @@
#!/usr/bin/env python3
import sys
import os
import time
import argparse
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
class RefreshController:
def __init__(self, debug_port: int, verbose: bool = False):
"""Initialize connection to Chrome instance"""
self.verbose = verbose
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 = ""
with open(os.path.join(js_dir, "viewport.js")) as f:
self.js_code = f.read()
def log(self, msg: str):
"""Log message if verbose mode is enabled"""
if self.verbose:
print(f"DEBUG: {msg}", file=sys.stderr)
def get_page_state(self) -> dict:
"""Get current page state for comparison"""
return {
'url': self.driver.current_url,
'title': self.driver.title,
'scroll_position': self.driver.execute_script(
self.js_code + "; return getViewportInfo();"
)
}
def wait_for_page_change(self, original_state: dict, timeout: int = 30) -> bool:
"""Wait for page to refresh and stabilize"""
try:
start_time = time.time()
while time.time() - start_time < timeout:
current_state = self.get_page_state()
# Check if page has reloaded
if (current_state['url'] == original_state['url'] and
current_state['title'] == original_state['title']):
# Check if DOM has finished updating
is_stable = self.driver.execute_script("""
return (typeof window._lastMutation === 'undefined' ||
Date.now() - window._lastMutation > 1000);
""")
if is_stable:
return True
time.sleep(0.1)
self.log("Timeout waiting for page to stabilize")
return False
except Exception as e:
self.log(f"Error waiting for page change: {str(e)}")
return False
def setup_mutation_tracking(self):
"""Setup mutation observer to track DOM changes"""
self.driver.execute_script("""
// Track last mutation time
window._lastMutation = Date.now();
// Create mutation observer
const observer = new MutationObserver(() => {
window._lastMutation = Date.now();
});
// Observe entire document for any changes
observer.observe(document, {
childList: true,
subtree: true,
attributes: true,
characterData: true
});
""")
def refresh(self, hard: bool = False, timeout: int = 30) -> bool:
"""
Refresh the current page
Args:
hard (bool): If True, bypass cache with Ctrl+F5
timeout (int): Maximum time to wait for refresh in seconds
Returns:
bool: True if refresh was successful
"""
try:
# Get initial page state
original_state = self.get_page_state()
self.log(f"Original page state: {original_state}")
# Setup mutation tracking
self.setup_mutation_tracking()
if hard:
# Force reload bypassing cache
self.log("Executing hard refresh (Ctrl+F5)")
self.driver.execute_script("""
window.location.reload(true);
""")
else:
# Normal reload
self.log("Executing normal refresh")
self.driver.refresh()
# Wait for page to stabilize
success = self.wait_for_page_change(original_state, timeout)
if success:
# Restore scroll position
self.driver.execute_script(
"window.scrollTo(arguments[0], arguments[1]);",
original_state['scroll_position']['scroll_x'],
original_state['scroll_position']['scroll_y']
)
# Final verification
new_state = self.get_page_state()
self.log(f"New page state: {new_state}")
return True
return False
except Exception as e:
self.log(f"Refresh failed: {str(e)}")
return False
def main():
parser = argparse.ArgumentParser(description='Refresh page in Chrome')
parser.add_argument('debug_port', type=int, help='Chrome debug port')
parser.add_argument('--hard', action='store_true', help='Force refresh bypassing cache')
parser.add_argument('--timeout', type=int, default=30, help='Refresh timeout in seconds')
parser.add_argument('-v', '--verbose', action='store_true', help='Enable debug output')
args = parser.parse_args()
try:
controller = RefreshController(args.debug_port, args.verbose)
success = controller.refresh(args.hard, args.timeout)
sys.exit(0 if success else 1)
except Exception as e:
if args.verbose:
print(f"Fatal error: {str(e)}", file=sys.stderr)
sys.exit(1)
if __name__ == "__main__":
main()

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()