From 25e9b5e120a962ddc4d09dedc3546b07c592ee7d Mon Sep 17 00:00:00 2001 From: kaushik Date: Thu, 4 Jun 2026 15:48:38 +0530 Subject: [PATCH] feat: implement async command execution, improve client status tracking, and expand takeleap diagnostic commands. --- central_api_prototype.py | 32 ++++++- commands.json | 20 +++-- deploy_agent.sh | 186 +++++++++++++++++++++++++++++++++------ 3 files changed, 199 insertions(+), 39 deletions(-) diff --git a/central_api_prototype.py b/central_api_prototype.py index dd39df7..fb5c4f9 100644 --- a/central_api_prototype.py +++ b/central_api_prototype.py @@ -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"} diff --git a/commands.json b/commands.json index 1929bcc..4347c4d 100644 --- a/commands.json +++ b/commands.json @@ -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"] } } diff --git a/deploy_agent.sh b/deploy_agent.sh index 7f876d0..6e7f7e8 100644 --- a/deploy_agent.sh +++ b/deploy_agent.sh @@ -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)