from netmiko import ConnectHandler
import os
import time
import threading
import queue

LogTime = time.strftime('%Y-%m-%d_%H-%M-%S') 
Devices_Info_Path = 'DEVICES.cfg' #指定存放设备信息配置文件的路径   
CMD_Info_Path = 'CMD.cfg' #指定存放命令信息配置文件的路径
Threading_Num = 10 #控制线程数量

Q = queue.Queue(Threading_Num)
queueLock = threading.Lock() #控制标准输出

def show_info(msg): #多线程导致print输出信息不友好，使用此函数控制
    queueLock.acquire()
    print(msg)
    queueLock.release()

#读取设备信息，配置文件格式：设备IP 用户名 密码
def get_devices_info(): 
    try:
        with open(Devices_Info_Path, 'r') as f:
            for line in f:
                if len(line) == 1: #忽略空行(len('\n')==1)
                    continue                                
                info_line = line.strip().split()  #去除首尾无用信息，以空格分段
                if info_line[0][0] == '#':   #忽略注释行
                    continue
                else:
                    info_dict = {'ip':info_line[0],'username':info_line[1],'password':info_line[2]}
                    yield info_dict
    except FileNotFoundError as e:
        show_info('找不到设备的配置文件')
    except IndexError as e:
        show_info('设备配置文件内容格式有误')        

#读取巡检命令，每行1条命令
def get_cmd_info():
    CMD = []
    try:
        with open(CMD_Info_Path, 'r') as f:
            for line in f:
                if len(line) == 1: #忽略空行
                    continue
                cmd_line = line.strip()#去除首尾无用信息
                if cmd_line[0] == '#':   #忽略注释行
                    continue
                else:
                    CMD.append(cmd_line)
                return CMD
    except FileNotFoundError as e:
        show_info('找不到设备的配置文件')     

def run_cmd(host,commands):
    Q.put(host)
    show_info('{}开始巡检...'.format(host['ip']) )
    t1=time.time()
    try:
        net_connect = ConnectHandler(device_type='cisco_ios',**host)
        current_ip = host['ip']
        #创建以设备名+时间为文件夹名的目录
        dirname = current_ip + '_' + LogTime
        os.mkdir(dirname)
        for cmd in commands:
            output = net_connect.send_command(cmd)
            time.sleep(1)#等待1秒        
            with open(dirname+'\\'+cmd+'.txt', "w") as f:
                f.write(output)
        t2=time.time()        
        show_info('{}设备完成巡检,用时:{}s'.format(host['ip'],round(t2-t1,2)))
        net_connect.disconnect() #关闭会话
    except:
        show_info('{}设备巡检失败,记录保存在当前目录中的fail.txt'.format(host['ip']))
        with open('fail.txt', "a+") as f:
                f.write(host['ip']+'\n')  
    finally:
        Q.get(host)
        Q.task_done()
 
if __name__=='__main__':
    if os.path.exists('fail.txt') == True:#先清除巡检失败记录
        os.remove('fail.txt')
        
    host = get_devices_info()
    commands = get_cmd_info()  
######################################
#多线程
    t3=time.time()
    show_info('[巡检开始]')

    for i in host:
        task = threading.Thread(target=run_cmd,args=(i,commands))        
        task.setDaemon(True)
        task.start()        
    
    Q.join()
                
    show_info('[巡检结束]')
    t4=time.time()
    show_info('用时：{}s'.format(round(t4-t3,2)))
######################################
#单线程
#    t1=time.time()
#    print('[巡检开始]')   
#    for i in host:
#        run_cmd(i,commands)
#    print('[巡检结束]')
#    t2=time.time()
#    print(t2-t1)    
######################################