Broke Apart Monolithic Agent into Linked Modules

This commit is contained in:
2025-09-02 22:43:20 -06:00
parent a1753bd8b1
commit 4fdcd2e3c5
4 changed files with 948 additions and 425 deletions

View File

@@ -48,6 +48,10 @@ except Exception:
pass
from PIL import ImageGrab
# New modularized components
import agent_info
import agent_roles
# //////////////////////////////////////////////////////////////////////////
# CORE SECTION: CONFIG MANAGER
# //////////////////////////////////////////////////////////////////////////
@@ -271,25 +275,17 @@ def detect_agent_os():
print(f"[WARN] OS detection failed: {e}")
return "Unknown"
CONFIG.data['agent_operating_system'] = detect_agent_os()
CONFIG.data['agent_operating_system'] = agent_info.detect_agent_os()
CONFIG._write()
# //////////////////////////////////////////////////////////////////////////
# CORE SECTION: MACRO AUTOMATION
# //////////////////////////////////////////////////////////////////////////
MACRO_ENGINE_PATH = os.path.join(os.path.dirname(__file__), "Python_API_Endpoints", "macro_engines.py")
spec = importlib.util.spec_from_file_location("macro_engines", MACRO_ENGINE_PATH)
macro_engines = importlib.util.module_from_spec(spec)
spec.loader.exec_module(macro_engines)
# //////////////////////////////////////////////////////////////////////////
# CORE SECTION: ASYNC TASK / WEBSOCKET
# //////////////////////////////////////////////////////////////////////////
sio = socketio.AsyncClient(reconnection=True, reconnection_attempts=0, reconnection_delay=5)
role_tasks = {}
overlay_widgets = {}
background_tasks = []
roles_ctx = None
async def stop_all_roles():
print("[DEBUG] Stopping all roles.")
@@ -297,13 +293,11 @@ async def stop_all_roles():
print(f"[DEBUG] Cancelling task for node: {task}")
task.cancel()
role_tasks.clear()
for node_id, widget in overlay_widgets.items():
print(f"[DEBUG] Closing overlay widget: {node_id}")
try:
widget.close()
except Exception as e:
print(f"[WARN] Error closing widget: {e}")
overlay_widgets.clear()
# Close overlays managed in agent_roles module
try:
agent_roles.close_all_overlays()
except Exception:
pass
# ---------------- Heartbeat ----------------
async def send_heartbeat():
@@ -319,7 +313,7 @@ async def send_heartbeat():
payload = {
"agent_id": AGENT_ID,
"hostname": socket.gethostname(),
"agent_operating_system": CONFIG.data.get("agent_operating_system", detect_agent_os()),
"agent_operating_system": CONFIG.data.get("agent_operating_system", agent_info.detect_agent_os()),
"last_seen": int(time.time())
}
await sio.emit("agent_heartbeat", payload)
@@ -329,153 +323,19 @@ async def send_heartbeat():
await asyncio.sleep(60)
# ---------------- Detailed Agent Data ----------------
def _get_internal_ip():
try:
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
s.connect(("8.8.8.8", 80))
ip = s.getsockname()[0]
s.close()
return ip
except Exception:
return "unknown"
## Moved to agent_info module
def collect_summary():
try:
username = getpass.getuser()
domain = os.environ.get("USERDOMAIN") or socket.gethostname()
last_user = f"{domain}\\{username}" if username else "unknown"
except Exception:
last_user = "unknown"
try:
last_reboot = "unknown"
if psutil:
try:
last_reboot = time.strftime(
"%Y-%m-%d %H:%M:%S",
time.localtime(psutil.boot_time()),
)
except Exception:
last_reboot = "unknown"
if last_reboot == "unknown":
plat = platform.system().lower()
if plat == "windows":
try:
out = subprocess.run(
["wmic", "os", "get", "lastbootuptime"],
capture_output=True,
text=True,
timeout=60,
)
raw = "".join(out.stdout.splitlines()[1:]).strip()
if raw:
boot = datetime.datetime.strptime(raw.split(".")[0], "%Y%m%d%H%M%S")
last_reboot = boot.strftime("%Y-%m-%d %H:%M:%S")
except FileNotFoundError:
ps_cmd = "(Get-CimInstance Win32_OperatingSystem).LastBootUpTime"
out = subprocess.run(
["powershell", "-NoProfile", "-Command", ps_cmd],
capture_output=True,
text=True,
timeout=60,
)
raw = out.stdout.strip()
if raw:
try:
boot = datetime.datetime.strptime(raw.split(".")[0], "%Y%m%d%H%M%S")
last_reboot = boot.strftime("%Y-%m-%d %H:%M:%S")
except Exception:
pass
else:
try:
out = subprocess.run(
["uptime", "-s"], capture_output=True, text=True, timeout=30
)
val = out.stdout.strip()
if val:
last_reboot = val
except Exception:
pass
except Exception:
last_reboot = "unknown"
created = CONFIG.data.get("created")
if not created:
created = time.strftime("%Y-%m-%d %H:%M:%S")
CONFIG.data["created"] = created
CONFIG._write()
try:
external_ip = requests.get("https://api.ipify.org", timeout=5).text.strip()
except Exception:
external_ip = "unknown"
return {
"hostname": socket.gethostname(),
"operating_system": CONFIG.data.get("agent_operating_system", detect_agent_os()),
"last_user": last_user,
"internal_ip": _get_internal_ip(),
"external_ip": external_ip,
"last_reboot": last_reboot,
"created": created,
}
# Moved to agent_info.collect_summary
return agent_info.collect_summary(CONFIG)
def collect_software():
items = []
plat = platform.system().lower()
try:
if plat == "windows":
try:
out = subprocess.run(["wmic", "product", "get", "name,version"],
capture_output=True, text=True, timeout=60)
for line in out.stdout.splitlines():
if line.strip() and not line.lower().startswith("name"):
parts = line.strip().split(" ")
name = parts[0].strip()
version = parts[-1].strip() if len(parts) > 1 else ""
if name:
items.append({"name": name, "version": version})
except FileNotFoundError:
ps_cmd = (
"Get-ItemProperty "
"'HKLM:\\Software\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\*',"
"'HKLM:\\Software\\WOW6432Node\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\*' "
"| Where-Object { $_.DisplayName } "
"| Select-Object DisplayName,DisplayVersion "
"| ConvertTo-Json"
)
out = subprocess.run(
["powershell", "-NoProfile", "-Command", ps_cmd],
capture_output=True,
text=True,
timeout=60,
)
data = json.loads(out.stdout or "[]")
if isinstance(data, dict):
data = [data]
for pkg in data:
name = pkg.get("DisplayName")
if name:
items.append({
"name": name,
"version": pkg.get("DisplayVersion", "")
})
elif plat == "linux":
out = subprocess.run(["dpkg-query", "-W", "-f=${Package}\t${Version}\n"], capture_output=True, text=True)
for line in out.stdout.splitlines():
if "\t" in line:
name, version = line.split("\t", 1)
items.append({"name": name, "version": version})
else:
out = subprocess.run([sys.executable, "-m", "pip", "list", "--format", "json"], capture_output=True, text=True)
data = json.loads(out.stdout or "[]")
for pkg in data:
items.append({"name": pkg.get("name"), "version": pkg.get("version")})
except Exception as e:
print(f"[WARN] collect_software failed: {e}")
return items[:100]
# Moved to agent_info.collect_software
return agent_info.collect_software()
def collect_memory():
# Delegated to agent_info module
return agent_info.collect_memory()
entries = []
plat = platform.system().lower()
try:
@@ -564,6 +424,8 @@ def collect_memory():
return entries
def collect_storage():
# Delegated to agent_info module
return agent_info.collect_storage()
disks = []
plat = platform.system().lower()
try:
@@ -680,6 +542,8 @@ def collect_storage():
return disks
def collect_network():
# Delegated to agent_info module
return agent_info.collect_network()
adapters = []
plat = platform.system().lower()
try:
@@ -760,7 +624,7 @@ async def connect():
await sio.emit("agent_heartbeat", {
"agent_id": AGENT_ID,
"hostname": socket.gethostname(),
"agent_operating_system": CONFIG.data.get("agent_operating_system", detect_agent_os()),
"agent_operating_system": CONFIG.data.get("agent_operating_system", agent_info.detect_agent_os()),
"last_seen": int(time.time())
})
except Exception as e:
@@ -796,12 +660,10 @@ async def on_agent_config(cfg):
for rid in removed:
print(f"[DEBUG] Removing node {rid} from regions/overlays.")
CONFIG.data['regions'].pop(rid, None)
w = overlay_widgets.pop(rid, None)
if w:
try:
w.close()
except:
pass
try:
agent_roles.close_overlay(rid)
except Exception:
pass
if removed:
CONFIG._write()
@@ -814,275 +676,21 @@ async def on_agent_config(cfg):
role = role_cfg.get('role')
if role == 'screenshot':
print(f"[DEBUG] Starting screenshot task for {nid}")
task = asyncio.create_task(screenshot_task(role_cfg))
task = asyncio.create_task(agent_roles.screenshot_task(roles_ctx, role_cfg))
role_tasks[nid] = task
elif role == 'macro':
print(f"[DEBUG] Starting macro task for {nid}")
task = asyncio.create_task(macro_task(role_cfg))
task = asyncio.create_task(agent_roles.macro_task(roles_ctx, role_cfg))
role_tasks[nid] = task
@sio.on('list_agent_windows')
async def handle_list_agent_windows(data):
windows = macro_engines.list_windows()
windows = agent_roles.get_window_list()
await sio.emit('agent_window_list', {
'agent_id': AGENT_ID,
'windows': windows
})
# ---------------- Overlay Widget ----------------
overlay_green_thickness = 4
overlay_gray_thickness = 2
handle_size = overlay_green_thickness * 2
extra_top_padding = overlay_green_thickness * 2 + 4 # give space above the top-center green bar
class ScreenshotRegion(QtWidgets.QWidget):
def __init__(self, node_id, x=100, y=100, w=300, h=200, alias=None):
super().__init__()
self.node_id = node_id
self.alias = alias
self.setGeometry(x - handle_size, y - handle_size - extra_top_padding, w + handle_size*2, h + handle_size*2 + extra_top_padding)
self.setWindowFlags(QtCore.Qt.FramelessWindowHint | QtCore.Qt.WindowStaysOnTopHint)
self.setAttribute(QtCore.Qt.WA_TranslucentBackground)
self.resize_dir = None
self.drag_offset = None
self._start_geom = None
self._start_pos = None
self.setMouseTracking(True)
def paintEvent(self, event):
p = QtGui.QPainter(self)
p.setRenderHint(QtGui.QPainter.Antialiasing)
w = self.width()
h = self.height()
# draw gray capture box
p.setPen(QtGui.QPen(QtGui.QColor(130,130,130), overlay_gray_thickness))
p.drawRect(handle_size, handle_size + extra_top_padding, w-handle_size*2, h-handle_size*2 - extra_top_padding)
p.setPen(QtCore.Qt.NoPen)
p.setBrush(QtGui.QBrush(QtGui.QColor(0,191,255)))
edge = overlay_green_thickness*3
# corner handles
p.drawRect(0,extra_top_padding,edge,overlay_green_thickness)
p.drawRect(0,extra_top_padding,overlay_green_thickness,edge)
p.drawRect(w-edge,extra_top_padding,edge,overlay_green_thickness)
p.drawRect(w-overlay_green_thickness,extra_top_padding,overlay_green_thickness,edge)
p.drawRect(0,h-overlay_green_thickness,edge,overlay_green_thickness)
p.drawRect(0,h-edge,overlay_green_thickness,edge)
p.drawRect(w-edge,h-overlay_green_thickness,edge,overlay_green_thickness)
p.drawRect(w-overlay_green_thickness,h-edge,overlay_green_thickness,edge)
# side handles
long = overlay_green_thickness*6
p.drawRect((w-long)//2,extra_top_padding,long,overlay_green_thickness)
p.drawRect((w-long)//2,h-overlay_green_thickness,long,overlay_green_thickness)
p.drawRect(0,(h+extra_top_padding-long)//2,overlay_green_thickness,long)
p.drawRect(w-overlay_green_thickness,(h+extra_top_padding-long)//2,overlay_green_thickness,long)
# draw grabber bar (same size as top-center bar, but above it)
bar_width = overlay_green_thickness * 6
bar_height = overlay_green_thickness
bar_x = (w - bar_width) // 2
bar_y = 6 # 6-8 px down from top
p.setBrush(QtGui.QColor(0,191,255)) # Borealis Blue
p.drawRect(bar_x, bar_y - bar_height - 10, bar_width, bar_height * 4) # 2px padding above green bar
def get_geometry(self):
g = self.geometry()
return (g.x() + handle_size, g.y() + handle_size + extra_top_padding, g.width() - handle_size*2, g.height() - handle_size*2 - extra_top_padding)
def mousePressEvent(self,e):
if e.button()==QtCore.Qt.LeftButton:
pos=e.pos()
bar_width = overlay_green_thickness * 6
bar_height = overlay_green_thickness
bar_x = (self.width() - bar_width) // 2
bar_y = 2
bar_rect = QtCore.QRect(bar_x, bar_y, bar_width, bar_height)
if bar_rect.contains(pos):
self.drag_offset = e.globalPos() - self.frameGeometry().topLeft()
return
x1,y1,self_w,self_h = self.geometry().getRect()
m=handle_size
dirs = []
if pos.x()<=m: dirs.append('left')
if pos.x()>=self.width()-m: dirs.append('right')
if pos.y()<=m+extra_top_padding: dirs.append('top')
if pos.y()>=self.height()-m: dirs.append('bottom')
if dirs:
self.resize_dir = '_'.join(dirs)
self._start_geom = self.geometry()
self._start_pos = e.globalPos()
else:
self.drag_offset = e.globalPos()-self.frameGeometry().topLeft()
def mouseMoveEvent(self,e):
if self.resize_dir and self._start_geom and self._start_pos:
dx = e.globalX() - self._start_pos.x()
dy = e.globalY() - self._start_pos.y()
geom = QtCore.QRect(self._start_geom)
if 'left' in self.resize_dir:
new_x = geom.x() + dx
new_w = geom.width() - dx
geom.setX(new_x)
geom.setWidth(new_w)
if 'right' in self.resize_dir:
geom.setWidth(self._start_geom.width() + dx)
if 'top' in self.resize_dir:
new_y = geom.y() + dy
new_h = geom.height() - dy
geom.setY(new_y)
geom.setHeight(new_h)
if 'bottom' in self.resize_dir:
geom.setHeight(self._start_geom.height() + dy)
self.setGeometry(geom)
elif self.drag_offset and e.buttons() & QtCore.Qt.LeftButton:
self.move(e.globalPos()-self.drag_offset)
def mouseReleaseEvent(self,e):
self.drag_offset=None
self.resize_dir=None
self._start_geom=None
self._start_pos=None
x,y,w,h=self.get_geometry()
CONFIG.data['regions'][self.node_id]={'x':x,'y':y,'w':w,'h':h}
CONFIG._write()
asyncio.create_task(sio.emit('agent_screenshot_task',{ 'agent_id':AGENT_ID,'node_id':self.node_id,'image_base64':'','x':x,'y':y,'w':w,'h':h}))
# ---------------- Screenshot Task ----------------
async def screenshot_task(cfg):
nid=cfg.get('node_id')
alias=cfg.get('alias','')
r=CONFIG.data['regions'].get(nid)
if r:
region=(r['x'],r['y'],r['w'],r['h'])
else:
region=(cfg.get('x',100),cfg.get('y',100),cfg.get('w',300),cfg.get('h',200))
CONFIG.data['regions'][nid]={'x':region[0],'y':region[1],'w':region[2],'h':region[3]}
CONFIG._write()
if nid not in overlay_widgets:
widget=ScreenshotRegion(nid,*region,alias=alias)
overlay_widgets[nid]=widget; widget.show()
await sio.emit('agent_screenshot_task',{'agent_id':AGENT_ID,'node_id':nid,'image_base64':'','x':region[0],'y':region[1],'w':region[2],'h':region[3]})
interval=cfg.get('interval',1000)/1000.0
loop=asyncio.get_event_loop()
executor=concurrent.futures.ThreadPoolExecutor(max_workers=CONFIG.data.get('max_task_workers',8))
try:
while True:
x,y,w,h=overlay_widgets[nid].get_geometry()
grab=partial(ImageGrab.grab,bbox=(x,y,x+w,y+h))
img=await loop.run_in_executor(executor,grab)
buf=BytesIO(); img.save(buf,format='PNG'); encoded=base64.b64encode(buf.getvalue()).decode('utf-8')
await sio.emit('agent_screenshot_task',{'agent_id':AGENT_ID,'node_id':nid,'image_base64':encoded,'x':x,'y':y,'w':w,'h':h})
await asyncio.sleep(interval)
except asyncio.CancelledError:
print(f"[TASK] Screenshot role {nid} cancelled.")
except Exception as e:
print(f"[ERROR] Screenshot task {nid} failed: {e}")
traceback.print_exc()
# ---------------- Macro Task ----------------
async def macro_task(cfg):
"""
Improved macro_task supporting all operation modes, live config, error reporting, and UI feedback.
"""
nid = cfg.get('node_id')
# Track trigger state for edge/level changes
last_trigger_value = 0
has_run_once = False
while True:
# Always re-fetch config (hot reload support)
window_handle = cfg.get('window_handle')
macro_type = cfg.get('macro_type', 'keypress')
operation_mode = cfg.get('operation_mode', 'Continuous')
key = cfg.get('key')
text = cfg.get('text')
interval_ms = int(cfg.get('interval_ms', 1000))
randomize = cfg.get('randomize_interval', False)
random_min = int(cfg.get('random_min', 750))
random_max = int(cfg.get('random_max', 950))
active = cfg.get('active', True)
trigger = int(cfg.get('trigger', 0))
async def emit_macro_status(success, message=""):
await sio.emit('macro_status', {
"agent_id": AGENT_ID,
"node_id": nid,
"success": success,
"message": message,
"timestamp": int(asyncio.get_event_loop().time() * 1000)
})
if not (active is True or str(active).lower() == "true"):
await asyncio.sleep(0.2)
continue
try:
send_macro = False
if operation_mode == "Run Once":
if not has_run_once:
send_macro = True
has_run_once = True
elif operation_mode == "Continuous":
send_macro = True
elif operation_mode == "Trigger-Continuous":
send_macro = (trigger == 1)
elif operation_mode == "Trigger-Once":
if last_trigger_value == 0 and trigger == 1:
send_macro = True
last_trigger_value = trigger
else:
send_macro = True
if send_macro:
if macro_type == 'keypress' and key:
result = macro_engines.send_keypress_to_window(window_handle, key)
elif macro_type == 'typed_text' and text:
result = macro_engines.type_text_to_window(window_handle, text)
else:
await emit_macro_status(False, "Invalid macro type or missing key/text")
await asyncio.sleep(0.2)
continue
if isinstance(result, tuple):
success, err = result
else:
success, err = bool(result), ""
if success:
await emit_macro_status(True, f"Macro sent: {macro_type}")
else:
await emit_macro_status(False, err or "Unknown macro engine failure")
else:
await asyncio.sleep(0.05)
if send_macro:
if randomize:
ms = random.randint(random_min, random_max)
else:
ms = interval_ms
await asyncio.sleep(ms / 1000.0)
else:
await asyncio.sleep(0.1)
except asyncio.CancelledError:
print(f"[TASK] Macro role {nid} cancelled.")
break
except Exception as e:
print(f"[ERROR] Macro task {nid} failed: {e}")
import traceback
traceback.print_exc()
await emit_macro_status(False, str(e))
await asyncio.sleep(0.5)
# ---------------- Config Watcher ----------------
async def config_watcher():
while True:
@@ -1128,13 +736,15 @@ if __name__=='__main__':
app=QtWidgets.QApplication(sys.argv)
loop=QEventLoop(app); asyncio.set_event_loop(loop)
dummy_window=PersistentWindow(); dummy_window.show()
# Initialize roles context for role tasks
roles_ctx = agent_roles.RolesContext(sio=sio, agent_id=AGENT_ID, config=CONFIG)
try:
background_tasks.append(loop.create_task(config_watcher()))
background_tasks.append(loop.create_task(connect_loop()))
background_tasks.append(loop.create_task(idle_task()))
# Start periodic heartbeats
background_tasks.append(loop.create_task(send_heartbeat()))
background_tasks.append(loop.create_task(send_agent_details()))
background_tasks.append(loop.create_task(agent_info.send_agent_details(AGENT_ID, CONFIG)))
loop.run_forever()
except Exception as e:
print(f"[FATAL] Event loop crashed: {e}")