135 lines
4.7 KiB
Python
135 lines
4.7 KiB
Python
import os
|
|
import sys
|
|
|
|
|
|
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
sys.path.insert(0, PROJECT_ROOT)
|
|
|
|
from remote_worker_controller import RemoteWorkerController
|
|
from config import (
|
|
CLEAN_BACKUP_PATH,
|
|
LOCAL_IMPORT_DIR,
|
|
MUMU_MANAGER_PATH,
|
|
MUMU_RECOVER_AHK_EXE,
|
|
MUMU_RECOVER_SCRIPT_PATH,
|
|
NET_BRIDGE_CARD,
|
|
)
|
|
|
|
|
|
def create_controller() -> RemoteWorkerController:
|
|
return RemoteWorkerController([])
|
|
|
|
|
|
def test_cmd_task_command_wraps_inner_command_with_cmd_safe_quotes():
|
|
command = RemoteWorkerController._cmd_task_command(
|
|
'"C:\\Program Files\\Netease\\MuMu\\nx_main\\MuMuManager.exe" control -v 1 restart'
|
|
)
|
|
|
|
assert (
|
|
command
|
|
== 'C:\\Windows\\System32\\cmd.exe /d /q /s /c ""C:\\Program Files\\Netease\\MuMu\\nx_main\\MuMuManager.exe" control -v 1 restart"'
|
|
)
|
|
|
|
|
|
def test_cmd_switch_with_command_keeps_original_inner_quotes():
|
|
command = RemoteWorkerController._cmd_switch_with_command(
|
|
"/c",
|
|
'"C:\\Program Files\\Netease\\MuMu\\nx_main\\MuMuManager.exe" control -v 2 restart',
|
|
)
|
|
|
|
assert command == '/c "C:\\Program Files\\Netease\\MuMu\\nx_main\\MuMuManager.exe" control -v 2 restart'
|
|
|
|
|
|
def test_restart_mumu_task_command_keeps_schtasks_tr_argument_balanced():
|
|
controller = create_controller()
|
|
|
|
command = controller._restart_mumu_task_command({})
|
|
|
|
assert isinstance(command, dict)
|
|
assert command["command"] == "powershell -NoProfile -ExecutionPolicy Bypass -Command -"
|
|
assert "$taskName = 'RestartMuMu';" in command["stdin"]
|
|
assert "$taskExecute = 'C:\\Windows\\System32\\cmd.exe';" in command["stdin"]
|
|
assert '$taskArguments = \'/c "C:\\Program Files\\Netease\\MuMu\\nx_main\\MuMuManager.exe" control -v 2 restart\';' in command["stdin"]
|
|
|
|
|
|
def test_direct_wait_command_quotes_inner_cmd_payload_with_redirection():
|
|
controller = create_controller()
|
|
|
|
command = controller._recover_mumu_command(
|
|
{"ssh_target": "192.168.2.71:22", "repo_dir": "D:\\autool"}
|
|
)
|
|
|
|
assert isinstance(command, str)
|
|
assert command.startswith("cmd.exe /d /q /s /c ")
|
|
assert " & call " in command
|
|
assert MUMU_RECOVER_SCRIPT_PATH in command
|
|
assert "Recover script not found" in command
|
|
assert CLEAN_BACKUP_PATH in command
|
|
assert LOCAL_IMPORT_DIR in command
|
|
assert MUMU_MANAGER_PATH in command
|
|
assert NET_BRIDGE_CARD in command
|
|
assert MUMU_RECOVER_AHK_EXE in command
|
|
assert "schtasks /create /tn" not in command
|
|
|
|
|
|
def test_start_worker_task_uses_single_wrapped_cmd_argument():
|
|
controller = create_controller()
|
|
|
|
command = controller._start_worker_command(
|
|
{
|
|
"repo_dir": "C:\\Program Files\\AutoTool Dispatcher",
|
|
"python_exe": "C:\\Python39\\python.exe",
|
|
}
|
|
)
|
|
|
|
assert isinstance(command, dict)
|
|
assert command["command"] == "powershell -NoProfile -ExecutionPolicy Bypass -Command -"
|
|
assert "$taskName = 'RunBatch';" in command["stdin"]
|
|
assert "$taskExecute = 'C:\\Windows\\System32\\cmd.exe';" in command["stdin"]
|
|
assert '$taskArguments = \'/k C:\\Program Files\\AutoTool Dispatcher\\venv\\Scripts\\python.exe C:\\Program Files\\AutoTool Dispatcher\\batch_run.py\';' in command["stdin"]
|
|
assert "& schtasks /end /tn $taskName" in command["stdin"]
|
|
assert "taskkill /F /T /IM" not in command["stdin"]
|
|
|
|
|
|
def test_stop_worker_uses_simple_taskkill_for_console_hosts():
|
|
controller = create_controller()
|
|
|
|
command = controller._stop_worker_command({})
|
|
|
|
assert isinstance(command, str)
|
|
assert command.startswith("cmd.exe /d /v:on /q /s /c ")
|
|
assert 'schtasks /end /tn ""RunBatch""' in command
|
|
assert "taskkill /F /T /IM OpenConsole.exe" in command
|
|
assert "taskkill /F /T /IM WindowsTerminal.exe" in command
|
|
assert "taskkill /F /T /IM conhost.exe" in command
|
|
assert "taskkill /F /T /IM powershell.exe" in command
|
|
assert "start /min taskkill /F /T /IM cmd.exe" in command
|
|
assert "Stop-Process" not in command
|
|
assert "Get-CimInstance Win32_Process" not in command
|
|
|
|
|
|
def test_stop_worker_uses_legacy_end_for_protected_workers():
|
|
controller = create_controller()
|
|
|
|
command = controller._stop_worker_command({"worker_id": "192.168.1.51"})
|
|
|
|
assert isinstance(command, str)
|
|
assert 'schtasks /end /tn ""RunBatch""' in command
|
|
assert "taskkill /F /T /IM" not in command
|
|
|
|
|
|
def test_start_worker_uses_legacy_end_before_run_for_protected_workers():
|
|
controller = create_controller()
|
|
|
|
command = controller._start_worker_command(
|
|
{
|
|
"worker_id": "192.168.2.101",
|
|
"repo_dir": "C:\\autool",
|
|
"python_exe": "C:\\Python39\\python.exe",
|
|
}
|
|
)
|
|
|
|
assert isinstance(command, dict)
|
|
assert "& schtasks /end /tn $taskName" in command["stdin"]
|
|
assert "taskkill /F /T /IM" not in command["stdin"]
|