500 lines
16 KiB
Python
500 lines
16 KiB
Python
#!/usr/bin/env python3
|
||
# -*- encoding=utf8 -*-
|
||
"""
|
||
Worker 批量添加脚本
|
||
|
||
扫描指定 IP 范围内的设备，通过 SSH 获取 MAC 地址，
|
||
自动生成 Worker 条目并合并写入 config/config.json 的 WORKER_INVENTORY。
|
||
|
||
使用示例:
|
||
# 扫描 192.168.1.10 到 192.168.1.50
|
||
python scripts/add_workers.py --range 192.168.1.10-50
|
||
|
||
# 扫描多个范围
|
||
python scripts/add_workers.py --range 192.168.1.10-50 192.168.2.50-80
|
||
|
||
# 指定标签
|
||
python scripts/add_workers.py --range 192.168.1.10-50 --tags group_b
|
||
|
||
# 仅扫描，不写入
|
||
python scripts/add_workers.py --range 192.168.1.10-50 --dry-run
|
||
|
||
# 跳过连通性检测
|
||
python scripts/add_workers.py --range 192.168.1.10-50 --skip-ping
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import re
|
||
import socket
|
||
import sys
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from typing import Dict, List, Optional, Tuple
|
||
|
||
# 项目根目录
|
||
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
sys.path.insert(0, BASE_DIR)
|
||
|
||
from config import (
|
||
CONFIG_PATH,
|
||
SSH_DEFAULT_USER,
|
||
SSH_DEFAULT_PASSWORD,
|
||
SSH_DEFAULT_PORT,
|
||
WORKER_INVENTORY_CONFIG_KEY,
|
||
WORKER_INVENTORY_PATH,
|
||
)
|
||
|
||
# 默认 Worker 配置
|
||
DEFAULT_REPO_DIR = "D:/autool"
|
||
DEFAULT_PYTHON_EXE = "python"
|
||
|
||
# 并发设置
|
||
MAX_PING_WORKERS = 50
|
||
MAX_SSH_WORKERS = 10
|
||
PING_TIMEOUT = 2 # 秒
|
||
SSH_TIMEOUT = 10 # 秒
|
||
|
||
|
||
def parse_ip_range(range_str: str) -> List[str]:
|
||
"""
|
||
解析 IP 范围字符串，支持以下格式：
|
||
- 单个 IP: 192.168.1.10
|
||
- 末段范围: 192.168.1.10-50
|
||
- CIDR: 192.168.1.0/24
|
||
"""
|
||
range_str = range_str.strip()
|
||
|
||
# CIDR 格式: 192.168.1.0/24
|
||
cidr_match = re.fullmatch(r"(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})/(\d{1,2})", range_str)
|
||
if cidr_match:
|
||
base_ip = cidr_match.group(1)
|
||
prefix_len = int(cidr_match.group(2))
|
||
if prefix_len < 24 or prefix_len > 32:
|
||
raise ValueError(f"仅支持 /24 ~ /32 的 CIDR 范围: {range_str}")
|
||
parts = base_ip.split(".")
|
||
base = int(parts[3])
|
||
host_bits = 32 - prefix_len
|
||
count = 2 ** host_bits
|
||
# 排除网络地址和广播地址
|
||
start = max(base, 1)
|
||
end = min(base + count - 1, 254)
|
||
return [f"{parts[0]}.{parts[1]}.{parts[2]}.{i}" for i in range(start, end + 1)]
|
||
|
||
# 末段范围格式: 192.168.1.10-50
|
||
range_match = re.fullmatch(r"(\d{1,3}\.\d{1,3}\.\d{1,3})\.(\d{1,3})-(\d{1,3})", range_str)
|
||
if range_match:
|
||
prefix = range_match.group(1)
|
||
start = int(range_match.group(2))
|
||
end = int(range_match.group(3))
|
||
if start > end:
|
||
raise ValueError(f"起始地址大于结束地址: {range_str}")
|
||
if end > 254:
|
||
raise ValueError(f"IP 末段超出范围 (最大 254): {range_str}")
|
||
return [f"{prefix}.{i}" for i in range(start, end + 1)]
|
||
|
||
# 单个 IP: 192.168.1.10
|
||
ip_match = re.fullmatch(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", range_str)
|
||
if ip_match:
|
||
return [range_str]
|
||
|
||
raise ValueError(f"无法解析 IP 范围格式: {range_str}")
|
||
|
||
|
||
def check_ssh_port(ip: str, port: int = SSH_DEFAULT_PORT, timeout: float = PING_TIMEOUT) -> bool:
|
||
"""尝试通过 TCP 连接 SSH 端口检查主机是否在线"""
|
||
try:
|
||
with socket.create_connection((ip, port), timeout=timeout):
|
||
return True
|
||
except (socket.timeout, ConnectionRefusedError, OSError):
|
||
return False
|
||
|
||
|
||
def scan_reachable_hosts(ip_list: List[str], port: int = SSH_DEFAULT_PORT, skip_ping: bool = False) -> List[str]:
|
||
"""并发端口扫描，返回 SSH 端口可达的主机列表"""
|
||
if skip_ping:
|
||
print(f" 跳过连通性检测，将直接尝试 SSH 登录 {len(ip_list)} 台设备")
|
||
return ip_list
|
||
|
||
print(f" 正在扫描 {len(ip_list)} 个 IP 地址的 SSH 端口...")
|
||
reachable = []
|
||
|
||
with ThreadPoolExecutor(max_workers=MAX_PING_WORKERS) as executor:
|
||
futures = {executor.submit(check_ssh_port, ip, port): ip for ip in ip_list}
|
||
for future in as_completed(futures):
|
||
ip = futures[future]
|
||
try:
|
||
if future.result():
|
||
reachable.append(ip)
|
||
print(f" ✓ {ip} 在线 (端口可达)")
|
||
else:
|
||
print(f" ✗ {ip} 不可达 (端口关闭或超时)")
|
||
except Exception as e:
|
||
print(f" ✗ {ip} 检测异常: {e}")
|
||
|
||
# 按 IP 排序
|
||
reachable.sort(key=lambda x: tuple(int(p) for p in x.split(".")))
|
||
print(f" 扫描完成，{len(reachable)} 台设备 SSH 端口可连接\n")
|
||
return reachable
|
||
|
||
|
||
def get_mac_via_ssh(
|
||
ip: str,
|
||
user: str = SSH_DEFAULT_USER,
|
||
password: str = SSH_DEFAULT_PASSWORD,
|
||
port: int = SSH_DEFAULT_PORT,
|
||
) -> Optional[str]:
|
||
"""
|
||
通过 SSH 连接 Windows 设备并获取物理网卡 MAC 地址。
|
||
使用 getmac /fo csv /nh 获取 MAC 列表，选取以太网适配器的 MAC。
|
||
"""
|
||
try:
|
||
import paramiko
|
||
except ImportError:
|
||
print(" [错误] 需要 paramiko 库，请运行: pip install paramiko")
|
||
sys.exit(1)
|
||
|
||
ssh = paramiko.SSHClient()
|
||
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||
|
||
try:
|
||
ssh.connect(
|
||
hostname=ip,
|
||
port=port,
|
||
username=user,
|
||
password=password,
|
||
timeout=SSH_TIMEOUT,
|
||
auth_timeout=SSH_TIMEOUT,
|
||
banner_timeout=SSH_TIMEOUT,
|
||
look_for_keys=True,
|
||
allow_agent=True,
|
||
)
|
||
|
||
# 获取 MAC 地址列表
|
||
_, stdout, _ = ssh.exec_command("getmac /fo csv /nh", timeout=SSH_TIMEOUT)
|
||
output = stdout.read().decode("utf-8", errors="ignore").strip()
|
||
|
||
if not output:
|
||
return None
|
||
|
||
# 解析 getmac 输出，格式: "MAC地址","传输名称","..."
|
||
# 优先选取以太网/有线网卡的 MAC，跳过虚拟网卡和断开的连接
|
||
best_mac = None
|
||
for line in output.splitlines():
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
# 去掉引号分割
|
||
parts = [p.strip().strip('"') for p in line.split(",")]
|
||
if len(parts) < 2:
|
||
continue
|
||
mac = parts[0].strip()
|
||
transport = parts[1].strip() if len(parts) > 1 else ""
|
||
|
||
# 跳过无效 MAC
|
||
if not re.match(r"([0-9A-Fa-f]{2}[:-]){5}[0-9A-Fa-f]{2}", mac):
|
||
continue
|
||
|
||
# 跳过已断开的连接
|
||
if "已断开" in transport or "Disconnected" in transport.lower() or "Media disconnected" in transport.lower():
|
||
continue
|
||
|
||
# 统一 MAC 格式为 XX:XX:XX:XX:XX:XX（大写、冒号分隔）
|
||
mac = mac.upper().replace("-", ":")
|
||
|
||
# 优先选取硬件以太网适配器
|
||
transport_lower = transport.lower()
|
||
if any(kw in transport_lower for kw in ["ethernet", "以太网", "realtek", "intel"]):
|
||
best_mac = mac
|
||
break
|
||
|
||
# 记录第一个有效的 MAC 作为备选
|
||
if best_mac is None:
|
||
best_mac = mac
|
||
|
||
return best_mac
|
||
except Exception as e:
|
||
print(f" ✗ {ip} SSH 连接失败: {e}")
|
||
return None
|
||
finally:
|
||
ssh.close()
|
||
|
||
|
||
def collect_worker_info(
|
||
ip_list: List[str],
|
||
tags: List[str],
|
||
user: str = SSH_DEFAULT_USER,
|
||
password: str = SSH_DEFAULT_PASSWORD,
|
||
port: int = SSH_DEFAULT_PORT,
|
||
) -> List[Dict]:
|
||
"""批量通过 SSH 收集设备信息，生成 worker 条目"""
|
||
print(f" 正在通过 SSH 获取 {len(ip_list)} 台设备的 MAC 地址...")
|
||
workers = []
|
||
|
||
def _collect_one(ip: str) -> Optional[Dict]:
|
||
mac = get_mac_via_ssh(ip, user=user, password=password, port=port)
|
||
if mac:
|
||
worker_id = f"{ip}_{mac}"
|
||
print(f" ✓ {ip} → MAC: {mac} → worker_id: {worker_id}")
|
||
return {
|
||
"worker_id": worker_id,
|
||
"ssh_target": ip,
|
||
"repo_dir": DEFAULT_REPO_DIR,
|
||
"python_exe": DEFAULT_PYTHON_EXE,
|
||
"tags": list(tags),
|
||
}
|
||
else:
|
||
print(f" ✗ {ip} 无法获取 MAC 地址")
|
||
return None
|
||
|
||
with ThreadPoolExecutor(max_workers=MAX_SSH_WORKERS) as executor:
|
||
futures = {executor.submit(_collect_one, ip): ip for ip in ip_list}
|
||
for future in as_completed(futures):
|
||
result = future.result()
|
||
if result:
|
||
workers.append(result)
|
||
|
||
# 按 IP 排序
|
||
workers.sort(key=lambda w: tuple(int(p) for p in w["ssh_target"].split(".")))
|
||
print(f" 成功获取 {len(workers)} 台设备信息\n")
|
||
return workers
|
||
|
||
|
||
def load_existing_inventory(path: str) -> Tuple[List[Dict], set, set]:
|
||
"""加载现有 Worker 列表，返回列表、已存在 worker_id 和已存在 IP 集合"""
|
||
if not os.path.exists(path):
|
||
return [], set(), set()
|
||
|
||
with open(path, "r", encoding="utf-8") as f:
|
||
payload = json.load(f)
|
||
|
||
data = payload.get(WORKER_INVENTORY_CONFIG_KEY, []) if isinstance(payload, dict) else payload
|
||
|
||
if not isinstance(data, list):
|
||
print(f" [警告] {path} 格式异常（非数组），将创建新文件")
|
||
return [], set(), set()
|
||
|
||
existing_ids = set()
|
||
# 同时收集已存在的 ssh_target（IP）用于去重
|
||
existing_ips = set()
|
||
for item in data:
|
||
if isinstance(item, dict):
|
||
wid = item.get("worker_id", "")
|
||
if wid:
|
||
existing_ids.add(wid)
|
||
ip = item.get("ssh_target", "")
|
||
if ip:
|
||
existing_ips.add(ip)
|
||
|
||
return data, existing_ids, existing_ips
|
||
|
||
|
||
def save_inventory(path: str, workers: List[Dict]) -> None:
|
||
"""写入统一配置；如果传入旧数组文件路径，则按旧格式写回。"""
|
||
payload = {}
|
||
if os.path.exists(path):
|
||
with open(path, "r", encoding="utf-8") as f:
|
||
try:
|
||
payload = json.load(f)
|
||
except json.JSONDecodeError:
|
||
payload = {}
|
||
|
||
if isinstance(payload, dict):
|
||
payload[WORKER_INVENTORY_CONFIG_KEY] = workers
|
||
with open(path, "w", encoding="utf-8") as f:
|
||
json.dump(payload, f, indent=2, ensure_ascii=False)
|
||
f.write("\n")
|
||
return
|
||
|
||
with open(path, "w", encoding="utf-8") as f:
|
||
json.dump(workers, f, indent=2, ensure_ascii=False)
|
||
f.write("\n")
|
||
|
||
|
||
def merge_and_save(
|
||
existing: List[Dict],
|
||
existing_ids: set,
|
||
existing_ips: set,
|
||
new_workers: List[Dict],
|
||
path: str,
|
||
dry_run: bool = False,
|
||
) -> int:
|
||
"""合并新旧 Worker 列表并写入文件"""
|
||
to_add = []
|
||
skipped_dup_id = 0
|
||
skipped_dup_ip = 0
|
||
|
||
for worker in new_workers:
|
||
wid = worker["worker_id"]
|
||
ip = worker["ssh_target"]
|
||
if wid in existing_ids:
|
||
print(f" [跳过] worker_id 已存在: {wid}")
|
||
skipped_dup_id += 1
|
||
continue
|
||
if ip in existing_ips:
|
||
print(f" [跳过] IP 已存在于其他 worker: {ip}")
|
||
skipped_dup_ip += 1
|
||
continue
|
||
to_add.append(worker)
|
||
|
||
if not to_add:
|
||
print(" 没有新的 Worker 需要添加")
|
||
return 0
|
||
|
||
print(f"\n 将添加 {len(to_add)} 台新 Worker:")
|
||
for w in to_add:
|
||
tag_str = f" tags={w['tags']}" if w["tags"] else ""
|
||
print(f" + {w['worker_id']}{tag_str}")
|
||
|
||
if skipped_dup_id > 0:
|
||
print(f" （已跳过 {skipped_dup_id} 台 worker_id 重复的设备）")
|
||
if skipped_dup_ip > 0:
|
||
print(f" （已跳过 {skipped_dup_ip} 台 IP 重复的设备）")
|
||
|
||
if dry_run:
|
||
print("\n [Dry Run] 未写入文件")
|
||
return len(to_add)
|
||
|
||
# 合并并写入
|
||
merged = existing + to_add
|
||
# 备份原文件
|
||
if os.path.exists(path):
|
||
backup_path = path + ".bak"
|
||
with open(path, "r", encoding="utf-8") as f:
|
||
backup_content = f.read()
|
||
with open(backup_path, "w", encoding="utf-8") as f:
|
||
f.write(backup_content)
|
||
print(f"\n 已备份原文件到: {backup_path}")
|
||
|
||
save_inventory(path, merged)
|
||
|
||
print(f" 已写入 {path}，共 {len(merged)} 台 Worker")
|
||
return len(to_add)
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description="Worker 批量添加工具 — 扫描 IP 范围并自动添加到 config/config.json",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog="""
|
||
示例:
|
||
python scripts/add_workers.py --range 192.168.1.10-50
|
||
python scripts/add_workers.py --range 192.168.1.10-50 192.168.2.50-80
|
||
python scripts/add_workers.py --range 192.168.1.10-50 --tags group_b
|
||
python scripts/add_workers.py --range 192.168.1.10-50 --dry-run
|
||
python scripts/add_workers.py --range 192.168.1.10-50 --skip-ping
|
||
""",
|
||
)
|
||
parser.add_argument(
|
||
"--range",
|
||
nargs="+",
|
||
required=True,
|
||
dest="ip_ranges",
|
||
help="IP 范围，支持格式: 192.168.1.10 / 192.168.1.10-50 / 192.168.1.0/24",
|
||
)
|
||
parser.add_argument(
|
||
"--tags",
|
||
nargs="*",
|
||
default=[],
|
||
help="为新添加的 Worker 设置标签",
|
||
)
|
||
parser.add_argument(
|
||
"--dry-run",
|
||
action="store_true",
|
||
help="仅扫描和预览，不写入文件",
|
||
)
|
||
parser.add_argument(
|
||
"--skip-ping",
|
||
action="store_true",
|
||
help="跳过连通性检测，直接尝试 SSH 登录",
|
||
)
|
||
parser.add_argument(
|
||
"--inventory",
|
||
default=WORKER_INVENTORY_PATH,
|
||
help=f"配置文件或旧 worker_inventory.json 路径 (默认: {CONFIG_PATH})",
|
||
)
|
||
parser.add_argument(
|
||
"--user",
|
||
default=SSH_DEFAULT_USER,
|
||
help=f"SSH 用户名 (默认: {SSH_DEFAULT_USER})",
|
||
)
|
||
parser.add_argument(
|
||
"--password",
|
||
default=SSH_DEFAULT_PASSWORD,
|
||
help="SSH 密码",
|
||
)
|
||
parser.add_argument(
|
||
"--port",
|
||
type=int,
|
||
default=SSH_DEFAULT_PORT,
|
||
help=f"SSH 端口 (默认: {SSH_DEFAULT_PORT})",
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
print("=" * 60)
|
||
print("Worker 批量添加工具")
|
||
print("=" * 60)
|
||
|
||
# 1. 解析 IP 范围
|
||
all_ips = []
|
||
for r in args.ip_ranges:
|
||
try:
|
||
ips = parse_ip_range(r)
|
||
print(f" 范围 {r} → {len(ips)} 个 IP")
|
||
all_ips.extend(ips)
|
||
except ValueError as e:
|
||
print(f" [错误] {e}")
|
||
sys.exit(1)
|
||
|
||
# 去重并排序
|
||
all_ips = sorted(set(all_ips), key=lambda x: tuple(int(p) for p in x.split(".")))
|
||
print(f"\n 共计 {len(all_ips)} 个唯一 IP 地址\n")
|
||
|
||
if not all_ips:
|
||
print(" 没有有效的 IP 地址")
|
||
return
|
||
|
||
# 2. 端口扫描
|
||
print("[1/4] 连通性检测")
|
||
reachable = scan_reachable_hosts(all_ips, port=args.port, skip_ping=args.skip_ping)
|
||
|
||
if not reachable:
|
||
print(" 没有在线的设备")
|
||
return
|
||
|
||
# 3. SSH 获取 MAC 地址
|
||
print("[2/4] 获取设备信息")
|
||
new_workers = collect_worker_info(
|
||
reachable,
|
||
tags=args.tags,
|
||
user=args.user,
|
||
password=args.password,
|
||
port=args.port,
|
||
)
|
||
|
||
if not new_workers:
|
||
print(" 未能获取任何设备信息")
|
||
return
|
||
|
||
# 4. 加载现有 inventory
|
||
print("[3/4] 加载现有 Worker 列表")
|
||
existing, existing_ids, existing_ips = load_existing_inventory(args.inventory)
|
||
print(f" 当前共有 {len(existing)} 台 Worker\n")
|
||
|
||
# 5. 合并写入
|
||
print("[4/4] 合并并写入")
|
||
added = merge_and_save(
|
||
existing, existing_ids, existing_ips, new_workers, args.inventory, dry_run=args.dry_run,
|
||
)
|
||
|
||
print("\n" + "=" * 60)
|
||
print(f"完成！新增 {added} 台 Worker")
|
||
if args.dry_run:
|
||
print("（Dry Run 模式，未实际写入）")
|
||
print("=" * 60)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|