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

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
+    mark_client_active(client_id, clients)
-    clients[client_id]["last_seen"] = datetime.now().isoformat()
     
     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()
+    mark_client_active(client_id, clients)
-    clients[client_id]["active"] = True
     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"}
 

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_size": ["python", "client_agent_prototype.py", "--list-takeleap-size"],
-    "list_takeleap_errors": ["python", "client_agent_prototype.py", "--list-takeleap-sub", "Error_Videos"],
+    "list_takeleap_shift_size": ["python", "client_agent_prototype.py", "--list-takeleap-sub-size", "SHIFT"],
-    "list_takeleap_upload": ["python", "client_agent_prototype.py", "--list-takeleap-sub", "UPLOAD_FOLDER"]
+    "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_size": ["python3", "/opt/seekright-agent/client_agent_prototype.py", "--list-takeleap-size"],
-    "list_takeleap_errors": ["python3", "/opt/seekright-agent/client_agent_prototype.py", "--list-takeleap-sub", "Error_Videos"],
+    "list_takeleap_shift_size": ["python3", "/opt/seekright-agent/client_agent_prototype.py", "--list-takeleap-sub-size", "SHIFT"],
-    "list_takeleap_upload": ["python3", "/opt/seekright-agent/client_agent_prototype.py", "--list-takeleap-sub", "UPLOAD_FOLDER"]
+    "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"]
   }
 }

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:
+                    command_in_progress = True
-                        result = subprocess.run(safe_cmd_list, capture_output=True, text=True)
+                    t = threading.Thread(target=run_command_async, args=(command_key, safe_cmd_list, is_shutdown_or_restart))
-                        output = result.stdout.strip()
+                    t.daemon = True
-                        error_output = result.stderr.strip() if result.stderr else ""
+                    t.start()
-                        
-                        # 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")
                 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:
+                    for d in subdirs:
-                        subpath = os.path.join(root, d)
+                        sub_size = subdir_sizes.get(d, 0)
-                        sub_size = get_dir_size(subpath)
                         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)