devops/HelperScripts/PartTools.py

#!/usr/bin/env python2
#-*- coding: utf-8 -*-

"""
依赖：
	Python模块：
	        sys, argparse, json, subprocess
	Shell命令:
		lsblk, blkid, parted
"""

import os, sys
import argparse
import json
import subprocess
import re

bytes_per_gigabytes = 1024 * 1024 * 1024

DIRNAME_BASE = "disk"

def parse_output(lsblk_output):
    """
    这个函数用于处理lsblk的非json格式输出, 在CentOS7上面lsblk不支持json输出

    主要是处理parent设备和child设备
    lsblk输出固定为: NAME, SIZE, MOUNTPOINT(可为空)
    """
    last_parent = None
    ret_list = []
    if type(lsblk_output) == bytes:
        lsblk_output = lsblk_output.decode('utf-8')
    for l in lsblk_output.splitlines():
        print("Parsing line: %s" % l)
        name, size, mount_point = re.split("\s+", l)
        if (not last_parent) or (not name.startswith(last_parent)) :
            last_parent = name
            ret_list.append(
                {
                    "name": name,
                    "size": size,
                    "mountpoint": mount_point,
                    "path": os.path.join("/dev/", name)
                }
            )
        else:
            if not "children" in ret_list[-1]:
                ret_list[-1].update({"children": []})
            ret_list[-1]["children"].append(
                {
                    "name": name,
                    "size": size,
                    "mountpoint": mount_point,
                    "path": os.path.join("/dev/", name)
                }
            )
        print("Retlist is: %s" % ret_list)
    return ret_list
    

def check_directory():
    # 所有磁盘将被挂载到根目录下的 __DIRNAME_BASE__N 文件夹(比如当前的 DIRNAME_BASE 是“disk”那么就是 disk0 disk1 ...)
    # 需要保证此文件夹不存在, 由本脚本负责创建
    dirs = os.listdir("/")
    print(dirs)
    pattern = re.compile("^disk[0-9]*")
    for d in dirs:
        if pattern.search(d):
            print("请删除目录: /%s" % d)
            exit(-1)

def parse_args():
    parser = argparse.ArgumentParser(description='本脚本用于辅助快速格式化磁盘并挂载，支持按照大小筛选')
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('-s', '--size', dest='size', type=int, help='单位为G, 大于此Size的磁盘将被格式化, 会跳过系统盘所在分区')
    group.add_argument('-d', '--device', dest='device',type=str, nargs='*', help='多个磁盘列表, 形如 sda sdb sdc ..., 会跳过系统盘所在分区')
    return parser.parse_args()

def bytes_to_gigabytes(byte_count):
    return byte_count / 1024 / 1024 / 1024

def find_root_device():
    # 思路：查找挂载点是 "/" 的分区，对应的磁盘就是根设备，不能格式化
    # -n 不打印标题行, -J 使用过Json输出, -b 以bytes输出大小, -o 选择column.
    # blkid 可以查看mountpoint
    dev_list_json = subprocess.check_output("lsblk -nlb -o NAME,SIZE,MOUNTPOINT".split())
    dev_list = parse_output(dev_list_json)
    # dev_list = json.loads(dev_list_json)
    root_dev = None
    for block in dev_list:
        if "children" in block:
            for part in block["children"]:
                if part["mountpoint"] == "/":
                    root_dev = block
    return root_dev

def list_disk_by_size(filter_size):
    # 使用lsblk列出所有磁盘
    # -n 不打印标题行, -J 使用过Json输出, -b 以bytes输出大小, -d 不显示子设备如sda1, -o 选择column.
    # CentOS7 不支持 -J 参数, 使用 -l 参数
    dev_list_json = subprocess.check_output("lsblk -nlb -o NAME,SIZE,MOUNTPOINT".split())
    dev_list = parse_output(dev_list_json)
    # dev_list = json.loads(dev_list_json)
    # dev_list = dev_list["blockdevices"]
    return list(filter(lambda x: int(x["size"]) > (int(filter_size) * bytes_per_gigabytes),dev_list))

def list_disk_by_name(filter_name):
    dev_list_json = subprocess.check_output("lsblk -nlb -o NAME,SIZE,MOUNTPOINT".split())
    dev_list = parse_output(dev_list_json)
    # dev_list = json.loads(dev_list_json)
    # dev_list = dev_list["blockdevices"]
    return list(filter(lambda x: x['name'] in filter_name, dev_list))

