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

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)