diff --git a/tools/mass_update.py b/tools/mass_update.py index b49d72cc..431b48b5 100755 --- a/tools/mass_update.py +++ b/tools/mass_update.py @@ -7,30 +7,64 @@ Usage: python3 mass_update.py [fast|full] [-c config] [-j jobs] Config format (JSON): { "servers": [ - {"host": "10.0.0.1", "workdir": "/opt/utun"}, - {"host": "10.0.0.2", "port": 2222, "user": "admin", "workdir": "/root/utun"} + {"host": "10.0.0.1"}, + {"host": "10.0.0.2", "port": 2222, "user": "admin", "ssh_key": "keys/server2_rsa", "workdir": "/opt/utun"} ] } -Default config: ~/.utun_servers.json +Default config: tools/utun_servers.json (in script directory) +Fallback: ~/.utun_servers.json +SSH keys: tools/keys/ (in script directory) """ import argparse import json +import os import subprocess import sys from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path -DEFAULT_CONFIG = Path.home() / ".utun_servers.json" -SSH_TIMEOUT = 300 # 5 minutes +SSH_TIMEOUT = 300 -def load_servers(config_path): +def get_script_dir(): + """Get script directory (tools/).""" + return Path(__file__).parent.resolve() + + +def get_default_config(): + """Get config path: tools/utun_servers.json or fallback to home.""" + script_dir = get_script_dir() + config_in_tools = script_dir / "utun_servers.json" + if config_in_tools.exists(): + return config_in_tools + config_in_home = Path.home() / ".utun_servers.json" + if config_in_home.exists(): + return config_in_home + return config_in_tools + + +def resolve_ssh_key(ssh_key, script_dir): + """Resolve SSH key path.""" + if not ssh_key: + return None + key_path = Path(ssh_key) + if key_path.is_absolute(): + return str(key_path) + keys_dir = script_dir / "keys" + return str(keys_dir / ssh_key) + + +def load_servers(config_path, script_dir): """Load server list from JSON config.""" with open(config_path, 'r') as f: data = json.load(f) - return data.get("servers", []) + servers = data.get("servers", []) + for server in servers: + if "ssh_key" in server: + server["_resolved_key"] = resolve_ssh_key(server["ssh_key"], script_dir) + return servers def update_server(server, command): @@ -39,16 +73,19 @@ def update_server(server, command): port = server.get("port", 22) workdir = server.get("workdir", "~/utun") user = server.get("user", "root") + resolved_key = server.get("_resolved_key") - ssh_cmd = [ - "ssh", + ssh_cmd = ["ssh"] + if resolved_key and Path(resolved_key).exists(): + ssh_cmd.extend(["-i", resolved_key]) + ssh_cmd.extend([ "-p", str(port), "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=10", "-o", "BatchMode=yes", f"{user}@{host}", f"cd {workdir} && ./tools/utun_update.sh {command}" - ] + ]) try: result = subprocess.run( @@ -80,6 +117,8 @@ def update_server(server, command): def main(): + script_dir = get_script_dir() + parser = argparse.ArgumentParser( description="Mass update uTun servers via SSH" ) @@ -90,8 +129,7 @@ def main(): ) parser.add_argument( "-c", "--config", - default=str(DEFAULT_CONFIG), - help=f"Servers config file (default: {DEFAULT_CONFIG})" + help=f"Servers config file (default: auto from script dir)" ) parser.add_argument( "-j", "--jobs", @@ -102,20 +140,27 @@ def main(): args = parser.parse_args() - config_path = Path(args.config) + if args.config: + config_path = Path(args.config) + else: + config_path = get_default_config() + if not config_path.exists(): print(f"ERROR: Config not found: {config_path}", file=sys.stderr) - print("\nCreate it with format:", file=sys.stderr) + print("\nCreate config in script directory:", file=sys.stderr) + print(f" {script_dir / 'utun_servers.json'}", file=sys.stderr) + print("\nOr use -c to specify custom config.", file=sys.stderr) + print("\nFormat:", file=sys.stderr) example = { "servers": [ - {"host": "10.0.0.1", "workdir": "/opt/utun"}, - {"host": "10.0.0.2", "port": 2222, "user": "admin"} + {"host": "10.0.0.1"}, + {"host": "10.0.0.2", "port": 2222, "user": "admin", "ssh_key": "keys/server2_rsa"} ] } print(json.dumps(example, indent=2), file=sys.stderr) sys.exit(1) - servers = load_servers(config_path) + servers = load_servers(config_path, script_dir) if not servers: print("ERROR: No servers found in config", file=sys.stderr) sys.exit(1)