kaushik/rmm-backend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: implement async command execution, improve client status tracking, and expand takeleap diagnostic commands.

Browse Source
This commit is contained in:
kaushik
2026-06-04 15:48:38 +05:30
parent c4a96adfc1
commit 25e9b5e120
3 changed files with 199 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
central_api_prototype.py
View File

@@ -123,8 +123,25 @@ def append_to_logs(log_type: str, client_id: str, detail: Any):
except Exception as e:
print(f"[!] Error writing log to MongoDB: {e}")
def mark_client_active(client_id: str, clients: Dict[str, Any]):
if client_id in clients:
was_active = clients[client_id].get("active", False)
if not was_active:
if client_id in last_down_notification_times:
send_rocketchat_notification(
text=f"✅ **System UP:** Client `{client_id}` has recovered and is back online.",
color="#2ecc71"
)
last_down_notification_times.pop(client_id, None)
clients[client_id]["active"] = True
clients[client_id]["last_seen"] = datetime.now().isoformat()
def load_clients() -> Dict[str, Any]:
try:
# Clean up any ghost clients with empty or whitespace-only keys
clients_collection.delete_many({"_id": {"$in": ["", None]}})
clients_collection.delete_many({"_id": {"$regex": "^\\s*$"}})
# If collection is empty, seed it with default registry
if clients_collection.count_documents({}) == 0:
default_data = {
@@ -248,6 +265,9 @@ async def get_command(client_id: str):
Endpoint for clients to poll for pending commands.
Client Agent hits this endpoint to ask: "Do I have any work to do?"
"""
if not client_id or not client_id.strip():
raise HTTPException(status_code=400, detail="client_id cannot be empty")
clients = load_clients()
# 1. Auto-register new client if we haven't seen it yet
@@ -264,8 +284,7 @@ async def get_command(client_id: str):
# 3. Reset the command to "none", mark active, and update last_seen timestamp
clients[client_id]["pending_command"] = "none"
clients[client_id]["active"] = True
clients[client_id]["last_seen"] = datetime.now().isoformat()
mark_client_active(client_id, clients)
save_clients(clients)
return {"command": cmd}
@@ -297,6 +316,9 @@ async def receive_telemetry(client_id: str, payload: TelemetryPayload):
"""
Endpoint for agents to push their live hardware telemetry.
"""
if not client_id or not client_id.strip():
raise HTTPException(status_code=400, detail="client_id cannot be empty")
clients = load_clients()
# Auto-register new client if we see it through telemetry first
@@ -305,8 +327,7 @@ async def receive_telemetry(client_id: str, payload: TelemetryPayload):
clients[client_id] = {"pending_command": "none", "last_seen": None, "active": True}
# Update last_seen timestamp, active status, and store the current telemetry details
clients[client_id]["last_seen"] = datetime.now().isoformat()
clients[client_id]["active"] = True
mark_client_active(client_id, clients)
clients[client_id]["telemetry"] = payload.dict()
save_clients(clients)
@@ -397,6 +418,9 @@ async def receive_command_result(client_id: str, payload: CommandResultPayload):
"""
Endpoint for remote agents to push execution results back to the central server.
"""
if not client_id or not client_id.strip():
raise HTTPException(status_code=400, detail="client_id cannot be empty")
append_to_logs("command_result", client_id, payload.dict())
return {"status": "success"}

20
commands.json
View File

@@ -10,9 +10,13 @@
"check_space": ["powershell.exe", "-Command", "Get-PSDrive -PSProvider FileSystem"],
"check_takeleap": ["python", "client_agent_prototype.py", "--check-takeleap"],
"list_takeleap": ["python", "client_agent_prototype.py", "--list-takeleap"],
"list_takeleap_shift": ["python", "client_agent_prototype.py", "--list-takeleap-sub", "SHIFT"],
"list_takeleap_errors": ["python", "client_agent_prototype.py", "--list-takeleap-sub", "Error_Videos"],
"list_takeleap_upload": ["python", "client_agent_prototype.py", "--list-takeleap-sub", "UPLOAD_FOLDER"]
"list_takeleap_size": ["python", "client_agent_prototype.py", "--list-takeleap-size"],
"list_takeleap_shift_size": ["python", "client_agent_prototype.py", "--list-takeleap-sub-size", "SHIFT"],
"list_takeleap_shift_files": ["python", "client_agent_prototype.py", "--list-takeleap-sub-files", "SHIFT"],
"list_takeleap_errors_size": ["python", "client_agent_prototype.py", "--list-takeleap-sub-size", "Error_Videos"],
"list_takeleap_errors_files": ["python", "client_agent_prototype.py", "--list-takeleap-sub-files", "Error_Videos"],
"list_takeleap_upload_size": ["python", "client_agent_prototype.py", "--list-takeleap-sub-size", "UPLOAD_FOLDER"],
"list_takeleap_upload_files": ["python", "client_agent_prototype.py", "--list-takeleap-sub-files", "UPLOAD_FOLDER"]
},
"posix": {
"whoami": ["whoami"],
@@ -55,8 +59,12 @@
"netstat": ["ss", "-tulnp"],
"check_takeleap": ["python3", "/opt/seekright-agent/client_agent_prototype.py", "--check-takeleap"],
"list_takeleap": ["python3", "/opt/seekright-agent/client_agent_prototype.py", "--list-takeleap"],
"list_takeleap_shift": ["python3", "/opt/seekright-agent/client_agent_prototype.py", "--list-takeleap-sub", "SHIFT"],
"list_takeleap_errors": ["python3", "/opt/seekright-agent/client_agent_prototype.py", "--list-takeleap-sub", "Error_Videos"],
"list_takeleap_upload": ["python3", "/opt/seekright-agent/client_agent_prototype.py", "--list-takeleap-sub", "UPLOAD_FOLDER"]
"list_takeleap_size": ["python3", "/opt/seekright-agent/client_agent_prototype.py", "--list-takeleap-size"],
"list_takeleap_shift_size": ["python3", "/opt/seekright-agent/client_agent_prototype.py", "--list-takeleap-sub-size", "SHIFT"],
"list_takeleap_shift_files": ["python3", "/opt/seekright-agent/client_agent_prototype.py", "--list-takeleap-sub-files", "SHIFT"],
"list_takeleap_errors_size": ["python3", "/opt/seekright-agent/client_agent_prototype.py", "--list-takeleap-sub-size", "Error_Videos"],
"list_takeleap_errors_files": ["python3", "/opt/seekright-agent/client_agent_prototype.py", "--list-takeleap-sub-files", "Error_Videos"],
"list_takeleap_upload_size": ["python3", "/opt/seekright-agent/client_agent_prototype.py", "--list-takeleap-sub-size", "UPLOAD_FOLDER"],
"list_takeleap_upload_files": ["python3", "/opt/seekright-agent/client_agent_prototype.py", "--list-takeleap-sub-files", "UPLOAD_FOLDER"]
}
}

186
deploy_agent.sh
View File

@@ -48,6 +48,7 @@ import psutil
import os
import sys
import string
import threading
CLIENT_ID = "TEMPLATE_CLIENT_ID"
CENTRAL_BASE = "TEMPLATE_CENTRAL_URL"
@@ -59,6 +60,32 @@ platform_name = "nt" if os.name == "nt" else "posix"
CENTRAL_COMMANDS_URL = f"{CENTRAL_BASE}/api/get-commands?platform={platform_name}"
CACHED_COMMANDS = {}
command_in_progress = False
def run_command_async(command_key, safe_cmd_list, is_shutdown_or_restart):
global command_in_progress
try:
result = subprocess.run(safe_cmd_list, capture_output=True, text=True)
output = result.stdout.strip()
error_output = result.stderr.strip() if result.stderr else ""
# Post execution results to Central Server
send_command_result_to_server(command_key, result.returncode, output, error_output)
if not is_shutdown_or_restart:
if result.returncode == 0:
print(f" -> Output: {output}")
send_rocket_chat_notification(f"✅ Success: {command_key}", output, "#00FF00")
else:
print(f" -> Error: {error_output}")
send_rocket_chat_notification(f"❌ Failed: {command_key}", error_output, "#FF0000")
except Exception as ex:
print(f" -> [!] Execution Error: {ex}")
send_command_result_to_server(command_key, -1, "", str(ex))
if not is_shutdown_or_restart:
send_rocket_chat_notification(f"❌ Execution Error: {command_key}", str(ex), "#FF0000")
finally:
command_in_progress = False
def load_allowed_commands():
global CACHED_COMMANDS
@@ -185,6 +212,9 @@ def send_command_result_to_server(command: str, returncode: int, stdout: str, st
print(f" -> [!] Failed to send command result to server: {e}")
def poll_server():
global command_in_progress
if command_in_progress:
return
try:
print(f"[*] Checking for commands...")
req = urllib.request.Request(CENTRAL_API_URL)
@@ -211,26 +241,10 @@ def poll_server():
)
time.sleep(3)
try:
result = subprocess.run(safe_cmd_list, capture_output=True, text=True)
output = result.stdout.strip()
error_output = result.stderr.strip() if result.stderr else ""
# Post execution results to Central Server
send_command_result_to_server(command_key, result.returncode, output, error_output)
if not is_shutdown_or_restart:
if result.returncode == 0:
print(f" -> Output: {output}")
send_rocket_chat_notification(f"✅ Success: {command_key}", output, "#00FF00")
else:
print(f" -> Error: {error_output}")
send_rocket_chat_notification(f"❌ Failed: {command_key}", error_output, "#FF0000")
except Exception as ex:
print(f" -> [!] Execution Error: {ex}")
send_command_result_to_server(command_key, -1, "", str(ex))
if not is_shutdown_or_restart:
send_rocket_chat_notification(f"❌ Execution Error: {command_key}", str(ex), "#FF0000")
command_in_progress = True
t = threading.Thread(target=run_command_async, args=(command_key, safe_cmd_list, is_shutdown_or_restart))
t.daemon = True
t.start()
else:
warning_msg = f"Server requested unknown/malicious command: '{command_key}'. Ignored."
print(f" -> [!] SECURITY WARNING: {warning_msg}")
@@ -239,6 +253,13 @@ def poll_server():
print(f" -> [!] Failed to connect to central server: {e}")
def get_dir_size(path):
if os.name != "nt":
try:
result = subprocess.run(["du", "-sb", path], capture_output=True, text=True, timeout=30)
if result.returncode == 0:
return int(result.stdout.split()[0])
except Exception:
pass
total_size = 0
try:
for dirpath, dirnames, filenames in os.walk(path):
@@ -252,6 +273,30 @@ def get_dir_size(path):
pass
return total_size
def get_takeleap_sizes(root_path, subdirs):
sizes = {d: 0 for d in subdirs}
total_size = 0
subdir_abs_paths = {}
for d in subdirs:
subdir_abs_paths[os.path.join(root_path, d)] = d
try:
for dirpath, dirnames, filenames in os.walk(root_path):
dir_size = 0
for f in filenames:
fp = os.path.join(dirpath, f)
try:
dir_size += os.path.getsize(fp)
except Exception:
pass
total_size += dir_size
for sub_abs, sub_name in subdir_abs_paths.items():
if dirpath == sub_abs or dirpath.startswith(sub_abs + os.sep):
sizes[sub_name] += dir_size
break
except Exception:
pass
return total_size, sizes
def format_size(size_in_bytes):
for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
if size_in_bytes < 1024.0:
@@ -305,7 +350,7 @@ def check_takeleap_folders():
else:
print("true" if files_found else "false")
def list_takeleap_parent():
def list_takeleap_parent_folders():
found_any = False
if os.name == "nt":
@@ -327,13 +372,51 @@ def list_takeleap_parent():
if os.path.basename(root).upper() == "TAKELEAP":
found_any = True
total_size = get_dir_size(root)
print(f"\n[+] Found TAKELEAP at: {root}")
try:
items = os.listdir(root)
if items:
print(" Contents:")
for idx, item in enumerate(items, 1):
is_dir = "Folder" if os.path.isdir(os.path.join(root, item)) else "File"
print(f" {idx}. {item} ({is_dir})")
else:
print(" (Empty)")
except Exception as e:
print(f" (Error listing directory: {e})")
if not found_any:
print("No TAKELEAP folders found on this system.")
def list_takeleap_parent_size():
found_any = False
if os.name == "nt":
drives = [f"{letter}:\\" for letter in string.ascii_uppercase if os.path.exists(f"{letter}:\\")]
else:
drives = ["/mnt", "/media", "/home"]
if os.path.exists("/TAKELEAP"):
drives.append("/TAKELEAP")
for drive in drives:
for root, dirs, files in os.walk(drive, topdown=True):
if os.name == "nt":
dirs[:] = [d for d in dirs if d.lower() not in {"windows", "program files", "program files (x86)", "appdata", "programdata", "$recycle.bin", "system volume information", "node_modules", ".git"}]
else:
if root == "/":
dirs[:] = [d for d in dirs if d.lower() not in {"proc", "sys", "dev", "var", "lib", "run", "boot", "snap", "node_modules", ".git"}]
elif root == "/mnt":
dirs[:] = [d for d in dirs if not (len(d) == 1 and d.isalpha()) and d.lower() not in {"wslg", "wsl"}]
if os.path.basename(root).upper() == "TAKELEAP":
found_any = True
subdirs = list(dirs)
total_size, subdir_sizes = get_takeleap_sizes(root, subdirs)
print(f"\n[+] Found TAKELEAP at: {root} (Total Size: {format_size(total_size)})")
if dirs:
if subdirs:
print(" Subfolders present:")
for d in dirs:
subpath = os.path.join(root, d)
sub_size = get_dir_size(subpath)
for d in subdirs:
sub_size = subdir_sizes.get(d, 0)
print(f" - {d} (Size: {format_size(sub_size)})")
else:
print(" (No subfolders found inside)")
@@ -341,7 +424,7 @@ def list_takeleap_parent():
if not found_any:
print("No TAKELEAP folders found on this system.")
def list_takeleap_subfolder(subfolder_name):
def list_takeleap_subfolder_size(subfolder_name):
found_any = False
if os.name == "nt":
@@ -368,8 +451,41 @@ def list_takeleap_subfolder(subfolder_name):
try:
found_any = True
sub_size = get_dir_size(subpath)
items = [item for item in os.listdir(subpath) if os.path.isfile(os.path.join(subpath, item))]
print(f"\n[+] Folder: {subpath} (Total Size: {format_size(sub_size)})")
except Exception as e:
print(f" [!] Error accessing folder {subpath}: {e}")
if not found_any:
print(f"No TAKELEAP/{subfolder_name} folders found on this system.")
def list_takeleap_subfolder_files(subfolder_name):
found_any = False
if os.name == "nt":
drives = [f"{letter}:\\" for letter in string.ascii_uppercase if os.path.exists(f"{letter}:\\")]
else:
drives = ["/mnt", "/media", "/home"]
if os.path.exists("/TAKELEAP"):
drives.append("/TAKELEAP")
for drive in drives:
for root, dirs, files in os.walk(drive, topdown=True):
if os.name == "nt":
dirs[:] = [d for d in dirs if d.lower() not in {"windows", "program files", "program files (x86)", "appdata", "programdata", "$recycle.bin", "system volume information", "node_modules", ".git"}]
else:
if root == "/":
dirs[:] = [d for d in dirs if d.lower() not in {"proc", "sys", "dev", "var", "lib", "run", "boot", "snap", "node_modules", ".git"}]
elif root == "/mnt":
dirs[:] = [d for d in dirs if not (len(d) == 1 and d.isalpha()) and d.lower() not in {"wslg", "wsl"}]
if os.path.basename(root).upper() == "TAKELEAP":
for sub in dirs:
if sub.lower() == subfolder_name.lower():
subpath = os.path.join(root, sub)
try:
found_any = True
items = [item for item in os.listdir(subpath) if os.path.isfile(os.path.join(subpath, item))]
print(f"\n[+] Folder Contents: {subpath}")
if items:
for idx, item in enumerate(items, 1):
print(f" {idx}. {item}")
@@ -381,17 +497,29 @@ def list_takeleap_subfolder(subfolder_name):
if not found_any:
print(f"No TAKELEAP/{subfolder_name} folders found on this system.")
def list_takeleap_subfolder(subfolder_name):
list_takeleap_subfolder_files(subfolder_name)
if __name__ == "__main__":
if len(sys.argv) > 1:
if sys.argv[1] == "--check-takeleap":
check_takeleap_folders()
sys.exit(0)
elif sys.argv[1] == "--list-takeleap":
list_takeleap_parent()
list_takeleap_parent_folders()
sys.exit(0)
elif sys.argv[1] == "--list-takeleap-size":
list_takeleap_parent_size()
sys.exit(0)
elif sys.argv[1] == "--list-takeleap-sub" and len(sys.argv) > 2:
list_takeleap_subfolder(sys.argv[2])
sys.exit(0)
elif sys.argv[1] == "--list-takeleap-sub-size" and len(sys.argv) > 2:
list_takeleap_subfolder_size(sys.argv[2])
sys.exit(0)
elif sys.argv[1] == "--list-takeleap-sub-files" and len(sys.argv) > 2:
list_takeleap_subfolder_files(sys.argv[2])
sys.exit(0)
else:
print(f"Unknown argument or missing parameters: {sys.argv[1]}")
sys.exit(1)