switch_config.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
交换机配置自动化工具
支持华为交换机的批量配置
支持多设备并发配置和可视化进度显示
"""

import json
import paramiko
import time
import re
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm
import queue

def load_devices(filename):
    with open(filename, 'r', encoding='utf-8') as f:
        data = json.load(f)
        return data.get('devices', [])

def wait_for_prompt(shell, timeout=10):
    """等待命令提示符出现"""
    start_time = time.time()
    while time.time() - start_time < timeout:
        if shell.recv_ready():
            output = shell.recv(65535).decode('utf-8')
            # 检查是否出现提示符
            if re.search(r'[<>]\[.*\]', output) or re.search(r'[<>].*[>#]', output):
                return output
        time.sleep(0.1)
    return ""

def send_command_and_wait(shell, command, timeout=10):
    """发送命令并等待响应"""
    # 发送命令
    shell.send((command + '\n').encode('utf-8'))
    time.sleep(0.5)
    
    # 等待响应
    start_time = time.time()
    output = ""
    
    while time.time() - start_time < timeout:
        if shell.recv_ready():
            chunk = shell.recv(65535).decode('utf-8')
            output += chunk
            
            # 检查是否出现提示符（命令执行完成）
            if re.search(r'[<>]\[.*\]', chunk) or re.search(r'[<>].*[>#]', chunk):
                break
        time.sleep(0.1)
    
    return output

def configure_switch(device, commands, result_queue=None, device_index=None):
    """配置单个交换机"""
    result = {
        'device': device['hostname'],
        'ip': device['ip'],
        'success': False,
        'error': None,
        'start_time': time.time(),
        'end_time': None,
        'commands_executed': 0
    }
    
    try:
        # 创建SSH客户端
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        
        # 连接到交换机
        ssh.connect(device['ip'], 
                   username=device['username'], 
                   password=device['password'], 
                   port=device.get('port', 22),
                   timeout=10)
        
        # 获取交互式shell
        shell = ssh.invoke_shell()
        time.sleep(2)
        
        # 读取初始输出（登录信息等）
        if shell.recv_ready():
            shell.recv(65535).decode('utf-8')
        
        # 等待提示符出现
        wait_for_prompt(shell)
        
        # 发送配置命令
        for i, command in enumerate(commands):
            if command.strip():
                output = send_command_and_wait(shell, command)
                result['commands_executed'] = i + 1
        
        result['success'] = True
        result['end_time'] = time.time()
        
        # 关闭连接
        ssh.close()
        
    except Exception as e:
        result['error'] = str(e)
        result['end_time'] = time.time()
    
    # 将结果放入队列
    if result_queue:
        result_queue.put(result)
    
    return result

def configure_switches_concurrent(devices, commands, max_workers=10):
    """并发配置多台交换机"""
    print(f"开始并发配置 {len(devices)} 台设备，最大并发数: {max_workers}")
    print("=" * 60)
    
    results = []
    result_queue = queue.Queue()
    
    # 创建进度条
    with tqdm(total=len(devices), desc="配置进度", unit="台") as pbar:
        # 使用线程池执行并发配置
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            # 提交所有任务
            future_to_device = {
                executor.submit(configure_switch, device, commands, result_queue, i): device 
                for i, device in enumerate(devices)
            }
            
            # 处理完成的任务
            for future in as_completed(future_to_device):
                device = future_to_device[future]
                try:
                    result = future.result()
                    results.append(result)
                    
                    # 更新进度条
                    pbar.update(1)
                    
                    # 显示实时状态
                    if result['success']:
                        pbar.set_postfix({
                            '当前': device['hostname'],
                            '状态': '✓ 成功',
                            '命令': f"{result['commands_executed']}/{len(commands)}"
                        })
                    else:
                        pbar.set_postfix({
                            '当前': device['hostname'],
                            '状态': '✗ 失败',
                            '错误': result['error'][:20] + '...' if len(result['error']) > 20 else result['error']
                        })
                        
                except Exception as e:
                    result = {
                        'device': device['hostname'],
                        'ip': device['ip'],
                        'success': False,
                        'error': str(e),
                        'start_time': time.time(),
                        'end_time': time.time(),
                        'commands_executed': 0
                    }
                    results.append(result)
                    pbar.update(1)
                    pbar.set_postfix({
                        '当前': device['hostname'],
                        '状态': '✗ 异常',
                        '错误': str(e)[:20] + '...' if len(str(e)) > 20 else str(e)
                    })
    
    return results

def print_summary(results, commands):
    """打印配置结果摘要"""
    print("\n" + "=" * 60)
    print("配置结果摘要")
    print("=" * 60)
    
    success_count = sum(1 for r in results if r['success'])
    total_count = len(results)
    
    print(f"总设备数: {total_count}")
    print(f"成功配置: {success_count}")
    print(f"失败配置: {total_count - success_count}")
    print(f"成功率: {success_count/total_count*100:.1f}%")
    
    # 计算平均配置时间
    if results:
        avg_time = sum(r['end_time'] - r['start_time'] for r in results) / len(results)
        print(f"平均配置时间: {avg_time:.2f}秒")
    
    print("\n详细结果:")
    print("-" * 60)
    
    for result in results:
        status = "✓ 成功" if result['success'] else "✗ 失败"
        duration = result['end_time'] - result['start_time']
        
        print(f"{result['device']:<25} {result['ip']:<15} {status:<8} "
              f"{result['commands_executed']}/{len(commands)} 命令 "
              f"({duration:.2f}秒)")
        
        if not result['success'] and result['error']:
            print(f"  错误: {result['error']}")

def main():
    # 加载设备信息
    devices = load_devices('devices.json')
    
    if not devices:
        print("没有找到设备配置")
        return
    
    print(f"找到 {len(devices)} 台设备")
    
    # 配置命令列表（华为交换机命令）
    commands = [
        "system-view",
        "acl name vtyallowlogin 2024",
        "rule 55 permit source 10.128.2.3 0",
        "quit",
        "quit",
        "save",
        "Y"
    ]
    
    print("配置命令:")
    for i, cmd in enumerate(commands, 1):
        print(f"  {i}. {cmd}")
    
    # 设置并发数（可以根据网络和设备性能调整）
    max_workers = min(10, len(devices))  # 最多10个并发，不超过设备数量
    
    # 开始并发配置
    start_time = time.time()
    results = configure_switches_concurrent(devices, commands, max_workers)
    total_time = time.time() - start_time
    
    # 打印结果摘要
    print_summary(results, commands)
    
    print(f"\n总耗时: {total_time:.2f}秒")
    
    # 如果有失败的设备，询问是否重试
    failed_devices = [r for r in results if not r['success']]
    if failed_devices:
        print(f"\n有 {len(failed_devices)} 台设备配置失败")
        retry = input("是否重试失败的设备？(y/N): ").strip().lower()
        if retry == 'y':
            print("\n重试失败的设备...")
            retry_results = configure_switches_concurrent(
                [d for d in devices if d['hostname'] in [r['device'] for r in failed_devices]], 
                commands, 
                max_workers
            )
            print_summary(retry_results, commands)

if __name__ == "__main__":
    main()