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


import subprocess
import shlex
import sys
import re
from pathlib import Path

blkid = "blkid"
fstab = Path("/etc/fstab")
mount = "mount -a"


def runcmd(cmd, shell=False, encoding='utf-8'):
    args = cmd if shell else shlex.split(cmd)
    r = subprocess.run(args, shell=shell, stdout=subprocess.PIPE, stderr=subprocess.PIPE, encoding=encoding)
    return r.stdout, r.returncode


class edit_fstab:
    def __init__(self, filename: Path, disks: dict):
        self.filename = filename
        self.disks = disks

    def read_vars(self):
        __vars = []
        r = r'UUID=.*/opt/storage/data.*'
        n = 1
        with open(self.filename, 'r', encoding='utf-8', newline='') as f:
            for line in f:
                if re.match(r, line, re.M | re.I):
                    continue
                elif line:
                    __vars.append(line.replace("\n", ""))
        for uuid in self.disks.values():
            line = "UUID={} /opt/storage/data{}     ext4     defaults,nofail 0 2".format(uuid, n)
            n += 1
            __vars.append(line)
        return __vars

    def write_vars(self, __vars: list):
        with open(self.filename, 'w', encoding='utf-8', newline='') as f:
            for i in __vars:
                f.write("{}\n".format(i))


def get_uuid(blkids: str, raid=False):
    dev = "md" if raid else "sd"
    r = r'(/dev/{}.*): UUID="(.*)" BLOCK_SIZE=.* TYPE="ext4".*'.format(dev)
    uuids = blkids.split("\n")
    disks = {}
    udev = get_usbdev()
    for line in uuids:
        skip = False
        if len(udev):
            for u in udev:
                if u in line:
                    print(f"{u}是移动设备，不会自动挂载。")
                    skip = True
                    break
        mach = re.match(r, line, re.M | re.I)
        if mach and not skip:
            disks[mach.group(1)] = mach.group(2)
    return disks


def get_usbdev():
    cmd = "lsblk -d -o NAME,RM,ROTA | egrep \"sd.*1.*1\" | awk '{print $1}'"
    ret = runcmd(cmd, shell=True)
    return ret[0].replace("\n", " ").split()


if __name__ == '__main__':
    raid = False
    if len(sys.argv) == 2:
        if sys.argv[1] == "--raid":
            raid = True
    ret = runcmd(blkid)
    disks = get_uuid(ret[0], raid=raid)
    if len(disks) in (1, 4, 8):
        new_fstab = edit_fstab(fstab, disks)
        new_fstab.write_vars(new_fstab.read_vars())
        runcmd(mount)

