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

import json
import datetime
import logging
import os
import subprocess

logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s %(process)d %(thread)d %(levelname)s %(message)s @%(pathname)s:%(lineno)d')


def custom_subprocess_run(cmd):
    logging.info('execute cmd %s' % cmd)
    return subprocess.check_output(cmd, shell=True)


class DiskTool(object):

    def __init__(self, prefix, skip_disk_with_mount=True):
        self.skip_disk_with_mount = skip_disk_with_mount
        self.prefix = prefix

    @staticmethod
    def ls_blk():
        output = custom_subprocess_run("lsblk --json --output-all")
        blk_info = json.loads(output)
        return blk_info['blockdevices']

    def find_mount_line_by_dev(self, dev):
        mount_points = []
        for disk in self.ls_blk():

            if disk['kname'] == dev and disk.get('mountpoint'):
                mount_points.append(disk['mountpoint'])

            if disk.get('children'):
                for part in disk['children']:
                    if part['kname'].startswith(dev) and part.get('mountpoint'):
                        mount_points.append(part['mountpoint'])

        return mount_points

    def is_dev_need_process(self, dev):
        if dev.get('type') not in ('disk', 'part'):
            # 只處理磁盤和分區(qū), 不處理loop設(shè)備等
            logging.info("dev {} type is {}, no need process".format(dev['kname'], dev['type']))
            return False

        if dev.get('mountpoint') == '/':  # 不處理系統(tǒng)盤
            logging.info("dev {} mount /, no need process".format(dev['kname']))
            return False

        if dev['mountpoint'] and self.skip_disk_with_mount:
            logging.info("dev {} mount {}, no need process".format(dev['kname'], dev['mountpoint']))
            return False

        return True

    def find_disk_to_process(self):
        ret = []
        for disk in self.ls_blk():
            if not self.is_dev_need_process(disk):
                continue

            for part in disk.get('children') or []:
                if not self.is_dev_need_process(part):
                    break
            else:
                ret.append(disk)
        return ret

    @staticmethod
    def get_disk_and_part_device_from_disks(disks):
        ret = []
        for disk in disks:
            ret.append(disk.get('kname'))
            for part in disk.get('children') or []:
                ret.append(part.get('kname'))
        return ret

    @staticmethod
    def format_dev(dev, label="ahaha"):
        try_counter = 20
        while try_counter > 0:
            try:
                try_counter -= 1
                cmd = 'mkfs.ext4 -F -L {} /dev/{}'.format(label, dev)
                return custom_subprocess_run(cmd)
            except Exception as e:
                logging.exception('format error: {}'.format(e))
                raise

    def get_uuids(self, disk_dev=None):
        uuids = []
        for dev in self.ls_blk():
            if not disk_dev or dev['kname'] == disk_dev:
                if dev['uuid']:
                    uuids.append(dev['uuid'])
                if dev.get('children'):
                    for part in dev['children']:
                        uuids.append(part['uuid'])
        return uuids

    def is_fstab_line_valid(self, disk_dev, line):
        line = line.strip()
        items = [item.strip() for item in line.split()]
        if not line:
            return False
        if items[1] == "/":  # / mount reserved
            return True

        if line.startswith('#') or line.find(disk_dev) != -1:
            return False

        dev_uuids = self.get_uuids(disk_dev)
        for uuid in dev_uuids:
            if line.find(uuid) != -1:
                return False

        if items[0].startswith('UUID'):  # clear invalid uuid
            mount_uuid = items[0].split('=')[1]
            blk_uuids = self.get_uuids()
            if mount_uuid not in set(blk_uuids):
                return False

        for point in self.find_mount_line_by_dev(disk_dev):
            if items[1].strip() == point:
                return False

        return True

    def cleanup_fstab(self, disk_dev):
        new_lines = []
        with open('/etc/fstab', 'r') as f:
            lines = f.readlines()
        for line in lines:
            if self.is_fstab_line_valid(disk_dev, line):
                new_lines.append(line)
            else:
                logging.info("we remove fstab line {}".format(line))
        with open('/etc/fstab', 'w') as f:
            f.writelines(new_lines)

    def add_disk_to_fstab(self, disk_dev, mountpoint):
        uuids = self.get_uuids(disk_dev)
        if len(uuids) > 1:
            raise Exception("{} has partition, format failed".format(disk_dev))
        with open('/etc/fstab', 'a') as f:
            line = 'UUID={} {} ext4 nofail,noatime 0 2\n'.format(uuids[0], mountpoint)
            logging.info('adding line to /etc/fstab, content=%s' % line.strip())
            f.write(line)

    def find_next_available_dataxx_mountpoint_from_fstab(self):
        with open('/etc/fstab', 'r') as f:
            fstab = f.read()
        for i in range(0, 100):
            mountpoint = self.prefix + str(i).zfill(2)
            if mountpoint not in fstab:
                logging.info('found available mount point from /etc/fstab mount point=%s' % mountpoint)
                return mountpoint
        return None

    def umount_repartition_format_cleanup_fstab_of_disk(self, disk):
        logging.info('processing disk.kname=%s', disk['kname'])
        self.umount_disk([disk])
        self.format_dev(disk['kname'])
        self.cleanup_fstab(disk['kname'])
        return disk['kname']

    def umount_dev_if_mounted(self, dev):
        mount_points = self.find_mount_line_by_dev(dev)
        logging.info(mount_points)
        if not mount_points:
            return True

        for mount in mount_points:
            cmd = 'umount {}'.format(mount)
            custom_subprocess_run(cmd)

    def umount_disk(self, disks):
        disk_and_disk_parts = self.get_disk_and_part_device_from_disks(disks)
        for dev in disk_and_disk_parts:
            self.umount_dev_if_mounted(dev)

    @staticmethod
    def backup_fstab():
        now_str = datetime.datetime.now().strftime("%y.%m.%d_%H.%M.%S")
        custom_subprocess_run("cp /etc/fstab /etc/fstab.{}.bak".format(now_str))

    def remount_all(self):
        self.backup_fstab()
        data_disks = self.find_disk_to_process()
        for disk in data_disks:
            disk_dev = self.umount_repartition_format_cleanup_fstab_of_disk(disk)
            mountpoint = self.find_next_available_dataxx_mountpoint_from_fstab()
            if not os.path.exists(mountpoint):
                os.makedirs(mountpoint)
            self.add_disk_to_fstab(disk_dev, mountpoint)
        custom_subprocess_run("mount -a")


def main():
    # 如果磁盤已經(jīng)掛載默認(rèn)不作處理, 設(shè)置環(huán)境變量skip_disk_with_mount=yes后執(zhí)行,會(huì)卸載掉已經(jīng)掛載的盤進(jìn)行重新格式化&&掛載: 高危操作
    skip_disk_with_mount = os.environ.get('skip_disk_with_mount', 'no') in ('true', 'yes', 'y', '1')
    # 默認(rèn)掛載點(diǎn)以/data 作為前綴, 掛載成/data00捻艳、/data01勾给、/data99這樣的格式
    prefix = os.environ.get("mount_prefix", "/data")

    try:
        disk_tool = DiskTool(prefix, skip_disk_with_mount)
        disk_tool.remount_all()
    except Exception as e:
        logging.exception("do remount data disk faild: {}".format(str(e)))


if __name__ == '__main__':
    main()