Update: mass_update.py - config in script dir, per-server SSH keys

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)