import os import asyncio import concurrent.futures from functools import partial from io import BytesIO import base64 import traceback import random import importlib.util from PyQt5 import QtCore, QtGui, QtWidgets from PIL import ImageGrab class RolesContext: def __init__(self, sio, agent_id, config): self.sio = sio self.agent_id = agent_id self.config = config # Load macro engines from the local Python_API_Endpoints directory MACRO_ENGINE_PATH = os.path.join(os.path.dirname(__file__), "Python_API_Endpoints", "macro_engines.py") spec = importlib.util.spec_from_file_location("macro_engines", MACRO_ENGINE_PATH) macro_engines = importlib.util.module_from_spec(spec) spec.loader.exec_module(macro_engines) # Overlay visuals overlay_green_thickness = 4 overlay_gray_thickness = 2 handle_size = overlay_green_thickness * 2 extra_top_padding = overlay_green_thickness * 2 + 4 # Track active screenshot overlay widgets per node_id overlay_widgets: dict[str, QtWidgets.QWidget] = {} def get_window_list(): """Return a list of windows from macro engines.""" try: return macro_engines.list_windows() except Exception: return [] def close_overlay(node_id: str): w = overlay_widgets.pop(node_id, None) if w: try: w.close() except Exception: pass def close_all_overlays(): for node_id, widget in list(overlay_widgets.items()): try: widget.close() except Exception: pass overlay_widgets.clear() class ScreenshotRegion(QtWidgets.QWidget): def __init__(self, ctx: RolesContext, node_id, x=100, y=100, w=300, h=200, alias=None): super().__init__() self.ctx = ctx self.node_id = node_id self.alias = alias self.setGeometry( x - handle_size, y - handle_size - extra_top_padding, w + handle_size * 2, h + handle_size * 2 + extra_top_padding, ) self.setWindowFlags(QtCore.Qt.FramelessWindowHint | QtCore.Qt.WindowStaysOnTopHint) self.setAttribute(QtCore.Qt.WA_TranslucentBackground) self.resize_dir = None self.drag_offset = None self._start_geom = None self._start_pos = None self.setMouseTracking(True) def paintEvent(self, event): p = QtGui.QPainter(self) p.setRenderHint(QtGui.QPainter.Antialiasing) w = self.width() h = self.height() # draw gray capture box p.setPen(QtGui.QPen(QtGui.QColor(130, 130, 130), overlay_gray_thickness)) p.drawRect(handle_size, handle_size + extra_top_padding, w - handle_size * 2, h - handle_size * 2 - extra_top_padding) p.setPen(QtCore.Qt.NoPen) p.setBrush(QtGui.QBrush(QtGui.QColor(0, 191, 255))) edge = overlay_green_thickness * 3 # corner handles p.drawRect(0, extra_top_padding, edge, overlay_green_thickness) p.drawRect(0, extra_top_padding, overlay_green_thickness, edge) p.drawRect(w - edge, extra_top_padding, edge, overlay_green_thickness) p.drawRect(w - overlay_green_thickness, extra_top_padding, overlay_green_thickness, edge) p.drawRect(0, h - overlay_green_thickness, edge, overlay_green_thickness) p.drawRect(0, h - edge, overlay_green_thickness, edge) p.drawRect(w - edge, h - overlay_green_thickness, edge, overlay_green_thickness) p.drawRect(w - overlay_green_thickness, h - edge, overlay_green_thickness, edge) # side handles long = overlay_green_thickness * 6 p.drawRect((w - long) // 2, extra_top_padding, long, overlay_green_thickness) p.drawRect((w - long) // 2, h - overlay_green_thickness, long, overlay_green_thickness) p.drawRect(0, (h + extra_top_padding - long) // 2, overlay_green_thickness, long) p.drawRect(w - overlay_green_thickness, (h + extra_top_padding - long) // 2, overlay_green_thickness, long) # grabber bar bar_width = overlay_green_thickness * 6 bar_height = overlay_green_thickness bar_x = (w - bar_width) // 2 bar_y = 6 p.setBrush(QtGui.QColor(0, 191, 255)) p.drawRect(bar_x, bar_y - bar_height - 10, bar_width, bar_height * 4) def get_geometry(self): g = self.geometry() return ( g.x() + handle_size, g.y() + handle_size + extra_top_padding, g.width() - handle_size * 2, g.height() - handle_size * 2 - extra_top_padding, ) def mousePressEvent(self, e): if e.button() == QtCore.Qt.LeftButton: pos = e.pos() bar_width = overlay_green_thickness * 6 bar_height = overlay_green_thickness bar_x = (self.width() - bar_width) // 2 bar_y = 2 bar_rect = QtCore.QRect(bar_x, bar_y, bar_width, bar_height) if bar_rect.contains(pos): self.drag_offset = e.globalPos() - self.frameGeometry().topLeft() return m = handle_size dirs = [] if pos.x() <= m: dirs.append('left') if pos.x() >= self.width() - m: dirs.append('right') if pos.y() <= m + extra_top_padding: dirs.append('top') if pos.y() >= self.height() - m: dirs.append('bottom') if dirs: self.resize_dir = '_'.join(dirs) self._start_geom = self.geometry() self._start_pos = e.globalPos() else: self.drag_offset = e.globalPos() - self.frameGeometry().topLeft() def mouseMoveEvent(self, e): if self.resize_dir and self._start_geom and self._start_pos: dx = e.globalX() - self._start_pos.x() dy = e.globalY() - self._start_pos.y() geom = QtCore.QRect(self._start_geom) if 'left' in self.resize_dir: new_x = geom.x() + dx new_w = geom.width() - dx geom.setX(new_x) geom.setWidth(new_w) if 'right' in self.resize_dir: geom.setWidth(self._start_geom.width() + dx) if 'top' in self.resize_dir: new_y = geom.y() + dy new_h = geom.height() - dy geom.setY(new_y) geom.setHeight(new_h) if 'bottom' in self.resize_dir: geom.setHeight(self._start_geom.height() + dy) self.setGeometry(geom) elif self.drag_offset and e.buttons() & QtCore.Qt.LeftButton: self.move(e.globalPos() - self.drag_offset) def mouseReleaseEvent(self, e): self.drag_offset = None self.resize_dir = None self._start_geom = None self._start_pos = None x, y, w, h = self.get_geometry() self.ctx.config.data['regions'][self.node_id] = {'x': x, 'y': y, 'w': w, 'h': h} try: self.ctx.config._write() except Exception: pass # Emit a zero-image update so the server knows new geometry immediately asyncio.create_task(self.ctx.sio.emit('agent_screenshot_task', { 'agent_id': self.ctx.agent_id, 'node_id': self.node_id, 'image_base64': '', 'x': x, 'y': y, 'w': w, 'h': h })) async def screenshot_task(ctx: RolesContext, cfg): nid = cfg.get('node_id') alias = cfg.get('alias', '') r = ctx.config.data['regions'].get(nid) if r: region = (r['x'], r['y'], r['w'], r['h']) else: region = ( cfg.get('x', 100), cfg.get('y', 100), cfg.get('w', 300), cfg.get('h', 200), ) ctx.config.data['regions'][nid] = {'x': region[0], 'y': region[1], 'w': region[2], 'h': region[3]} try: ctx.config._write() except Exception: pass if nid not in overlay_widgets: widget = ScreenshotRegion(ctx, nid, *region, alias=alias) overlay_widgets[nid] = widget widget.show() await ctx.sio.emit('agent_screenshot_task', { 'agent_id': ctx.agent_id, 'node_id': nid, 'image_base64': '', 'x': region[0], 'y': region[1], 'w': region[2], 'h': region[3] }) interval = cfg.get('interval', 1000) / 1000.0 loop = asyncio.get_event_loop() executor = concurrent.futures.ThreadPoolExecutor(max_workers=ctx.config.data.get('max_task_workers', 8)) try: while True: x, y, w, h = overlay_widgets[nid].get_geometry() grab = partial(ImageGrab.grab, bbox=(x, y, x + w, y + h)) img = await loop.run_in_executor(executor, grab) buf = BytesIO(); img.save(buf, format='PNG') encoded = base64.b64encode(buf.getvalue()).decode('utf-8') await ctx.sio.emit('agent_screenshot_task', { 'agent_id': ctx.agent_id, 'node_id': nid, 'image_base64': encoded, 'x': x, 'y': y, 'w': w, 'h': h }) await asyncio.sleep(interval) except asyncio.CancelledError: print(f"[TASK] Screenshot role {nid} cancelled.") except Exception as e: print(f"[ERROR] Screenshot task {nid} failed: {e}") traceback.print_exc() async def macro_task(ctx: RolesContext, cfg): """Improved macro task supporting operation modes, live config, and feedback.""" nid = cfg.get('node_id') last_trigger_value = 0 has_run_once = False while True: window_handle = cfg.get('window_handle') macro_type = cfg.get('macro_type', 'keypress') operation_mode = cfg.get('operation_mode', 'Continuous') key = cfg.get('key') text = cfg.get('text') interval_ms = int(cfg.get('interval_ms', 1000)) randomize = cfg.get('randomize_interval', False) random_min = int(cfg.get('random_min', 750)) random_max = int(cfg.get('random_max', 950)) active = cfg.get('active', True) trigger = int(cfg.get('trigger', 0)) async def emit_macro_status(success, message=""): await ctx.sio.emit('macro_status', { "agent_id": ctx.agent_id, "node_id": nid, "success": success, "message": message, "timestamp": int(asyncio.get_event_loop().time() * 1000) }) if not (active is True or str(active).lower() == "true"): await asyncio.sleep(0.2) continue try: send_macro = False if operation_mode == "Run Once": if not has_run_once: send_macro = True has_run_once = True elif operation_mode == "Continuous": send_macro = True elif operation_mode == "Trigger-Continuous": send_macro = (trigger == 1) elif operation_mode == "Trigger-Once": if last_trigger_value == 0 and trigger == 1: send_macro = True last_trigger_value = trigger else: send_macro = True if send_macro: if macro_type == 'keypress' and key: result = macro_engines.send_keypress_to_window(window_handle, key) elif macro_type == 'typed_text' and text: result = macro_engines.type_text_to_window(window_handle, text) else: await emit_macro_status(False, "Invalid macro type or missing key/text") await asyncio.sleep(0.2) continue if isinstance(result, tuple): success, err = result else: success, err = bool(result), "" if success: await emit_macro_status(True, f"Macro sent: {macro_type}") else: await emit_macro_status(False, err or "Unknown macro engine failure") else: await asyncio.sleep(0.05) if send_macro: if randomize: ms = random.randint(random_min, random_max) else: ms = interval_ms await asyncio.sleep(ms / 1000.0) else: await asyncio.sleep(0.1) except asyncio.CancelledError: print(f"[TASK] Macro role {nid} cancelled.") break except Exception as e: print(f"[ERROR] Macro task {nid} failed: {e}") traceback.print_exc() await emit_macro_status(False, str(e)) await asyncio.sleep(0.5)