From 41c52590f4173a545ef7d82e4381db54b8e5405a Mon Sep 17 00:00:00 2001 From: Nicole Rappe Date: Wed, 15 Oct 2025 04:40:21 -0600 Subject: [PATCH] Support selecting agent context for screenshots --- Data/Agent/Roles/role_Screenshot.py | 16 ++ Data/Agent/agent.py | 21 +- Data/Agent/role_manager.py | 4 + .../Node_Agent_Role_Screenshot.jsx | 71 +++++-- .../WebUI/src/nodes/Agent/Node_Agent.jsx | 186 +++++++++++++++--- Data/Server/server.py | 49 ++++- 6 files changed, 291 insertions(+), 56 deletions(-) diff --git a/Data/Agent/Roles/role_Screenshot.py b/Data/Agent/Roles/role_Screenshot.py index 45af8d8..7b0709d 100644 --- a/Data/Agent/Roles/role_Screenshot.py +++ b/Data/Agent/Roles/role_Screenshot.py @@ -87,6 +87,15 @@ def _coerce_text(value): return '' +def _normalize_mode(value): + text = _coerce_text(value).strip().lower() + if text in {'interactive', 'currentuser', 'user'}: + return 'currentuser' + if text in {'system', 'svc', 'service'}: + return 'system' + return '' + + class ScreenshotRegion(QtWidgets.QWidget): def __init__(self, ctx, node_id, x=100, y=100, w=300, h=200, alias=None): super().__init__() @@ -376,6 +385,8 @@ class Role: 'h': _coerce_int(cfg.get('h'), 200, minimum=1), 'visible': _coerce_bool(cfg.get('visible'), True), 'alias': _coerce_text(cfg.get('alias')), + 'target_agent_mode': _normalize_mode(cfg.get('target_agent_mode')), + 'target_agent_host': _coerce_text(cfg.get('target_agent_host')), } return norm except Exception: @@ -387,6 +398,11 @@ class Role: if not nid: return + target_mode = cfg.get('target_agent_mode') or '' + current_mode = getattr(self.ctx, 'service_mode', '') or '' + if target_mode and current_mode and target_mode != current_mode: + return + alias = cfg.get('alias', '') visible = cfg.get('visible', True) reg = self.ctx.config.data.setdefault('regions', {}) diff --git a/Data/Agent/agent.py b/Data/Agent/agent.py index 9140ad9..53afe21 100644 --- a/Data/Agent/agent.py +++ b/Data/Agent/agent.py @@ -72,6 +72,7 @@ def _bootstrap_log(msg: str): # Headless/service mode flag (skip Qt and interactive UI) SYSTEM_SERVICE_MODE = ('--system-service' in sys.argv) or (os.environ.get('BOREALIS_AGENT_MODE') == 'system') +SERVICE_MODE = 'system' if SYSTEM_SERVICE_MODE else 'currentuser' _bootstrap_log(f'agent.py loaded; SYSTEM_SERVICE_MODE={SYSTEM_SERVICE_MODE}; argv={sys.argv!r}') def _argv_get(flag: str, default: str = None): try: @@ -859,7 +860,8 @@ async def send_heartbeat(): "agent_id": AGENT_ID, "hostname": socket.gethostname(), "agent_operating_system": detect_agent_os(), - "last_seen": int(time.time()) + "last_seen": int(time.time()), + "service_mode": SERVICE_MODE, } await sio.emit("agent_heartbeat", payload) # Also report collector status alive ping. @@ -872,6 +874,7 @@ async def send_heartbeat(): 'agent_id': AGENT_ID, 'hostname': socket.gethostname(), 'active': True, + 'service_mode': SERVICE_MODE, 'last_user': f"{os.environ.get('USERDOMAIN') or socket.gethostname()}\\{getpass.getuser()}" }) else: @@ -879,6 +882,7 @@ async def send_heartbeat(): 'agent_id': AGENT_ID, 'hostname': socket.gethostname(), 'active': True, + 'service_mode': SERVICE_MODE, }) except Exception: pass @@ -1203,7 +1207,7 @@ async def send_agent_details_once(): async def connect(): print(f"[INFO] Successfully Connected to Borealis Server!") _log_agent('Connected to server.') - await sio.emit('connect_agent', {"agent_id": AGENT_ID}) + await sio.emit('connect_agent', {"agent_id": AGENT_ID, "service_mode": SERVICE_MODE}) # Send an immediate heartbeat so the UI can populate instantly. try: @@ -1211,7 +1215,8 @@ async def connect(): "agent_id": AGENT_ID, "hostname": socket.gethostname(), "agent_operating_system": detect_agent_os(), - "last_seen": int(time.time()) + "last_seen": int(time.time()), + "service_mode": SERVICE_MODE, }) except Exception as e: print(f"[WARN] initial heartbeat failed: {e}") @@ -1225,6 +1230,7 @@ async def connect(): 'agent_id': AGENT_ID, 'hostname': socket.gethostname(), 'active': True, + 'service_mode': SERVICE_MODE, 'last_user': f"{os.environ.get('USERDOMAIN') or socket.gethostname()}\\{getpass.getuser()}" }) else: @@ -1232,6 +1238,7 @@ async def connect(): 'agent_id': AGENT_ID, 'hostname': socket.gethostname(), 'active': True, + 'service_mode': SERVICE_MODE, }) except Exception: pass @@ -1542,9 +1549,10 @@ if __name__=='__main__': # Initialize roles context for role tasks # Initialize role manager and hot-load roles from Roles/ try: - hooks = {'send_service_control': send_service_control, 'get_server_url': get_server_url} + base_hooks = {'send_service_control': send_service_control, 'get_server_url': get_server_url} if not SYSTEM_SERVICE_MODE: # Load interactive-context roles (tray/UI, current-user execution, screenshot, etc.) + hooks_interactive = {**base_hooks, 'service_mode': 'currentuser'} ROLE_MANAGER = RoleManager( base_dir=os.path.dirname(__file__), context='interactive', @@ -1552,12 +1560,13 @@ if __name__=='__main__': agent_id=AGENT_ID, config=CONFIG, loop=loop, - hooks=hooks, + hooks=hooks_interactive, ) ROLE_MANAGER.load() # Load system roles only when running in SYSTEM service mode ROLE_MANAGER_SYS = None if SYSTEM_SERVICE_MODE: + hooks_system = {**base_hooks, 'service_mode': 'system'} ROLE_MANAGER_SYS = RoleManager( base_dir=os.path.dirname(__file__), context='system', @@ -1565,7 +1574,7 @@ if __name__=='__main__': agent_id=AGENT_ID, config=CONFIG, loop=loop, - hooks=hooks, + hooks=hooks_system, ) ROLE_MANAGER_SYS.load() except Exception as e: diff --git a/Data/Agent/role_manager.py b/Data/Agent/role_manager.py index 2f90af0..1731ed2 100644 --- a/Data/Agent/role_manager.py +++ b/Data/Agent/role_manager.py @@ -24,6 +24,10 @@ class RoleManager: self.config = config self.loop = loop self.hooks = hooks or {} + try: + self.service_mode = (self.hooks.get('service_mode') or '').strip().lower() + except Exception: + self.service_mode = '' def __init__(self, base_dir: str, context: str, sio, agent_id: str, config, loop, hooks: Optional[dict] = None): self.base_dir = base_dir diff --git a/Data/Server/WebUI/src/nodes/Agent Roles/Node_Agent_Role_Screenshot.jsx b/Data/Server/WebUI/src/nodes/Agent Roles/Node_Agent_Role_Screenshot.jsx index 50a87f7..fa18478 100644 --- a/Data/Server/WebUI/src/nodes/Agent Roles/Node_Agent_Role_Screenshot.jsx +++ b/Data/Server/WebUI/src/nodes/Agent Roles/Node_Agent_Role_Screenshot.jsx @@ -1,5 +1,5 @@ ////////// PROJECT FILE SEPARATION LINE ////////// CODE AFTER THIS LINE ARE FROM: /Data/WebUI/src/nodes/Agent/Node_Agent_Role_Screenshot.jsx -import React, { useEffect, useRef, useState } from "react"; +import React, { useCallback, useEffect, useRef, useState } from "react"; import { Handle, Position, useReactFlow, useStore } from "reactflow"; import ShareIcon from "@mui/icons-material/Share"; import IconButton from "@mui/material/IconButton"; @@ -21,6 +21,17 @@ const AgentScreenshotNode = ({ id, data }) => { const { setNodes, getNodes } = useReactFlow(); const edges = useStore(state => state.edges); + const resolveAgentData = useCallback(() => { + try { + const agentEdge = edges.find(e => e.target === id && e.sourceHandle === "provisioner"); + const agentNode = getNodes().find(n => n.id === agentEdge?.source); + return agentNode?.data || null; + } catch (err) { + return null; + } + }, [edges, getNodes, id]); + + // Core config values pulled from sidebar config (with defaults) const interval = parseInt(data?.interval || 1000, 10) || 1000; const region = { @@ -32,6 +43,11 @@ const AgentScreenshotNode = ({ id, data }) => { const visible = (data?.visible ?? "true") === "true"; const alias = data?.alias || ""; const [imageBase64, setImageBase64] = useState(data?.value || ""); + const agentData = resolveAgentData(); + const targetModeLabel = ((agentData?.agent_mode || "").toString().toLowerCase() === "system") + ? "SYSTEM Agent" + : "CURRENTUSER Agent"; + const targetHostLabel = (agentData?.agent_host || "").toString(); // Always push current imageBase64 into BorealisValueBus at the global update rate useEffect(() => { @@ -56,14 +72,9 @@ const AgentScreenshotNode = ({ id, data }) => { const handleScreenshot = (payload) => { if (payload?.node_id !== id) return; // Additionally ensure payload is from the agent connected upstream of this node - try { - const agentEdge = edges.find(e => e.target === id && e.sourceHandle === "provisioner"); - const agentNode = getNodes().find(n => n.id === agentEdge?.source); - const selectedAgentId = agentNode?.data?.agent_id; - if (!selectedAgentId || payload?.agent_id !== selectedAgentId) return; - } catch (err) { - return; // fail-closed if we cannot resolve upstream agent - } + const agentData = resolveAgentData(); + const selectedAgentId = agentData?.agent_id; + if (!selectedAgentId || payload?.agent_id !== selectedAgentId) return; if (payload.image_base64) { setImageBase64(payload.image_base64); @@ -86,24 +97,30 @@ const AgentScreenshotNode = ({ id, data }) => { socket.on("agent_screenshot_task", handleScreenshot); return () => socket.off("agent_screenshot_task", handleScreenshot); - }, [id, setNodes, edges, getNodes]); + }, [id, setNodes, resolveAgentData]); // Register this node for the agent provisioning sync window.__BorealisInstructionNodes = window.__BorealisInstructionNodes || {}; - window.__BorealisInstructionNodes[id] = () => ({ - node_id: id, - role: "screenshot", - interval, - visible, - alias, - ...region - }); + window.__BorealisInstructionNodes[id] = () => { + const agentData = resolveAgentData() || {}; + const modeRaw = (agentData.agent_mode || "").toString().toLowerCase(); + const targetMode = modeRaw === "system" ? "system" : "currentuser"; + return { + node_id: id, + role: "screenshot", + interval, + visible, + alias, + target_agent_mode: targetMode, + target_agent_host: agentData.agent_host || "", + ...region + }; + }; // Manual live view copy button const handleCopyLiveViewLink = () => { - const agentEdge = edges.find(e => e.target === id && e.sourceHandle === "provisioner"); - const agentNode = getNodes().find(n => n.id === agentEdge?.source); - const selectedAgentId = agentNode?.data?.agent_id; + const agentData = resolveAgentData(); + const selectedAgentId = agentData?.agent_id; if (!selectedAgentId) { alert("No valid agent connection found."); @@ -132,6 +149,17 @@ const AgentScreenshotNode = ({ id, data }) => {
Interval: {interval} ms
+
+ Agent Context: {targetModeLabel} +
+
+ Target Host:{" "} + {targetHostLabel ? ( + targetHostLabel + ) : ( + unknown + )} +
Overlay: {visible ? "Yes" : "No"}
@@ -165,6 +193,7 @@ Capture a live screenshot of a defined region from a remote Borealis Agent. - Optionally show a visual overlay with a label - Pushes base64 PNG stream to downstream nodes - Use copy button to share live view URL +- Targets the CURRENTUSER or SYSTEM agent context selected upstream `.trim(), content: "Capture screenshot region via agent", component: AgentScreenshotNode, diff --git a/Data/Server/WebUI/src/nodes/Agent/Node_Agent.jsx b/Data/Server/WebUI/src/nodes/Agent/Node_Agent.jsx index 2dd39db..063149b 100644 --- a/Data/Server/WebUI/src/nodes/Agent/Node_Agent.jsx +++ b/Data/Server/WebUI/src/nodes/Agent/Node_Agent.jsx @@ -8,22 +8,61 @@ const BorealisAgentNode = ({ id, data }) => { const edges = useStore((state) => state.edges); const [agents, setAgents] = useState({}); const [selectedAgent, setSelectedAgent] = useState(data.agent_id || ""); + const [selectedHost, setSelectedHost] = useState(data.agent_host || ""); + const initialMode = (data.agent_mode || "currentuser").toLowerCase(); + const [selectedMode, setSelectedMode] = useState( + initialMode === "system" ? "system" : "currentuser" + ); const [isConnected, setIsConnected] = useState(false); const prevRolesRef = useRef([]); - // Agent List Sorted (Online First) - const agentList = useMemo(() => { - if (!agents || typeof agents !== "object") return []; - return Object.entries(agents) - .map(([aid, info]) => ({ - id: aid, - status: info?.status || "offline", - last_seen: info?.last_seen || 0 - })) - .filter(({ status }) => status !== "offline") - .sort((a, b) => b.last_seen - a.last_seen); + // Group agents by hostname and execution context + const agentsByHostname = useMemo(() => { + if (!agents || typeof agents !== "object") return {}; + const grouped = {}; + Object.entries(agents).forEach(([aid, info]) => { + if (!info || typeof info !== "object") return; + const status = (info.status || "").toString().toLowerCase(); + if (status === "offline") return; + const host = (info.hostname || info.agent_hostname || "").trim() || "unknown"; + const modeRaw = (info.service_mode || "").toString().toLowerCase(); + const mode = modeRaw === "system" ? "system" : "currentuser"; + if (!grouped[host]) { + grouped[host] = { currentuser: null, system: null }; + } + grouped[host][mode] = { + agent_id: aid, + status: info.status || "offline", + last_seen: info.last_seen || 0, + info, + }; + }); + return grouped; }, [agents]); + const hostOptions = useMemo(() => { + const entries = Object.entries(agentsByHostname) + .map(([host, contexts]) => { + const candidates = [contexts.currentuser, contexts.system].filter(Boolean); + if (!candidates.length) return null; + const badge = (record) => { + if (!record) return "✕"; + const st = (record.status || "").toString().toLowerCase(); + if (st === "provisioned") return "✓"; + return "•"; + }; + const label = `${host} (CURRENTUSER ${badge(contexts.currentuser)}, SYSTEM ${badge(contexts.system)})`; + const latest = Math.max(...candidates.map((r) => r.last_seen || 0)); + return { host, label, contexts, latest }; + }) + .filter(Boolean) + .sort((a, b) => { + if (b.latest !== a.latest) return b.latest - a.latest; + return a.host.localeCompare(b.host); + }); + return entries; + }, [agentsByHostname]); + // Fetch Agents Periodically useEffect(() => { const fetchAgents = () => { @@ -37,15 +76,79 @@ const BorealisAgentNode = ({ id, data }) => { return () => clearInterval(interval); }, []); + // Ensure host selection stays aligned with available agents + useEffect(() => { + const hostExists = hostOptions.some((opt) => opt.host === selectedHost); + if (hostExists) return; + + if (selectedAgent && agents[selectedAgent]) { + const info = agents[selectedAgent]; + const inferredHost = (info?.hostname || info?.agent_hostname || "").trim() || "unknown"; + if (inferredHost && inferredHost !== selectedHost) { + setSelectedHost(inferredHost); + return; + } + } + + const fallbackHost = hostOptions[0]?.host || ""; + if (fallbackHost !== selectedHost) { + setSelectedHost(fallbackHost); + } + if (!fallbackHost && selectedAgent) { + setSelectedAgent(""); + } + }, [hostOptions, selectedHost, selectedAgent, agents]); + + // Align agent selection with host/mode choice + useEffect(() => { + if (!selectedHost) { + if (selectedMode !== "currentuser") setSelectedMode("currentuser"); + if (selectedAgent) setSelectedAgent(""); + return; + } + const contexts = agentsByHostname[selectedHost]; + if (!contexts) { + if (selectedMode !== "currentuser") setSelectedMode("currentuser"); + if (selectedAgent) setSelectedAgent(""); + return; + } + if (!contexts[selectedMode]) { + const fallbackMode = contexts.currentuser + ? "currentuser" + : contexts.system + ? "system" + : selectedMode; + if (fallbackMode !== selectedMode) { + setSelectedMode(fallbackMode); + return; + } + } + const activeContext = contexts[selectedMode]; + const targetAgentId = activeContext?.agent_id || ""; + if (targetAgentId !== selectedAgent) { + setSelectedAgent(targetAgentId); + } + }, [selectedHost, selectedMode, agentsByHostname, selectedAgent]); + // Sync node data with sidebar changes useEffect(() => { setNodes((nds) => nds.map((n) => - n.id === id ? { ...n, data: { ...n.data, agent_id: selectedAgent } } : n + n.id === id + ? { + ...n, + data: { + ...n.data, + agent_id: selectedAgent, + agent_host: selectedHost, + agent_mode: selectedMode, + }, + } + : n ) ); setIsConnected(false); - }, [selectedAgent, setNodes, id]); + }, [selectedAgent, selectedHost, selectedMode, setNodes, id]); // Attached Roles logic const attachedRoleIds = useMemo( @@ -109,11 +212,19 @@ const BorealisAgentNode = ({ id, data }) => { // Status Label const selectedAgentStatus = useMemo(() => { - if (!selectedAgent) return "Unassigned"; - const agent = agents[selectedAgent]; - if (!agent) return "Reconnecting..."; - return agent.status === "provisioned" ? "Connected" : "Available"; - }, [agents, selectedAgent]); + if (!selectedHost) return "Unassigned"; + const contexts = agentsByHostname[selectedHost]; + if (!contexts) return "Offline"; + const activeContext = contexts[selectedMode]; + if (!selectedAgent || !activeContext) return "Unavailable"; + const status = (activeContext.status || "").toString().toLowerCase(); + if (status === "provisioned") return "Connected"; + if (status === "orphaned") return "Available"; + if (!status) return "Available"; + return status.charAt(0).toUpperCase() + status.slice(1); + }, [agentsByHostname, selectedHost, selectedMode, selectedAgent]); + + const activeHostContexts = selectedHost ? agentsByHostname[selectedHost] : null; // Render (Sidebar handles config) return ( @@ -128,20 +239,44 @@ const BorealisAgentNode = ({ id, data }) => {
Borealis Agent
- + + + + +
+ Target Agent ID:{" "} + {selectedAgent ? ( + {selectedAgent} + ) : ( + none + )} +
+ {isConnected ? (