321 lines
11 KiB
Python
321 lines
11 KiB
Python
import sys
|
|
import uuid
|
|
import time
|
|
import json
|
|
import base64
|
|
import threading
|
|
import requests
|
|
from io import BytesIO
|
|
import socket
|
|
|
|
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"
|
|
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"
|
|
|
|
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)
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.trigger.connect(self.handle)
|
|
|
|
def handle(self, x, y, w, h):
|
|
launch_region(x, y, w, h)
|
|
|
|
region_launcher = None
|
|
|
|
# ---------------- 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)
|
|
|
|
# ---------------- 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()
|
|
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
|
|
|
|
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
|
|
|
|
# ---------------- Region Overlay ----------------
|
|
class ScreenshotRegion(QtWidgets.QWidget):
|
|
def __init__(self, x=100, y=100, w=300, h=200):
|
|
super().__init__()
|
|
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(AGENT_ID)
|
|
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)
|
|
|
|
# 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,
|
|
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()
|
|
|
|
# ---------------- Threads ----------------
|
|
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:
|
|
x, y, w, h = region_widget.get_geometry()
|
|
try:
|
|
img = ImageGrab.grab(bbox=(x, y, x + w, y + h))
|
|
send_image_data(img)
|
|
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)
|
|
|
|
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
|
|
|
|
while True:
|
|
# If we aren't connected, reconnect
|
|
if not IS_CONNECTED:
|
|
reconnect()
|
|
|
|
# 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)
|
|
|
|
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()
|
|
|
|
# ---------------- Main ----------------
|
|
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()
|
|
|
|
# Enter Qt main event loop
|
|
sys.exit(app_instance.exec_())
|