diff --git a/Data/Agent/borealis-agent.py b/Data/Agent/borealis-agent.py index cc8974a..7623bbb 100644 --- a/Data/Agent/borealis-agent.py +++ b/Data/Agent/borealis-agent.py @@ -10,6 +10,7 @@ import concurrent.futures from functools import partial from io import BytesIO import base64 +import traceback import socketio from qasync import QEventLoop @@ -17,7 +18,7 @@ from PyQt5 import QtCore, QtGui, QtWidgets from PIL import ImageGrab # ////////////////////////////////////////////////////////////////////////// -# CORE SECTION: CONFIG MANAGER (do not modify unless you know what you’re doing) +# CORE SECTION: CONFIG MANAGER # ////////////////////////////////////////////////////////////////////////// CONFIG_PATH = os.path.join(os.path.dirname(__file__), "agent_settings.json") DEFAULT_CONFIG = { @@ -36,7 +37,9 @@ class ConfigManager: self.load() def load(self): + print("[DEBUG] Loading config from disk.") if not os.path.exists(self.path): + print("[DEBUG] Config file not found. Creating default.") self.data = DEFAULT_CONFIG.copy() self._write() else: @@ -44,6 +47,7 @@ class ConfigManager: with open(self.path, 'r') as f: loaded = json.load(f) self.data = {**DEFAULT_CONFIG, **loaded} + print("[DEBUG] Config loaded:", self.data) except Exception as e: print(f"[WARN] Failed to parse config: {e}") self.data = DEFAULT_CONFIG.copy() @@ -56,6 +60,7 @@ class ConfigManager: try: with open(self.path, 'w') as f: json.dump(self.data, f, indent=2) + print("[DEBUG] Config written to disk.") except Exception as e: print(f"[ERROR] Could not write config: {e}") @@ -71,83 +76,92 @@ class ConfigManager: return False CONFIG = ConfigManager(CONFIG_PATH) -CONFIG.data['regions'] = {} -CONFIG._write() -# ////////////////////////////////////////////////////////////////////////// -# END CORE SECTION: CONFIG MANAGER -# ////////////////////////////////////////////////////////////////////////// +CONFIG.load() -host = socket.gethostname().lower() -stored_id = CONFIG.data.get('agent_id') -if stored_id: - AGENT_ID = stored_id -else: - AGENT_ID = f"{host}-agent-{uuid.uuid4().hex[:8]}" - CONFIG.data['agent_id'] = AGENT_ID +def init_agent_id(): + if not CONFIG.data.get('agent_id'): + CONFIG.data['agent_id'] = f"{socket.gethostname().lower()}-agent-{uuid.uuid4().hex[:8]}" + CONFIG._write() + return CONFIG.data['agent_id'] + +AGENT_ID = init_agent_id() +print(f"[DEBUG] Using AGENT_ID: {AGENT_ID}") + +def clear_regions_only(): + CONFIG.data['regions'] = CONFIG.data.get('regions', {}) CONFIG._write() -# ////////////////////////////////////////////////////////////////////////// -# CORE SECTION: WEBSOCKET SETUP & HANDLERS (do not modify unless absolutely necessary) +clear_regions_only() + # ////////////////////////////////////////////////////////////////////////// sio = socketio.AsyncClient(reconnection=True, reconnection_attempts=0, reconnection_delay=5) role_tasks = {} overlay_widgets = {} +background_tasks = [] + +async def stop_all_roles(): + print("[DEBUG] Stopping all roles.") + for task in list(role_tasks.values()): + print(f"[DEBUG] Cancelling task for node: {task}") + task.cancel() + role_tasks.clear() + for node_id, widget in overlay_widgets.items(): + print(f"[DEBUG] Closing overlay widget: {node_id}") + try: widget.close() + except Exception as e: print(f"[WARN] Error closing widget: {e}") + overlay_widgets.clear() @sio.event async def connect(): - print(f"[WebSocket] Connected to Agent ID: {AGENT_ID}.") + print(f"[WebSocket] Connected to Borealis Server with Agent ID: {AGENT_ID}") await sio.emit('connect_agent', {"agent_id": AGENT_ID}) await sio.emit('request_config', {"agent_id": AGENT_ID}) @sio.event async def disconnect(): print("[WebSocket] Disconnected from Borealis server.") - for task in list(role_tasks.values()): - task.cancel() - role_tasks.clear() - for widget in list(overlay_widgets.values()): - try: widget.close() - except: pass - overlay_widgets.clear() + await stop_all_roles() CONFIG.data['regions'].clear() CONFIG._write() - CONFIG.load() @sio.on('agent_config') async def on_agent_config(cfg): - print(f"[CONNECTED] Received config with {len(cfg.get('roles',[]))} roles.") - new_ids = {r.get('node_id') for r in cfg.get('roles', []) if r.get('node_id')} + print("[DEBUG] agent_config event received.") + roles = cfg.get('roles', []) + if not roles: + print("[CONFIG] Config Reset by Borealis Server Operator - Awaiting New Config...") + await stop_all_roles() + return + + print(f"[CONFIG] Received New Agent Config with {len(roles)} Role(s).") + + new_ids = {r.get('node_id') for r in roles if r.get('node_id')} old_ids = set(role_tasks.keys()) removed = old_ids - new_ids - # Cancel removed roles for rid in removed: - if rid in CONFIG.data['regions']: - CONFIG.data['regions'].pop(rid, None) + print(f"[DEBUG] Removing node {rid} from regions/overlays.") + CONFIG.data['regions'].pop(rid, None) w = overlay_widgets.pop(rid, None) if w: try: w.close() except: pass - if removed: CONFIG._write() - # Cancel all existing to ensure clean state for task in list(role_tasks.values()): task.cancel() role_tasks.clear() - # Restart everything to ensure roles are re-applied - for role_cfg in cfg.get('roles', []): + for role_cfg in roles: nid = role_cfg.get('node_id') if role_cfg.get('role') == 'screenshot': + print(f"[DEBUG] Starting screenshot task for {nid}") task = asyncio.create_task(screenshot_task(role_cfg)) role_tasks[nid] = task -# ////////////////////////////////////////////////////////////////////////// -# END CORE SECTION: WEBSOCKET SETUP & HANDLERS -# ////////////////////////////////////////////////////////////////////////// +# ---------------- Overlay Widget ---------------- class ScreenshotRegion(QtWidgets.QWidget): def __init__(self, node_id, x=100, y=100, w=300, h=200): super().__init__() @@ -169,14 +183,14 @@ class ScreenshotRegion(QtWidgets.QWidget): p = QtGui.QPainter(self) p.setRenderHint(QtGui.QPainter.Antialiasing) p.setBrush(QtCore.Qt.transparent) - p.setPen(QtGui.QPen(QtGui.QColor(0,255,0),2)) + p.setPen(QtGui.QPen(QtGui.QColor(0, 255, 0), 2)) p.drawRect(self.rect()) hr = self.resize_handle_size - hrect = QtCore.QRect(self.width()-hr, self.height()-hr, hr, hr) - p.fillRect(hrect, QtGui.QColor(0,255,0)) + hrect = QtCore.QRect(self.width() - hr, self.height() - hr, hr, hr) + p.fillRect(hrect, QtGui.QColor(0, 255, 0)) def mousePressEvent(self, e): - if e.button()==QtCore.Qt.LeftButton: + if e.button() == QtCore.Qt.LeftButton: x, y = e.pos().x(), e.pos().y() if x > self.width() - self.resize_handle_size and y > self.height() - self.resize_handle_size: self.resizing = True @@ -202,15 +216,13 @@ class ScreenshotRegion(QtWidgets.QWidget): # ---------------- Screenshot Task ---------------- async def screenshot_task(cfg): nid = cfg.get('node_id') - # If existing region in config, honor that + print(f"[DEBUG] Running screenshot_task for {nid}") r = 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)) - CONFIG.data['regions'][nid] = { - 'x': region[0], 'y': region[1], 'w': region[2], 'h': region[3] - } + CONFIG.data['regions'][nid] = {'x': region[0], 'y': region[1], 'w': region[2], 'h': region[3]} CONFIG._write() if nid not in overlay_widgets: @@ -218,21 +230,25 @@ async def screenshot_task(cfg): overlay_widgets[nid] = widget widget.show() + await sio.emit('agent_screenshot_task', { + 'agent_id': 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=CONFIG.data.get('max_task_workers', DEFAULT_CONFIG['max_task_workers']) - ) + executor = concurrent.futures.ThreadPoolExecutor(max_workers=CONFIG.data.get('max_task_workers', 8)) try: while True: x, y, w, h = overlay_widgets[nid].get_geometry() - prev = CONFIG.data['regions'].get(nid) new_geom = {'x': x, 'y': y, 'w': w, 'h': h} - if prev != new_geom: + if CONFIG.data['regions'].get(nid) != new_geom: CONFIG.data['regions'][nid] = new_geom CONFIG._write() - grab = partial(ImageGrab.grab, bbox=(x, y, x+w, y+h)) + 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') @@ -241,43 +257,79 @@ async def screenshot_task(cfg): 'agent_id': AGENT_ID, 'node_id': nid, 'image_base64': encoded, - 'x': x, 'y': y, 'w': w, 'h': h # Bi-directional live-sync + 'x': x, 'y': y, 'w': w, 'h': h }) await asyncio.sleep(interval) except asyncio.CancelledError: - return + print(f"[TASK] Screenshot role {nid} cancelled.") except Exception as e: print(f"[ERROR] Screenshot task {nid} failed: {e}") + traceback.print_exc() # ---------------- Config Watcher ---------------- async def config_watcher(): + print("[DEBUG] Starting config watcher") while True: - if CONFIG.watch(): pass - await asyncio.sleep(CONFIG.data.get('config_file_watcher_interval', DEFAULT_CONFIG['config_file_watcher_interval'])) + CONFIG.watch() + await asyncio.sleep(CONFIG.data.get('config_file_watcher_interval', 2)) + +# ---------------- Persistent Idle Task ---------------- +async def idle_task(): + print("[Agent] Entering idle state. Awaiting instructions...") + try: + while True: + await asyncio.sleep(60) + print("[DEBUG] Idle task still alive.") + except asyncio.CancelledError: + print("[FATAL] Idle task was cancelled!") + except Exception as e: + print(f"[FATAL] Idle task crashed: {e}") + traceback.print_exc() + +# ---------------- Dummy Qt Widget to Prevent Exit ---------------- +class PersistentWindow(QtWidgets.QWidget): + def __init__(self): + super().__init__() + self.setWindowTitle("KeepAlive") + self.setGeometry(-1000, -1000, 1, 1) + self.setAttribute(QtCore.Qt.WA_DontShowOnScreen) + self.hide() # ////////////////////////////////////////////////////////////////////////// -# CORE SECTION: MAIN & EVENT LOOP (do not modify unless you know what you’re doing) +# MAIN & EVENT LOOP # ////////////////////////////////////////////////////////////////////////// async def connect_loop(): retry = 5 while True: try: - url = CONFIG.data.get('borealis_server_url', DEFAULT_CONFIG['borealis_server_url']) + url = CONFIG.data.get('borealis_server_url', "http://localhost:5000") print(f"[WebSocket] Connecting to {url}...") await sio.connect(url, transports=['websocket']) break - except: - print(f"[WebSocket] Server unavailable, retrying in {retry}s...") + except Exception as e: + print(f"[WebSocket] Server unavailable: {e}. Retrying in {retry}s...") await asyncio.sleep(retry) if __name__ == '__main__': + print("[DEBUG] Starting QApplication and QEventLoop") app = QtWidgets.QApplication(sys.argv) loop = QEventLoop(app) asyncio.set_event_loop(loop) - with loop: - loop.create_task(config_watcher()) - loop.create_task(connect_loop()) - loop.run_forever() -# ////////////////////////////////////////////////////////////////////////// -# END CORE SECTION: MAIN & EVENT LOOP -# ////////////////////////////////////////////////////////////////////////// + + # Dummy window to keep PyQt event loop alive + dummy_window = PersistentWindow() + dummy_window.show() + print("[DEBUG] Dummy window shown to prevent Qt exit") + + try: + with loop: + print("[DEBUG] Scheduling tasks into event loop") + background_tasks.append(loop.create_task(config_watcher())) + background_tasks.append(loop.create_task(connect_loop())) + background_tasks.append(loop.create_task(idle_task())) + loop.run_forever() + except Exception as e: + print(f"[FATAL] Event loop crashed: {e}") + traceback.print_exc() + finally: + print("[FATAL] Agent exited unexpectedly.") diff --git a/Data/Server/WebUI/src/Status_Bar.jsx b/Data/Server/WebUI/src/Status_Bar.jsx index 1a0d19e..e2ec75d 100644 --- a/Data/Server/WebUI/src/Status_Bar.jsx +++ b/Data/Server/WebUI/src/Status_Bar.jsx @@ -74,7 +74,7 @@ export default function StatusBar() { - API Server: + Backend API Server: { const [isConnected, setIsConnected] = useState(false); const prevRolesRef = useRef([]); + // ---------------- Agent List & Sorting ---------------- const agentList = useMemo(() => { if (!agents || typeof agents !== "object") return []; return Object.entries(agents) @@ -23,6 +24,7 @@ const BorealisAgentNode = ({ id, data }) => { .sort((a, b) => b.last_seen - a.last_seen); }, [agents]); + // ---------------- Periodic Agent Fetching ---------------- useEffect(() => { const fetchAgents = () => { fetch("/api/agents") @@ -35,6 +37,7 @@ const BorealisAgentNode = ({ id, data }) => { return () => clearInterval(interval); }, []); + // ---------------- Node Data Sync ---------------- useEffect(() => { setNodes((nds) => nds.map((n) => @@ -44,6 +47,7 @@ const BorealisAgentNode = ({ id, data }) => { setIsConnected(false); }, [selectedAgent]); + // ---------------- Attached Role Collection ---------------- const attachedRoleIds = useMemo( () => edges @@ -62,8 +66,9 @@ const BorealisAgentNode = ({ id, data }) => { .filter((r) => r); }, [attachedRoleIds, getNodes]); + // ---------------- Provision Role Logic ---------------- const provisionRoles = useCallback((roles) => { - if (!selectedAgent || roles.length === 0) return; + if (!selectedAgent) return; // Allow empty roles but require agent fetch("/api/agent/provision", { method: "POST", headers: { "Content-Type": "application/json" }, @@ -78,7 +83,7 @@ const BorealisAgentNode = ({ id, data }) => { const handleConnect = useCallback(() => { const roles = getAttachedRoles(); - provisionRoles(roles); + provisionRoles(roles); // Always call even with empty roles }, [getAttachedRoles, provisionRoles]); const handleDisconnect = useCallback(() => { @@ -95,6 +100,7 @@ const BorealisAgentNode = ({ id, data }) => { .catch(() => {}); }, [selectedAgent]); + // ---------------- Auto-Provision When Roles Change ---------------- useEffect(() => { const newRoles = getAttachedRoles(); const prevSerialized = JSON.stringify(prevRolesRef.current || []); @@ -102,8 +108,9 @@ const BorealisAgentNode = ({ id, data }) => { if (isConnected && newSerialized !== prevSerialized) { provisionRoles(newRoles); } - }, [attachedRoleIds, isConnected]); + }, [attachedRoleIds, isConnected, getAttachedRoles, provisionRoles]); + // ---------------- Status Label ---------------- const selectedAgentStatus = useMemo(() => { if (!selectedAgent) return "Unassigned"; const agent = agents[selectedAgent]; @@ -111,6 +118,7 @@ const BorealisAgentNode = ({ id, data }) => { return agent.status === "provisioned" ? "Connected" : "Available"; }, [agents, selectedAgent]); + // ---------------- Render ---------------- return (