87 lines
3.2 KiB
Python
Executable File
87 lines
3.2 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import sys
|
|
import os
|
|
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):
|
|
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 get_page_content(self) -> dict:
|
|
return self.driver.execute_script(self.js_code + "; return analyzePageContent();")
|
|
|
|
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 generate_xml(self) -> str:
|
|
page_data = self.get_page_content()
|
|
|
|
viewport = ET.Element("viewport")
|
|
for k, v in page_data['viewport'].items():
|
|
viewport.set(k, str(v))
|
|
|
|
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'
|
|
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:')
|
|
))
|
|
|
|
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)
|
|
|
|
process_node(page_data['content'], viewport)
|
|
return self.format_xml(ET.tostring(viewport, encoding='unicode'))
|
|
|
|
def main(debug_port: int):
|
|
try:
|
|
generator = ScreenshotGenerator(debug_port)
|
|
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]))
|