|
| 1 | +import sys |
| 2 | +import os |
| 3 | +from textwrap import dedent |
| 4 | +import requests |
| 5 | +import json |
| 6 | +from configobj import ConfigObj |
| 7 | +from pathlib import Path |
| 8 | +import traceback |
| 9 | +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../"))) |
| 10 | + |
| 11 | +print(f"Python Version: {sys.version}") |
| 12 | +print(f"Python Path: {sys.executable}") |
| 13 | +print(f"Current file path: {os.path.abspath(__file__)}") |
| 14 | + |
| 15 | +from rich import print as rich_print |
| 16 | +from rich.text import Text |
| 17 | +from rich.tree import Tree |
| 18 | +from colorama import Fore, init as colorama_init |
| 19 | +colorama_init(autoreset=True) |
| 20 | + |
| 21 | +import ctypes |
| 22 | +import objc |
| 23 | +from Foundation import NSObject |
| 24 | +from Quartz import ( |
| 25 | + CoreGraphics, |
| 26 | + CGEventSourceFlagsState, |
| 27 | + kCGEventSourceStateHIDSystemState, |
| 28 | + kCGEventFlagMaskControl, |
| 29 | + CGWindowListCopyWindowInfo, |
| 30 | + kCGWindowListOptionOnScreenOnly, |
| 31 | + kCGNullWindowID |
| 32 | +) |
| 33 | +from AppKit import NSEvent, NSControlKeyMask |
| 34 | +import time |
| 35 | + |
| 36 | +import xml.etree.ElementTree as ET |
| 37 | + |
| 38 | +# AXUIElement types |
| 39 | +AXUIElementRef = objc.objc_object |
| 40 | + |
| 41 | +# Load the AX API |
| 42 | +ApplicationServices = objc.loadBundle("ApplicationServices", |
| 43 | + globals(), |
| 44 | + bundle_path="/System/Library/Frameworks/ApplicationServices.framework" |
| 45 | +) |
| 46 | +# AX, _ = objc.loadBundleFunctions(ApplicationServices, globals(), [ |
| 47 | +# ("AXUIElementCreateApplication", b"^{__AXUIElement=}(i)") |
| 48 | +# ]) |
| 49 | + |
| 50 | +AX = objc.loadBundleFunctions(ApplicationServices, globals(), [ |
| 51 | + ("AXUIElementCreateSystemWide", b"^{__AXUIElement=}"), |
| 52 | + ("kAXFocusedUIElementAttribute", b"^{__CFString=}"), |
| 53 | + ("AXUIElementCopyAttributeValue", b"i^{__AXUIElement=}^{__CFString=}^@"), |
| 54 | + ("AXUIElementCopyAttributeNames", b"i^{__AXUIElement=}^{__CFArray=}"), |
| 55 | + ("AXUIElementCopyElementAtPosition", b"i^{__AXUIElement=}dd^@"), |
| 56 | + ("AXUIElementCreateApplication", b"^{__AXUIElement=}" + b"i") |
| 57 | +]) |
| 58 | + |
| 59 | +settings_conf_path = str(Path(__file__).parent.parent.parent / "Framework" / "settings.conf") |
| 60 | + |
| 61 | +def get_mouse_position(): |
| 62 | + event = CoreGraphics.CGEventCreate(None) |
| 63 | + loc = CoreGraphics.CGEventGetLocation(event) |
| 64 | + x, y = round(loc.x), round(loc.y) |
| 65 | + return x, y |
| 66 | + |
| 67 | + |
| 68 | +class App: |
| 69 | + def __init__(self, name: str, bundle_id: str, pid: int, window_title: str): |
| 70 | + self.name = name |
| 71 | + self.bundle_id = bundle_id |
| 72 | + self.pid = pid |
| 73 | + self.window_title = window_title |
| 74 | + |
| 75 | + def __str__(self): |
| 76 | + return Fore.GREEN + dedent(f""" |
| 77 | + App( |
| 78 | + name={self.name}, |
| 79 | + bundle_id={self.bundle_id}, |
| 80 | + pid={self.pid}, |
| 81 | + window_title={self.window_title}, |
| 82 | + )""") |
| 83 | + |
| 84 | +class Inspector: |
| 85 | + def __init__(self): |
| 86 | + self.x: int = -1 |
| 87 | + self.y: int = -1 |
| 88 | + self.app: App = App(name="", bundle_id="", pid=-1, window_title="") |
| 89 | + self.xml_str: str = "" |
| 90 | + self.xml_tree: ET.ElementTree = None |
| 91 | + |
| 92 | + self.server_address: str = "http://127.0.0.1" |
| 93 | + self.server_path: str = "/api/v1/mac/dump/driver" |
| 94 | + self.server_port: int = 18100 |
| 95 | + self.page_src: str = "" |
| 96 | + def wait_for_control_press(self): |
| 97 | + print("Hover over the element and press ⌃ Control key...") |
| 98 | + while True: |
| 99 | + flags = CGEventSourceFlagsState(kCGEventSourceStateHIDSystemState) |
| 100 | + if flags & kCGEventFlagMaskControl: |
| 101 | + point = NSEvent.mouseLocation() |
| 102 | + height = NSScreen.mainScreen().frame().size.height |
| 103 | + x = round(point.x) |
| 104 | + y = round(height - point.y) |
| 105 | + rich_print(f"Captured at x={x}, y={y}") |
| 106 | + self.x, self.y = x, y |
| 107 | + return |
| 108 | + time.sleep(0.1) |
| 109 | + |
| 110 | + def get_frontmost_app(self): |
| 111 | + window_list = CGWindowListCopyWindowInfo(kCGWindowListOptionOnScreenOnly, kCGNullWindowID) |
| 112 | + for window in window_list: |
| 113 | + if window.get("kCGWindowLayer") == 0 and window.get("kCGWindowOwnerName"): |
| 114 | + app_name = window["kCGWindowOwnerName"] |
| 115 | + pid = window["kCGWindowOwnerPID"] |
| 116 | + app = NSRunningApplication.runningApplicationWithProcessIdentifier_(pid) |
| 117 | + bundle_id = app.bundleIdentifier() |
| 118 | + window_title = window.get("kCGWindowName", "") |
| 119 | + self.app = App(name=app_name, bundle_id=bundle_id, pid=pid, window_title=window_title) |
| 120 | + print(self.app) |
| 121 | + break |
| 122 | + |
| 123 | + def get_server_port(self): |
| 124 | + config = ConfigObj(settings_conf_path) |
| 125 | + self.server_port = config["server"]["port"] |
| 126 | + |
| 127 | + def get_dump(self): |
| 128 | + url = f"{self.server_address}:{self.server_port}{self.server_path}" |
| 129 | + try: |
| 130 | + response = requests.get(url).json() |
| 131 | + except requests.exceptions.ConnectionError: |
| 132 | + print(Fore.RED + "Failed to connect to the server. Please launch the Zeuz Node first and launch an app") |
| 133 | + return |
| 134 | + if response["status"] == "ok": |
| 135 | + self.page_src = response["ui_xml"] |
| 136 | + print(Fore.GREEN + f"Successfully got dump from appium driver") |
| 137 | + elif response["status"] == "not_found": |
| 138 | + print(Fore.GREEN + f"You have not launched any app yet. Launch app with the following action:") |
| 139 | + action = [ |
| 140 | + { |
| 141 | + "action_name":f"Launch {self.app.name}", |
| 142 | + "action_disabled":"true", |
| 143 | + "step_actions":[ |
| 144 | + ["macos app bundle id","element parameter",self.app.bundle_id], |
| 145 | + ["launch","appium action","launch"] |
| 146 | + ] |
| 147 | + } |
| 148 | + ] |
| 149 | + print(Fore.CYAN + json.dumps(action, indent=4)) |
| 150 | + self.page_src = "" |
| 151 | + else: |
| 152 | + print(Fore.RED + f"Error: {response['error']}") |
| 153 | + self.page_src = "" |
| 154 | + |
| 155 | + def render_tree(self): |
| 156 | + if not self.page_src: |
| 157 | + return |
| 158 | + |
| 159 | + root = ET.fromstring(self.page_src) |
| 160 | + tree = Tree(f"[bold green]{self.app.name} ({self.app.bundle_id})") |
| 161 | + self.xml_tree = tree |
| 162 | + |
| 163 | + def check_bounding_box(element): |
| 164 | + if element.attrib.get('x') and element.attrib.get('y') and element.attrib.get('width') and element.attrib.get('height'): |
| 165 | + x = int(element.attrib.get('x')) |
| 166 | + y = int(element.attrib.get('y')) |
| 167 | + width = int(element.attrib.get('width')) |
| 168 | + height = int(element.attrib.get('height')) |
| 169 | + if (self.x >= x and |
| 170 | + self.x <= x + width and |
| 171 | + self.y >= y and |
| 172 | + self.y <= y + height |
| 173 | + ): |
| 174 | + return True |
| 175 | + return False |
| 176 | + |
| 177 | + def get_attribute_string(element): |
| 178 | + ignore = ['x', 'y', 'width', 'height'] |
| 179 | + return " ".join([f'{k}="{v}"' for k, v in element.attrib.items() if k not in ignore and v]) |
| 180 | + |
| 181 | + def set_single_zeuz_apiplugin(root): |
| 182 | + elements = root.findall(".//*[@zeuz='aiplugin']") |
| 183 | + if len(elements) > 1: |
| 184 | + element_areas = [] |
| 185 | + for element in elements: |
| 186 | + width = int(element.attrib.get('width', 0)) |
| 187 | + height = int(element.attrib.get('height', 0)) |
| 188 | + area = width * height |
| 189 | + element_areas.append((element, area)) |
| 190 | + |
| 191 | + element_areas.sort(key=lambda x: x[1]) |
| 192 | + for element, _ in element_areas[1:]: |
| 193 | + del element.attrib['zeuz'] |
| 194 | + |
| 195 | + def remove_coordinates(node): |
| 196 | + remove = ['x', 'y', 'width', 'height'] |
| 197 | + for child in node: |
| 198 | + for attrib in list(child.attrib): |
| 199 | + if attrib in remove: |
| 200 | + del child.attrib[attrib] |
| 201 | + remove_coordinates(child) |
| 202 | + |
| 203 | + def build_tree(element, parent_tree): |
| 204 | + element_tag = element.tag |
| 205 | + ignore = ['x', 'y', 'width', 'height'] |
| 206 | + element_attribs = get_attribute_string(element) |
| 207 | + element_coords = f"x={element.attrib.get('x', '')}, y={element.attrib.get('y', '')}, w={element.attrib.get('width', '')}, h={element.attrib.get('height', '')}" |
| 208 | + recorded_coords = f"self.x={self.x}, self.y={self.y}" |
| 209 | + |
| 210 | + if check_bounding_box(element): |
| 211 | + if not any(check_bounding_box(child) for child in element): |
| 212 | + area = int(element.attrib.get('width', '1')) * int(element.attrib.get('height', '1')) |
| 213 | + label = f"[bold blue]{element_tag}: [green]{element_attribs} [dim]({element_coords} Area: {area} {recorded_coords})" |
| 214 | + element.set('zeuz', 'aiplugin') |
| 215 | + else: |
| 216 | + label = f"[bold blue]{element_tag}: [yellow]{element_attribs}" |
| 217 | + |
| 218 | + else: |
| 219 | + label = f"[bold]{element_tag}: {element_attribs}" |
| 220 | + node = parent_tree.add(label) |
| 221 | + |
| 222 | + for child in element: |
| 223 | + if check_bounding_box(child): |
| 224 | + build_tree(child, node) |
| 225 | + else: |
| 226 | + node.add(f"[bold]{child.tag}: {get_attribute_string(child)}") |
| 227 | + |
| 228 | + build_tree(root, tree) |
| 229 | + set_single_zeuz_apiplugin(root) |
| 230 | + rich_print(tree) |
| 231 | + remove_coordinates(root) |
| 232 | + self.xml_str = ET.tostring(root).decode().encode('ascii', 'ignore').decode() |
| 233 | + |
| 234 | + |
| 235 | + ''' Comment out the below code to check if tree contains single zeuz apiplugin ''' |
| 236 | + # tree2 = Tree(f"[bold green]{self.app.name} ({self.app.bundle_id})") |
| 237 | + # build_tree(root, tree2) |
| 238 | + # rich_print(tree2) |
| 239 | + def send_to_server(self): |
| 240 | + config = ConfigObj(settings_conf_path) |
| 241 | + api_key = config["Authentication"]["api-key"].strip() |
| 242 | + server = config["Authentication"]["server_address"].strip() |
| 243 | + |
| 244 | + if not api_key or not server: |
| 245 | + print(Fore.RED + "API key or server address is not set. Please launch the Zeuz Node first and login") |
| 246 | + return |
| 247 | + url = f"{self.server_address}:{self.server_port}{self.server_path}" |
| 248 | + try: |
| 249 | + url = server + "/" if server[-1] != "/" else server |
| 250 | + url += "ai_record_single_action/" |
| 251 | + content = json.dumps({ |
| 252 | + 'page_src': self.xml_str, |
| 253 | + "action_type": "android", |
| 254 | + }) |
| 255 | + headers = { |
| 256 | + "X-Api-Key": api_key, |
| 257 | + } |
| 258 | + |
| 259 | + r = requests.request("POST", url, headers=headers, data=content, verify=False) |
| 260 | + response = r.json() |
| 261 | + if response["info"] == "success": |
| 262 | + r.ok and print("Element sent. You can " + Fore.GREEN + "'Add by AI' " + Fore.RESET + "from server") |
| 263 | + else: |
| 264 | + print(Fore.RED + response["info"]) |
| 265 | + except: |
| 266 | + traceback.print_exc() |
| 267 | + print(Fore.RED + "Failed to send content to AI Engine") |
| 268 | + return |
| 269 | + |
| 270 | + def run(self): |
| 271 | + while True: |
| 272 | + input("Press any key to start capturing...") |
| 273 | + self.wait_for_control_press() |
| 274 | + self.get_frontmost_app() |
| 275 | + self.get_server_port() |
| 276 | + if self.server_port == 0: |
| 277 | + print(Fore.RED + "Server port is not set. Please launch the Zeuz Node first and launch an app") |
| 278 | + continue |
| 279 | + self.get_dump() |
| 280 | + if not self.page_src: |
| 281 | + continue |
| 282 | + self.render_tree() |
| 283 | + self.send_to_server() |
| 284 | + |
| 285 | + time.sleep(0.2) |
| 286 | + |
| 287 | + |
| 288 | +def main(): |
| 289 | + inspector = Inspector() |
| 290 | + inspector.run() |
| 291 | + |
| 292 | +if __name__ == "__main__": |
| 293 | + main() |
0 commit comments