def format_disk(dev):
    """
    将磁盘分成一个大区并格式化为ext4
    返回格式化后到磁盘信息(包含磁盘自身及内部分区)
    """
    command_part = "parted -s -a optimal %s mklabel gpt -- mkpart primary ext4 1 -1" % dev["path"]
    part_info = subprocess.check_output(command_part.split())

    print(dev)
    path = dev['path']
    if type(path) == unicode:
        path = path.encode("utf-8")
    print("将磁盘%s分区" , path)
    # print(part_info)

    command_lsblk = "lsblk -nlb -o NAME,SIZE,MOUNTPOINT %s" % dev["path"]
    dev_info = subprocess.check_output(command_lsblk.split())
    dev_info = parse_output(dev_info)[0]
    # dev_info = json.loads(dev_info)
    # dev_info = dev_info['blockdevices'][0]

    command_mkfs = "mkfs.ext4 %s" % dev_info['children'][0]['path']
    mkfs_info = subprocess.check_output(command_mkfs.split())
    part_path = dev_info['children'][0]['path']
    if type(part_path) == unicode:
        part_path = part_path.encode("utf-8")
    print("将分区 %s 格式化:" % part_path)
    print(mkfs_info)
    
    return dev_info

def get_uuid_of_first_part(dev):
    first_part = dev['children'][0]['path']
    command_blkid = "blkid -s UUID -o value %s" % first_part
    blkid_info = subprocess.check_output(command_blkid.split())
    if type(blkid_info) == bytes:
        blkid_info = blkid_info.decode('utf-8')
    return blkid_info.strip()

def mount_disk(dev, mount_dir):
    first_part = dev['children'][0]['path']
    if not os.path.exists(mount_dir):
        os.mkdir(mount_dir)
    command_mount = "mount %s %s" % (first_part, mount_dir)
    subprocess.check_output(command_mount.split())

def write_fstab(uuid, mount_dir):
    """
    向 /etc/fstab 文件中追加一行: UUID=%s %s ext4 defaults 0 0 用以自动挂载
    """
    with open("/etc/fstab", 'a') as fp:
        fp.write("UUID=%s %s ext4 defaults 0 0\n" % (uuid, mount_dir))

def format_and_mount(dev_list):
    count = len(dev_list)
    if count == 1:
        dirs = [DIRNAME_BASE]
    else:
        dirs = list(map(lambda x: "%s%d" % (DIRNAME_BASE, x), range(count)))
    for i in range(count):
        dev = dev_list[i]
        # dev["path"] = os.path.join("/dev", dev['name'])
        mount_dir = os.path.join("/", dirs[i])
        _dev = format_disk(dev)
        dev_uuid = get_uuid_of_first_part(_dev)
        mount_disk(_dev, mount_dir)
        write_fstab(dev_uuid, mount_dir)

def check_permission():
    """
    由于用到了很多分区/查看分区信息到工具, 本脚本需要以root权限执行.
    """
    if os.getuid() != 0 or os.getgid() != 0:
        print("本脚本需要以root权限执行, 请 sudo 或 su -c 执行")
        exit(-1)

def check_commands():
    commands = ("lsblk", "blkid", "parted")
    try:
        for c in commands:
            command = "%s -h" % c
            subprocess.check_output(command.split())
    except Exception as e:
        print("出现错误了! 环境不满足! 请检查错误信息:")
        print(e)
        exit(-1)

def main():
    args = parse_args()
    check_permission()
    check_commands()
    check_directory()
    root_dev = find_root_device()
    print(root_dev)
    # print(bytes_to_gigabytes(int(root_dev['size'])))
    if args.size:
        # 按容量大小来格式化
        dev_list = list_disk_by_size(args.size)
    else:
        # 按照参数指定的磁盘来格式
        dev_list = list_disk_by_name(args.device)
        
    # 移除root_dev
    dev_list = list(filter(lambda x: x['name'] != root_dev['name'], dev_list))
    if len(dev_list) == 0:
        print("没有找到符合条件的磁盘, 请手动确认!")
        exit(-1)
    else:
        print("以下这些磁盘将要被格式化并挂载:")
        print(",".join(list(map(lambda x: x['path'], dev_list))))

    format_and_mount(dev_list)

if __name__ == "__main__":
    main()