bunny-lab/Borealis
2
0
Fork 0
Code Issues 10 Pull Requests Actions Packages Projects 1 Releases Wiki Activity

Added Reverse Proxy Example to Readme

Browse Source
This commit is contained in:
Nicole Rappe
2025-04-07 21:03:39 -06:00
parent bcf380e1fa
commit 31e899c842
4 changed files with 171 additions and 277 deletions
Download Patch File Download Diff File Expand all files Collapse all files

247
Data/Agent/borealis-agent.py
View File

@ -1,10 +1,9 @@
import sys
import uuid
import time
import json
import base64
import threading
import requests
import socketio
from io import BytesIO
import socket
@ -12,131 +11,56 @@ from PyQt5 import QtCore, QtGui, QtWidgets
from PIL import ImageGrab
# ---------------- Configuration ----------------
SERVER_URL = "https://borealis.bunny-lab.io" # Production URL Example
#SERVER_URL = "http://localhost:5000" # Development URL Example
CHECKIN_ENDPOINT = f"{SERVER_URL}/api/agent/checkin"
CONFIG_ENDPOINT = f"{SERVER_URL}/api/agent/config"
DATA_POST_ENDPOINT = f"{SERVER_URL}/api/agent/data"
HEARTBEAT_ENDPOINT = f"{SERVER_URL}/api/agent/heartbeat"
SERVER_URL = "http://localhost:5000" # WebSocket-enabled server URL
HOSTNAME = socket.gethostname().lower()
RANDOM_SUFFIX = uuid.uuid4().hex[:8]
AGENT_ID = f"{HOSTNAME}-agent-{RANDOM_SUFFIX}"
# Default poll interval for config. Adjust as needed.
CONFIG_POLL_INTERVAL = 5
# ---------------- State ----------------
app_instance = None
region_widget = None
capture_thread_started = False
current_interval = 1000
config_ready = threading.Event()
overlay_visible = True
heartbeat_thread_started = False
# Track if we have a valid connection to Borealis
IS_CONNECTED = False
CONNECTION_LOST_REPORTED = False
# Keep a copy of the last config to avoid repeated provisioning
LAST_CONFIG = {}
# ---------------- Signal Bridge ----------------
class RegionLauncher(QtCore.QObject):
trigger = QtCore.pyqtSignal(int, int, int, int)
# WebSocket client setup
sio = socketio.Client()
def __init__(self):
super().__init__()
self.trigger.connect(self.handle)
# ---------------- WebSocket Handlers ----------------
@sio.event
def connect():
print(f"[WS CONNECTED] 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})
def handle(self, x, y, w, h):
launch_region(x, y, w, h)
@sio.event
def disconnect():
print("[WS DISCONNECTED] Lost connection to Borealis server.")
region_launcher = None
@sio.on('agent_config')
def on_agent_config(config):
global current_interval, overlay_visible, LAST_CONFIG
# ---------------- Helper: Reconnect ----------------
def reconnect():
"""
Attempt to connect to Borealis until successful.
Sets IS_CONNECTED = True upon success.
"""
global IS_CONNECTED, CONNECTION_LOST_REPORTED
while not IS_CONNECTED:
try:
requests.post(CHECKIN_ENDPOINT, json={"agent_id": AGENT_ID, "hostname": HOSTNAME}, timeout=5)
IS_CONNECTED = True
CONNECTION_LOST_REPORTED = False
print(f"[INFO] Agent ID: {AGENT_ID} connected to Borealis.")
except Exception:
if not CONNECTION_LOST_REPORTED:
print(f"[CONNECTION LOST] Attempting to Reconnect to Borealis Server at {SERVER_URL}")
CONNECTION_LOST_REPORTED = True
time.sleep(10)
if config != LAST_CONFIG:
print("[PROVISIONING] Received new configuration from Borealis.")
x = config.get("x", 100)
y = config.get("y", 100)
w = config.get("w", 300)
h = config.get("h", 200)
current_interval = config.get("interval", 1000)
overlay_visible = config.get("visible", True)
# ---------------- Networking ----------------
def poll_for_config():
"""
Polls for agent configuration from Borealis.
Returns a config dict or None on failure.
"""
try:
res = requests.get(CONFIG_ENDPOINT, params={"agent_id": AGENT_ID}, timeout=5)
if res.status_code == 200:
return res.json()
if not region_widget:
region_launcher.trigger.emit(x, y, w, h)
else:
print(f"[ERROR] Config polling returned status: {res.status_code}")
except Exception:
# We'll let the config_loop handle setting IS_CONNECTED = False
pass
return None
region_widget.setGeometry(x, y, w, h)
region_widget.setVisible(overlay_visible)
def send_image_data(image):
"""
Attempts to POST screenshot data to Borealis if IS_CONNECTED is True.
"""
global IS_CONNECTED, CONNECTION_LOST_REPORTED
if not IS_CONNECTED:
return # Skip sending if not connected
try:
buffer = BytesIO()
image.save(buffer, format="PNG")
encoded = base64.b64encode(buffer.getvalue()).decode("utf-8")
response = requests.post(DATA_POST_ENDPOINT, json={
"agent_id": AGENT_ID,
"type": "screenshot",
"image_base64": encoded
}, timeout=5)
if response.status_code != 200:
print(f"[ERROR] Screenshot POST failed: {response.status_code} - {response.text}")
except Exception as e:
if IS_CONNECTED and not CONNECTION_LOST_REPORTED:
# Only report once
print(f"[CONNECTION LOST] Attempting to Reconnect to Borealis Server at {SERVER_URL}")
CONNECTION_LOST_REPORTED = True
IS_CONNECTED = False
def send_heartbeat():
"""
Attempts to send heartbeat if IS_CONNECTED is True.
"""
global IS_CONNECTED, CONNECTION_LOST_REPORTED
if not IS_CONNECTED:
return
try:
response = requests.get(HEARTBEAT_ENDPOINT, params={"agent_id": AGENT_ID}, timeout=5)
if response.status_code != 200:
print(f"[ERROR] Heartbeat returned status: {response.status_code}")
raise ValueError("Heartbeat not 200")
except Exception:
if IS_CONNECTED and not CONNECTION_LOST_REPORTED:
print(f"[CONNECTION LOST] Attempting to Reconnect to Borealis Server at {SERVER_URL}")
CONNECTION_LOST_REPORTED = True
IS_CONNECTED = False
LAST_CONFIG = config
config_ready.set()
# ---------------- Region Overlay ----------------
class ScreenshotRegion(QtWidgets.QWidget):
@ -161,12 +85,10 @@ class ScreenshotRegion(QtWidgets.QWidget):
painter = QtGui.QPainter(self)
painter.setRenderHint(QtGui.QPainter.Antialiasing)
# Transparent fill
painter.setBrush(QtCore.Qt.transparent)
painter.setPen(QtGui.QPen(QtGui.QColor(0, 255, 0), 2))
painter.drawRect(self.rect())
# Resize Handle Visual (Bottom-Right)
handle_rect = QtCore.QRect(
self.width() - self.resize_handle_size,
self.height() - self.resize_handle_size,
@ -199,112 +121,49 @@ class ScreenshotRegion(QtWidgets.QWidget):
geo = self.geometry()
return geo.x(), geo.y(), geo.width(), geo.height()
# ---------------- Threads ----------------
# ---------------- Screenshot Capture ----------------
def capture_loop():
"""
Continuously captures the user-defined region every current_interval ms if connected.
"""
global current_interval
print("[INFO] Screenshot capture loop started")
config_ready.wait()
while region_widget is None:
print("[WAIT] Waiting for region widget to initialize...")
time.sleep(0.2)
print(f"[INFO] Agent Capturing Region: x:{region_widget.x()} y:{region_widget.y()} w:{region_widget.width()} h:{region_widget.height()}")
while True:
if overlay_visible and IS_CONNECTED:
if overlay_visible:
x, y, w, h = region_widget.get_geometry()
try:
img = ImageGrab.grab(bbox=(x, y, x + w, y + h))
send_image_data(img)
buffer = BytesIO()
img.save(buffer, format="PNG")
encoded_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
sio.emit('screenshot', {
'agent_id': AGENT_ID,
'image_base64': encoded_image
})
except Exception as e:
print(f"[ERROR] Screenshot error: {e}")
time.sleep(current_interval / 1000)
def heartbeat_loop():
"""
Heartbeat every 10 seconds if connected.
If it fails, we set IS_CONNECTED=False, and rely on config_loop to reconnect.
"""
while True:
send_heartbeat()
time.sleep(10)
# ---------------- UI Launcher ----------------
class RegionLauncher(QtCore.QObject):
trigger = QtCore.pyqtSignal(int, int, int, int)
def config_loop():
"""
1) Reconnect (if needed) until the agent can contact Borealis
2) Poll for config. If new config is different from LAST_CONFIG, re-provision
3) If poll_for_config fails or we see connection issues, set IS_CONNECTED=False
and loop back to reconnect() on next iteration
"""
global capture_thread_started, heartbeat_thread_started
global current_interval, overlay_visible, LAST_CONFIG, IS_CONNECTED
def __init__(self):
super().__init__()
self.trigger.connect(self.handle)
while True:
# If we aren't connected, reconnect
if not IS_CONNECTED:
reconnect()
def handle(self, x, y, w, h):
launch_region(x, y, w, h)
# Attempt to get config
config = poll_for_config()
if config is None:
# This means we had a poll failure, so mark disconnected and retry.
IS_CONNECTED = False
continue
# If it has a "task" : "screenshot"
if config.get("task") == "screenshot":
# Compare to last known config
if config != LAST_CONFIG:
# Something changed, so provision
print("[PROVISIONING] Agent Provisioning Command Issued by Borealis")
x = config.get("x", 100)
y = config.get("y", 100)
w = config.get("w", 300)
h = config.get("h", 200)
current_interval = config.get("interval", 1000)
overlay_visible = config.get("visible", True)
print('[PROVISIONING] Agent Configured as "Screenshot" Collector')
print(f'[PROVISIONING] Polling Rate: {current_interval / 1000:.1f}s')
# Show or move region widget
if not region_widget:
region_launcher.trigger.emit(x, y, w, h)
else:
region_widget.setGeometry(x, y, w, h)
region_widget.setVisible(overlay_visible)
LAST_CONFIG = config
# Make sure capture thread is started
if not capture_thread_started:
threading.Thread(target=capture_loop, daemon=True).start()
capture_thread_started = True
# Make sure heartbeat thread is started
if not heartbeat_thread_started:
threading.Thread(target=heartbeat_loop, daemon=True).start()
heartbeat_thread_started = True
# Signal that provisioning is done so capture thread can run
config_ready.set()
# Sleep before next poll
time.sleep(CONFIG_POLL_INTERVAL)
region_launcher = None
def launch_region(x, y, w, h):
"""
Initializes the screenshot region overlay widget exactly once.
"""
global region_widget
if region_widget:
return
print("[INFO] Agent Starting...")
region_widget = ScreenshotRegion(x, y, w, h)
region_widget.show()
@ -313,8 +172,8 @@ if __name__ == "__main__":
app_instance = QtWidgets.QApplication(sys.argv)
region_launcher = RegionLauncher()
# Start the config loop in a background thread
threading.Thread(target=config_loop, daemon=True).start()
sio.connect(SERVER_URL)
threading.Thread(target=capture_loop, daemon=True).start()
# Enter Qt main event loop
sys.exit(app_instance.exec_())