#////////// PROJECT FILE SEPARATION LINE ////////// CODE AFTER THIS LINE ARE FROM: /Data/Agent/borealis-agent.py import sys import uuid import time import base64 import threading import socketio from io import BytesIO import socket import os from PyQt5 import QtCore, QtGui, QtWidgets from PIL import ImageGrab # ---------------- Configuration ---------------- SERVER_URL = os.environ.get("SERVER_URL", "http://localhost:5000") HOSTNAME = socket.gethostname().lower() RANDOM_SUFFIX = uuid.uuid4().hex[:8] AGENT_ID = f"{HOSTNAME}-agent-{RANDOM_SUFFIX}" # ---------------- App State ---------------- app_instance = None overlay_widgets = {} region_launchers = {} running_roles = {} running_threads = {} # ---------------- Socket Setup ---------------- # Enable automatic reconnection with retries in background sio = socketio.Client(reconnection=True, reconnection_attempts=0, reconnection_delay=5) @sio.event def connect(): print(f"[WebSocket] Agent ID: {AGENT_ID} connected to Borealis.") sio.emit('connect_agent', {"agent_id": AGENT_ID, "hostname": HOSTNAME}) sio.emit('request_config', {"agent_id": AGENT_ID}) @sio.event def disconnect(): print("[WebSocket] Lost connection to Borealis server.") @sio.on('agent_config') def on_agent_config(config): print("[PROVISIONED] Received new configuration from Borealis.") roles = config.get("roles", []) stop_all_roles() for role in roles: start_role_thread(role) # ---------------- Overlay Class ---------------- class ScreenshotRegion(QtWidgets.QWidget): def __init__(self, node_id, x=100, y=100, w=300, h=200): super().__init__() self.node_id = node_id self.setGeometry(x, y, w, h) self.setWindowFlags(QtCore.Qt.FramelessWindowHint | QtCore.Qt.WindowStaysOnTopHint) self.setAttribute(QtCore.Qt.WA_TranslucentBackground) self.drag_offset = None self.resizing = False self.resize_handle_size = 12 self.setVisible(True) self.label = QtWidgets.QLabel(self) self.label.setText(f"{node_id[:8]}") self.label.setStyleSheet("color: lime; background: transparent; font-size: 10px;") self.label.move(8, 4) self.setMouseTracking(True) def paintEvent(self, event): painter = QtGui.QPainter(self) painter.setRenderHint(QtGui.QPainter.Antialiasing) painter.setBrush(QtCore.Qt.transparent) painter.setPen(QtGui.QPen(QtGui.QColor(0, 255, 0), 2)) painter.drawRect(self.rect()) handle_rect = QtCore.QRect( self.width() - self.resize_handle_size, self.height() - self.resize_handle_size, self.resize_handle_size, self.resize_handle_size ) painter.fillRect(handle_rect, QtGui.QColor(0, 255, 0)) def mousePressEvent(self, event): if event.button() == QtCore.Qt.LeftButton: if event.pos().x() > self.width() - self.resize_handle_size and \ event.pos().y() > self.height() - self.resize_handle_size: self.resizing = True else: self.drag_offset = event.globalPos() - self.frameGeometry().topLeft() def mouseMoveEvent(self, event): if self.resizing: new_width = max(event.pos().x(), 100) new_height = max(event.pos().y(), 80) self.resize(new_width, new_height) elif event.buttons() & QtCore.Qt.LeftButton and self.drag_offset: self.move(event.globalPos() - self.drag_offset) def mouseReleaseEvent(self, event): self.resizing = False self.drag_offset = None def get_geometry(self): geo = self.geometry() return geo.x(), geo.y(), geo.width(), geo.height() # ---------------- Region UI Handler ---------------- class RegionLauncher(QtCore.QObject): trigger = QtCore.pyqtSignal(int, int, int, int) def __init__(self, node_id): super().__init__() self.node_id = node_id self.trigger.connect(self.handle) def handle(self, x, y, w, h): print(f"[Overlay] Launching overlay for {self.node_id} at ({x},{y},{w},{h})") if self.node_id in overlay_widgets: return widget = ScreenshotRegion(self.node_id, x, y, w, h) overlay_widgets[self.node_id] = widget widget.show() # ---------------- Role Management ---------------- def stop_all_roles(): for node_id, thread in running_threads.items(): if thread and thread.is_alive(): print(f"[Role] Terminating previous task: {node_id}") running_roles.clear() running_threads.clear() def start_role_thread(role_cfg): role = role_cfg.get("role") node_id = role_cfg.get("node_id") if not role or not node_id: print("[ERROR] Invalid role configuration (missing role or node_id).") return if role == "screenshot": thread = threading.Thread(target=run_screenshot_loop, args=(node_id, role_cfg), daemon=True) else: print(f"[SKIP] Unknown role: {role}") return running_roles[node_id] = role_cfg running_threads[node_id] = thread thread.start() print(f"[Role] Started task: {role} ({node_id})") # ---------------- Screenshot Role Loop ---------------- def run_screenshot_loop(node_id, cfg): interval = cfg.get("interval", 1000) visible = cfg.get("visible", True) x = cfg.get("x", 100) y = cfg.get("y", 100) w = cfg.get("w", 300) h = cfg.get("h", 200) if node_id not in region_launchers: launcher = RegionLauncher(node_id) region_launchers[node_id] = launcher QtCore.QTimer.singleShot(0, lambda: launcher.trigger.emit(x, y, w, h)) widget = overlay_widgets.get(node_id) if widget: widget.setGeometry(x, y, w, h) widget.setVisible(visible) while True: try: if node_id in overlay_widgets: widget = overlay_widgets[node_id] x, y, w, h = widget.get_geometry() print(f"[Capture] Screenshot task {node_id} at ({x},{y},{w},{h})") img = ImageGrab.grab(bbox=(x, y, x + w, y + h)) buffer = BytesIO() img.save(buffer, format="PNG") encoded = base64.b64encode(buffer.getvalue()).decode("utf-8") sio.emit("agent_screenshot_task", { "agent_id": AGENT_ID, "node_id": node_id, "image_base64": encoded }) except Exception as e: print(f"[ERROR] Screenshot task {node_id} failed: {e}") time.sleep(interval / 1000) # ---------------- Main ---------------- if __name__ == "__main__": app_instance = QtWidgets.QApplication(sys.argv) retry_interval = 5 # seconds between connection attempts def connect_loop(): while True: try: print(f"[WebSocket] Connecting to {SERVER_URL}...") sio.connect(SERVER_URL, transports=["websocket"], wait=False) break except Exception: print(f"[WebSocket] Borealis Server is Not Running - Retrying in {retry_interval} seconds...") time.sleep(retry_interval) threading.Thread(target=connect_loop, daemon=True).start() sys.exit(app_instance.exec_())