Borealis-Github-Replica/Data/Agent/api-collector-agent.py

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 = "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"

HOSTNAME = socket.gethostname().lower()
RANDOM_SUFFIX = uuid.uuid4().hex[:8]
AGENT_ID = f"{HOSTNAME}-agent-{RANDOM_SUFFIX}"

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

# ---------------- 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

# ---------------- Agent Networking ----------------
def check_in():
    try:
        requests.post(CHECKIN_ENDPOINT, json={"agent_id": AGENT_ID, "hostname": HOSTNAME})
        print(f"[INFO] Agent ID: {AGENT_ID}")
    except Exception as e:
        print(f"[ERROR] Check-in failed: {e}")

def poll_for_config():
    try:
        res = requests.get(CONFIG_ENDPOINT, params={"agent_id": AGENT_ID})
        if res.status_code == 200:
            return res.json()
    except Exception as e:
        print(f"[ERROR] Config polling failed: {e}")
    return None

def send_image_data(image):
    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
        })

        if response.status_code != 200:
            print(f"[ERROR] Screenshot POST failed: {response.status_code} - {response.text}")
    except Exception as e:
        print(f"[ERROR] Failed to send image: {e}")

# ---------------- 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():
    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:
            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 config_loop():
    global region_widget, capture_thread_started, current_interval, overlay_visible
    check_in()
    while True:
        config = poll_for_config()
        if config and config.get("task") == "screenshot":
            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(f"[PROVISIONING] Agent Configured as \"Screenshot\" Collector w/ Polling Rate of <{current_interval/1000:.1f}s>")

            if not region_widget:
                region_launcher.trigger.emit(x, y, w, h)
            elif region_widget:
                region_widget.setVisible(overlay_visible)

            if not capture_thread_started:
                threading.Thread(target=capture_loop, daemon=True).start()
                capture_thread_started = True

            config_ready.set()
        time.sleep(CONFIG_POLL_INTERVAL)

def launch_region(x, y, w, h):
    global region_widget
    if region_widget:
        return
    print(f"[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()
    threading.Thread(target=config_loop, daemon=True).start()
    sys.exit(app_instance.exec_())