选择系统盘

#!/usr/bin/python
#-*- config: UTF-8 -*-
import os
import re
import time
import datetime
import sys

def write_data(file, data):
    with open(file, mode='w+') as f:
        f.write(data)
        f.close()

def read_data(file):
    with open(file, mode='r') as f:
        data = f.readlines()
        f.close()
        return data
        
def read_cur_dir_data(file):
    return read_data(os.getcwd()+'\\'+file)
        
def diskpart_file(fileFrom, fileTo):
    cmd = 'call diskpart /s {0}>{1}'.format(fileFrom, fileTo)
    os.popen(cmd)

def get_disk_count():
    listFile = 'list.txt'
    listResFile = 'listRes.txt'
    listDisk = 'list disk'
    write_data(listFile, listDisk)
    diskpart_file(listFile, listResFile)
    listData = read_cur_dir_data(listResFile)
    print(listData)
    print("**************",listData[len(listData)-1].split(' ')[3],"****************")
    diskCount = int(listData[len(listData)-1].split(' ')[3])+1
    print('Disk Count:',diskCount)
    return diskCount

def assert_only_usb():
    diskFile = 'disk.txt'
    detFile = 'det.txt'
    if get_disk_count() == 1:
        data = 'sel disk {0}\r\ndet disk\r\nexit'.format(0)
        write_data(diskFile, data)
        diskpart_file(diskFile, detFile)
        resList = read_cur_dir_data(detFile)
        res = ','.join(resList)
        if (('类型   : USB' not in res.upper()) or ('Type   : USB' not in res)):
            print("Error 99: No disk found!")
            sys.exit(99)
    else:
        return False

def main():
    diskFile = 'disk.txt'
    detFile = 'det.txt'
    assert_only_usb()
    diskCount = get_disk_count()
    disk = 0
    while(disk < diskCount):
        data = 'sel disk {0}\r\ndet disk\r\nexit'.format(disk)
        write_data(diskFile, data)
        diskpart_file(diskFile, detFile)
        resList = read_cur_dir_data(detFile)
        res = ','.join(resList)
        if ('SSD' not in res.upper()):
            pass
        else:
            print("SSD: ",disk)
            sys.exit(disk)
        disk = disk + 1
    
    disk = 0  
    while(disk < diskCount):
        data = 'sel disk {0}\r\ndet disk\r\nexit'.format(disk)
        write_data(diskFile, data)
        diskpart_file(diskFile, detFile)
        resList = read_cur_dir_data(detFile)
        res = ','.join(resList)
        if (('类型   : USB' not in res.upper()) and ('Type   : USB' not in res.upper())):
            print("HDD: ",disk)
            sys.exit(disk)
        else:
            pass
        disk = disk + 1
        
if __name__ == "__main__":
    main()