567 lines
20 KiB
Python
567 lines
20 KiB
Python
"""
|
|
EasyBot Windows Agent — tool-based remote automation client.
|
|
|
|
The PC exposes granular tools (run_command, get_ui_tree, click, type, etc.)
|
|
and the EasyBot LLM agent controls everything directly. No LLM on client side.
|
|
|
|
Configuration via environment variables:
|
|
EASYBOT_SERVER WebSocket URL (e.g., wss://app.example.com/ws/windows-agent)
|
|
EASYBOT_TOKEN Authentication token (eb_wag_...)
|
|
LOG_LEVEL DEBUG | INFO | WARNING (default: INFO)
|
|
"""
|
|
|
|
import asyncio
|
|
import base64
|
|
import io
|
|
import json
|
|
import logging
|
|
import os
|
|
import platform
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from typing import Any
|
|
|
|
import websockets
|
|
from websockets.client import WebSocketClientProtocol
|
|
|
|
from . import __version__
|
|
|
|
logger = logging.getLogger("easybot-windows-agent")
|
|
|
|
# ── Configuration ────────────────────────────────────────────────────────────
|
|
|
|
RECONNECT_DELAY_BASE = 5
|
|
RECONNECT_DELAY_MAX = 120
|
|
TOOL_TIMEOUT = 60 # max seconds for a single tool call
|
|
|
|
|
|
# ── Tool Handlers ────────────────────────────────────────────────────────────
|
|
|
|
async def tool_run_command(args: dict[str, Any]) -> dict[str, Any]:
|
|
"""Execute a shell command and return output."""
|
|
command = args.get("command", "")
|
|
shell = args.get("shell", "powershell")
|
|
|
|
if not command:
|
|
return {"error": "No command provided"}
|
|
|
|
try:
|
|
if shell == "cmd":
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"cmd", "/c", command,
|
|
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
|
creationflags=subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0,
|
|
)
|
|
else:
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"powershell", "-NoProfile", "-NonInteractive", "-Command", command,
|
|
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
|
creationflags=subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0,
|
|
)
|
|
|
|
stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=TOOL_TIMEOUT)
|
|
return {
|
|
"stdout": stdout.decode("utf-8", errors="replace").strip(),
|
|
"stderr": stderr.decode("utf-8", errors="replace").strip(),
|
|
"exit_code": proc.returncode,
|
|
}
|
|
except asyncio.TimeoutError:
|
|
return {"error": "Command timed out", "exit_code": -1}
|
|
except Exception as e:
|
|
return {"error": str(e)}
|
|
|
|
|
|
async def tool_get_ui_tree(args: dict[str, Any]) -> dict[str, Any]:
|
|
"""Get the accessibility tree of a window or the desktop."""
|
|
import uiautomation as auto
|
|
|
|
window_title = args.get("window")
|
|
max_depth = args.get("depth", 3)
|
|
|
|
try:
|
|
if window_title:
|
|
root = auto.WindowControl(searchDepth=1, Name=window_title)
|
|
if not root.Exists(maxSearchSeconds=3):
|
|
# Try substring match
|
|
root = auto.WindowControl(searchDepth=1, SubName=window_title)
|
|
if not root.Exists(maxSearchSeconds=2):
|
|
return {"error": f"Window '{window_title}' not found"}
|
|
else:
|
|
root = auto.GetForegroundControl()
|
|
if not root:
|
|
root = auto.GetRootControl()
|
|
|
|
def walk(ctrl: Any, depth: int, max_d: int) -> list[dict]:
|
|
if depth > max_d:
|
|
return []
|
|
items = []
|
|
try:
|
|
name = ctrl.Name or ""
|
|
ctrl_type = ctrl.ControlTypeName or ""
|
|
auto_id = ctrl.AutomationId or ""
|
|
rect = ctrl.BoundingRectangle
|
|
value = ""
|
|
try:
|
|
value = ctrl.GetValuePattern().Value if ctrl.GetValuePattern() else ""
|
|
except Exception:
|
|
pass
|
|
|
|
item: dict[str, Any] = {
|
|
"type": ctrl_type,
|
|
"name": name,
|
|
}
|
|
if auto_id:
|
|
item["automation_id"] = auto_id
|
|
if value:
|
|
item["value"] = value
|
|
if rect and rect.width() > 0:
|
|
item["rect"] = {"x": rect.left, "y": rect.top, "w": rect.width(), "h": rect.height()}
|
|
|
|
children = []
|
|
if depth < max_d:
|
|
for child in ctrl.GetChildren():
|
|
children.extend(walk(child, depth + 1, max_d))
|
|
if children:
|
|
item["children"] = children
|
|
|
|
items.append(item)
|
|
except Exception:
|
|
pass
|
|
return items
|
|
|
|
tree = walk(root, 0, max_depth)
|
|
return {"window": root.Name or "Desktop", "tree": tree}
|
|
except Exception as e:
|
|
return {"error": str(e)}
|
|
|
|
|
|
async def tool_list_windows(args: dict[str, Any]) -> dict[str, Any]:
|
|
"""List all open top-level windows."""
|
|
import uiautomation as auto
|
|
|
|
try:
|
|
windows = []
|
|
for win in auto.GetRootControl().GetChildren():
|
|
try:
|
|
if win.ControlTypeName == "WindowControl" or win.ClassName == "Window":
|
|
name = win.Name or ""
|
|
if name and name != "Program Manager":
|
|
windows.append({
|
|
"title": name,
|
|
"class": win.ClassName or "",
|
|
"process_id": win.ProcessId,
|
|
})
|
|
except Exception:
|
|
continue
|
|
return {"windows": windows}
|
|
except Exception as e:
|
|
return {"error": str(e)}
|
|
|
|
|
|
async def tool_click_element(args: dict[str, Any]) -> dict[str, Any]:
|
|
"""Click a UI element by name, automation_id, or coordinates."""
|
|
name = args.get("name")
|
|
auto_id = args.get("automation_id")
|
|
x = args.get("x")
|
|
y = args.get("y")
|
|
button = args.get("button", "left")
|
|
double = args.get("double", False)
|
|
|
|
try:
|
|
if x is not None and y is not None:
|
|
import pyautogui
|
|
if double:
|
|
pyautogui.doubleClick(x, y)
|
|
elif button == "right":
|
|
pyautogui.rightClick(x, y)
|
|
else:
|
|
pyautogui.click(x, y)
|
|
return {"success": True, "method": "coordinates", "x": x, "y": y}
|
|
|
|
import uiautomation as auto
|
|
ctrl = None
|
|
if auto_id:
|
|
ctrl = auto.Control(searchDepth=8, AutomationId=auto_id)
|
|
elif name:
|
|
ctrl = auto.Control(searchDepth=8, Name=name)
|
|
|
|
if ctrl and ctrl.Exists(maxSearchSeconds=3):
|
|
if double:
|
|
ctrl.DoubleClick()
|
|
else:
|
|
ctrl.Click()
|
|
return {"success": True, "method": "ui_automation", "element": ctrl.Name}
|
|
else:
|
|
return {"error": f"Element not found: name={name}, automation_id={auto_id}"}
|
|
except Exception as e:
|
|
return {"error": str(e)}
|
|
|
|
|
|
async def tool_type_text(args: dict[str, Any]) -> dict[str, Any]:
|
|
"""Type text, optionally into a specific element."""
|
|
text = args.get("text", "")
|
|
element_name = args.get("element_name")
|
|
|
|
if not text:
|
|
return {"error": "No text provided"}
|
|
|
|
try:
|
|
if element_name:
|
|
import uiautomation as auto
|
|
ctrl = auto.Control(searchDepth=8, Name=element_name)
|
|
if ctrl.Exists(maxSearchSeconds=3):
|
|
try:
|
|
ctrl.GetValuePattern().SetValue(text)
|
|
return {"success": True, "method": "set_value", "element": element_name}
|
|
except Exception:
|
|
ctrl.Click()
|
|
await asyncio.sleep(0.2)
|
|
|
|
import pyautogui
|
|
pyautogui.typewrite(text, interval=0.02) if text.isascii() else pyautogui.write(text)
|
|
return {"success": True, "method": "keyboard"}
|
|
except Exception as e:
|
|
return {"error": str(e)}
|
|
|
|
|
|
async def tool_press_keys(args: dict[str, Any]) -> dict[str, Any]:
|
|
"""Press keyboard keys/shortcuts (e.g. 'ctrl+s', 'enter', 'alt+f4')."""
|
|
keys = args.get("keys", "")
|
|
if not keys:
|
|
return {"error": "No keys provided"}
|
|
|
|
try:
|
|
import pyautogui
|
|
parts = [k.strip() for k in keys.lower().split("+")]
|
|
if len(parts) > 1:
|
|
pyautogui.hotkey(*parts)
|
|
else:
|
|
pyautogui.press(parts[0])
|
|
return {"success": True, "keys": keys}
|
|
except Exception as e:
|
|
return {"error": str(e)}
|
|
|
|
|
|
async def tool_screenshot(args: dict[str, Any]) -> dict[str, Any]:
|
|
"""Take a screenshot and return as base64 PNG."""
|
|
from PIL import ImageGrab
|
|
|
|
try:
|
|
region = args.get("region")
|
|
if region:
|
|
bbox = (region["x"], region["y"], region["x"] + region["w"], region["y"] + region["h"])
|
|
img = ImageGrab.grab(bbox=bbox)
|
|
else:
|
|
img = ImageGrab.grab()
|
|
|
|
# Resize for efficiency (max 1280px wide)
|
|
if img.width > 1280:
|
|
ratio = 1280 / img.width
|
|
img = img.resize((1280, int(img.height * ratio)))
|
|
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="PNG", optimize=True)
|
|
b64 = base64.b64encode(buf.getvalue()).decode("ascii")
|
|
return {"image_base64": b64, "width": img.width, "height": img.height}
|
|
except Exception as e:
|
|
return {"error": str(e)}
|
|
|
|
|
|
async def tool_read_clipboard(args: dict[str, Any]) -> dict[str, Any]:
|
|
"""Read text from the clipboard."""
|
|
try:
|
|
import pyperclip
|
|
text = pyperclip.paste()
|
|
return {"text": text}
|
|
except ImportError:
|
|
# Fallback via PowerShell
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"powershell", "-NoProfile", "-Command", "Get-Clipboard",
|
|
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
|
)
|
|
stdout, _ = await proc.communicate()
|
|
return {"text": stdout.decode("utf-8", errors="replace").strip()}
|
|
except Exception as e:
|
|
return {"error": str(e)}
|
|
|
|
|
|
async def tool_set_clipboard(args: dict[str, Any]) -> dict[str, Any]:
|
|
"""Set clipboard text content."""
|
|
text = args.get("text", "")
|
|
try:
|
|
import pyperclip
|
|
pyperclip.copy(text)
|
|
return {"success": True}
|
|
except ImportError:
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"powershell", "-NoProfile", "-Command", f"Set-Clipboard -Value '{text}'",
|
|
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
|
)
|
|
await proc.communicate()
|
|
return {"success": True}
|
|
except Exception as e:
|
|
return {"error": str(e)}
|
|
|
|
|
|
async def tool_focus_window(args: dict[str, Any]) -> dict[str, Any]:
|
|
"""Bring a window to the foreground."""
|
|
title = args.get("title", "")
|
|
if not title:
|
|
return {"error": "No window title provided"}
|
|
|
|
try:
|
|
import uiautomation as auto
|
|
win = auto.WindowControl(searchDepth=1, SubName=title)
|
|
if win.Exists(maxSearchSeconds=3):
|
|
win.SetFocus()
|
|
win.SetActive()
|
|
return {"success": True, "window": win.Name}
|
|
else:
|
|
return {"error": f"Window '{title}' not found"}
|
|
except Exception as e:
|
|
return {"error": str(e)}
|
|
|
|
|
|
async def tool_file_read(args: dict[str, Any]) -> dict[str, Any]:
|
|
"""Read a text file."""
|
|
path = args.get("path", "")
|
|
if not path:
|
|
return {"error": "No path provided"}
|
|
try:
|
|
with open(path, "r", encoding="utf-8", errors="replace") as f:
|
|
content = f.read(100_000) # limit to 100KB
|
|
return {"content": content, "size": os.path.getsize(path)}
|
|
except Exception as e:
|
|
return {"error": str(e)}
|
|
|
|
|
|
async def tool_file_write(args: dict[str, Any]) -> dict[str, Any]:
|
|
"""Write content to a text file."""
|
|
path = args.get("path", "")
|
|
content = args.get("content", "")
|
|
if not path:
|
|
return {"error": "No path provided"}
|
|
try:
|
|
os.makedirs(os.path.dirname(path), exist_ok=True) if os.path.dirname(path) else None
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
f.write(content)
|
|
return {"success": True, "bytes_written": len(content.encode("utf-8"))}
|
|
except Exception as e:
|
|
return {"error": str(e)}
|
|
|
|
|
|
# ── Tool Registry ────────────────────────────────────────────────────────────
|
|
|
|
TOOLS: dict[str, Any] = {
|
|
"run_command": tool_run_command,
|
|
"get_ui_tree": tool_get_ui_tree,
|
|
"list_windows": tool_list_windows,
|
|
"click_element": tool_click_element,
|
|
"type_text": tool_type_text,
|
|
"press_keys": tool_press_keys,
|
|
"screenshot": tool_screenshot,
|
|
"read_clipboard": tool_read_clipboard,
|
|
"set_clipboard": tool_set_clipboard,
|
|
"focus_window": tool_focus_window,
|
|
"file_read": tool_file_read,
|
|
"file_write": tool_file_write,
|
|
}
|
|
|
|
|
|
# ── Agent Client ─────────────────────────────────────────────────────────────
|
|
|
|
class AgentClient:
|
|
def __init__(self, server_url: str, token: str):
|
|
self.server_url = server_url
|
|
self.token = token
|
|
self.ws: WebSocketClientProtocol | None = None
|
|
self.running = True
|
|
self.current_call_id: str | None = None
|
|
self.cancel_event = asyncio.Event()
|
|
self._reconnect_delay = RECONNECT_DELAY_BASE
|
|
|
|
async def run(self) -> None:
|
|
"""Main loop: connect → handle messages → reconnect on failure."""
|
|
while self.running:
|
|
try:
|
|
await self._connect_and_handle()
|
|
except (websockets.ConnectionClosed, OSError, ConnectionRefusedError) as e:
|
|
logger.warning("Connection lost: %s — reconnecting in %ds", e, self._reconnect_delay)
|
|
except Exception:
|
|
logger.exception("Unexpected error — reconnecting in %ds", self._reconnect_delay)
|
|
|
|
if not self.running:
|
|
break
|
|
await asyncio.sleep(self._reconnect_delay)
|
|
self._reconnect_delay = min(self._reconnect_delay * 2, RECONNECT_DELAY_MAX)
|
|
|
|
async def _connect_and_handle(self) -> None:
|
|
"""Single connection lifecycle."""
|
|
extra_headers = {"Authorization": f"Bearer {self.token}"}
|
|
|
|
async with websockets.connect(
|
|
self.server_url,
|
|
additional_headers=extra_headers,
|
|
ping_interval=None,
|
|
max_size=16 * 1024 * 1024,
|
|
close_timeout=10,
|
|
) as ws:
|
|
self.ws = ws
|
|
self._reconnect_delay = RECONNECT_DELAY_BASE
|
|
logger.info("Connected to %s", self.server_url)
|
|
|
|
# Send client info with list of available tools
|
|
await self._send({
|
|
"type": "client_info",
|
|
"hostname": platform.node(),
|
|
"os_version": f"{platform.system()} {platform.version()}",
|
|
"python_version": platform.python_version(),
|
|
"client_version": __version__,
|
|
"tools": list(TOOLS.keys()),
|
|
})
|
|
|
|
try:
|
|
async for raw in ws:
|
|
msg = json.loads(raw)
|
|
await self._handle_message(msg)
|
|
finally:
|
|
self.ws = None
|
|
self.current_call_id = None
|
|
|
|
async def _handle_message(self, msg: dict[str, Any]) -> None:
|
|
"""Dispatch incoming messages from the server."""
|
|
msg_type = msg.get("type")
|
|
|
|
if msg_type == "ping":
|
|
await self._send({"type": "pong"})
|
|
|
|
elif msg_type == "tool_call":
|
|
call_id = msg["call_id"]
|
|
tool_name = msg["tool_name"]
|
|
arguments = msg.get("arguments", {})
|
|
logger.info("Tool call: %s (call_id=%s)", tool_name, call_id[:8])
|
|
asyncio.create_task(self._execute_tool(call_id, tool_name, arguments))
|
|
|
|
elif msg_type == "cancel":
|
|
call_id = msg.get("call_id")
|
|
if self.current_call_id == call_id:
|
|
logger.info("Cancel requested for %s", call_id)
|
|
self.cancel_event.set()
|
|
|
|
else:
|
|
logger.debug("Unknown message type: %s", msg_type)
|
|
|
|
async def _execute_tool(self, call_id: str, tool_name: str, arguments: dict[str, Any]) -> None:
|
|
"""Execute a tool and return the result."""
|
|
self.current_call_id = call_id
|
|
self.cancel_event.clear()
|
|
start = time.time()
|
|
|
|
try:
|
|
handler = TOOLS.get(tool_name)
|
|
if not handler:
|
|
await self._send({
|
|
"type": "tool_result",
|
|
"call_id": call_id,
|
|
"error": f"Unknown tool: {tool_name}",
|
|
"duration_ms": 0,
|
|
})
|
|
return
|
|
|
|
result = await asyncio.wait_for(handler(arguments), timeout=TOOL_TIMEOUT)
|
|
duration_ms = int((time.time() - start) * 1000)
|
|
|
|
if "error" in result:
|
|
await self._send({
|
|
"type": "tool_result",
|
|
"call_id": call_id,
|
|
"error": result["error"],
|
|
"duration_ms": duration_ms,
|
|
})
|
|
else:
|
|
await self._send({
|
|
"type": "tool_result",
|
|
"call_id": call_id,
|
|
"result": result,
|
|
"duration_ms": duration_ms,
|
|
})
|
|
|
|
logger.info("Tool %s completed in %dms", tool_name, duration_ms)
|
|
|
|
except asyncio.TimeoutError:
|
|
await self._send({
|
|
"type": "tool_result",
|
|
"call_id": call_id,
|
|
"error": f"Tool '{tool_name}' timed out after {TOOL_TIMEOUT}s",
|
|
"duration_ms": int((time.time() - start) * 1000),
|
|
})
|
|
except Exception as e:
|
|
logger.error("Tool %s failed: %s", tool_name, e)
|
|
await self._send({
|
|
"type": "tool_result",
|
|
"call_id": call_id,
|
|
"error": f"{type(e).__name__}: {e}",
|
|
"duration_ms": int((time.time() - start) * 1000),
|
|
})
|
|
finally:
|
|
self.current_call_id = None
|
|
|
|
async def _send(self, msg: dict[str, Any]) -> None:
|
|
"""Send a JSON message to the server."""
|
|
if self.ws:
|
|
try:
|
|
await self.ws.send(json.dumps(msg))
|
|
except (websockets.ConnectionClosed, Exception):
|
|
pass
|
|
|
|
def stop(self) -> None:
|
|
"""Signal the client to stop."""
|
|
self.running = False
|
|
self.cancel_event.set()
|
|
|
|
|
|
# ── Entry Point ──────────────────────────────────────────────────────────────
|
|
|
|
def main() -> None:
|
|
server_url = os.environ.get("EASYBOT_SERVER")
|
|
token = os.environ.get("EASYBOT_TOKEN")
|
|
|
|
if not server_url or not token:
|
|
print("Error: EASYBOT_SERVER and EASYBOT_TOKEN environment variables are required.")
|
|
print()
|
|
print("Usage:")
|
|
print(" set EASYBOT_SERVER=wss://app.example.com/ws/windows-agent")
|
|
print(" set EASYBOT_TOKEN=eb_wag_...")
|
|
print(" easybot-windows-agent")
|
|
sys.exit(1)
|
|
|
|
log_level = os.environ.get("LOG_LEVEL", "INFO").upper()
|
|
logging.basicConfig(
|
|
level=getattr(logging, log_level, logging.INFO),
|
|
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
|
datefmt="%Y-%m-%d %H:%M:%S",
|
|
)
|
|
|
|
logger.info("EasyBot Windows Agent v%s starting", __version__)
|
|
logger.info("Server: %s", server_url)
|
|
logger.info("Token: %s...", token[:12] if len(token) > 12 else "***")
|
|
logger.info("Available tools: %s", ", ".join(TOOLS.keys()))
|
|
|
|
client = AgentClient(server_url, token)
|
|
|
|
def handle_signal(*_: Any) -> None:
|
|
logger.info("Shutdown signal received")
|
|
client.stop()
|
|
|
|
signal.signal(signal.SIGINT, handle_signal)
|
|
signal.signal(signal.SIGTERM, handle_signal)
|
|
|
|
asyncio.run(client.run())
|
|
logger.info("Agent stopped")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|