Make a zpool command for creating large pools

makepool.py

#!/usr/bin/env python3
"""
Make a `zpool create ...` command for a set of identical disks,
calculating vdev size etc. from the given configuration.
You will need to know the disk model name;
you should be able to identify it from the output of
`lsblk -d -e 7 -o NAME,SIZE,STATE,MODEL`.
"""
import getpass
from argparse import ArgumentParser
from itertools import cycle
import subprocess as sp
import logging

logger = logging.getLogger(__name__)

parser = ArgumentParser(description=__doc__)
parser.add_argument("name", help="Name of the new zpool")
parser.add_argument("diskmodel", help="Disk model as reported by lsblk")
parser.add_argument("-o", "--options", help="Arbitrary options to be passed to `zpool create`")
parser.add_argument("-d", "--vdevs", default=1, type=int, help="Number of vdevs to create, default 1")
parser.add_argument("-s", "--hot-spares", default=0, type=int, help="Number of hot spares to allow for, default 0") 
parser.add_argument("-p", "--parity", type=int, default=0, help="Parity (e.g. RAIDZ0/1/2)")
parser.add_argument("-v", "--verbose", action="count", default=0)

parsed = parser.parse_args()

level = {0: logging.WARNING, 1: logging.INFO, 2: logging.DEBUG}.get(parsed.verbose)
logging.basicConfig(level=level)

result = sp.run("lsblk --noheadings -d -e 7 -o MODEL,WWN".split(), capture_output=True, check=True, text=True)

wwns = dict()
for line in result.stdout.split("\n"):
    logger.debug("Processing line '%s'", line)
    if not line:
        continue
    wwn = line.split()[-1]
    model = line[:-len(wwn)].strip()
    if model == parsed.diskmodel:
        if wwn in wwns:
            logger.debug("Already found disc with this WWN")
            continue
        logger.debug("Using disk with wwn %s", wwn)
        wwns[wwn] = "/dev/disk/by-id/wwn-" + wwn
    else:
        logger.debug("Ignoring disk with wwn %s", wwn)

wwns = list(wwns.values())
logger.info("Found %s wwns with unique device names", len(wwns))

hot_spares = [wwns.pop() for _ in range(parsed.hot_spares)]
vdevs = [[] for _ in range(parsed.vdevs)]
for wwn, vdev in zip(wwns, cycle(vdevs)):
    vdev.append(wwn)

if len(vdevs) > 1 and not parsed.parity:
    logger.warning("Multiple vdevs but no parity given; ignoring vdevs")

shortest_vdev = len(vdevs[-1])
if shortest_vdev <= parsed.parity:
    raise RuntimeError(f"Shortest vdev has {shortest_vdev} drives; less than required for raidz{parsed.parity}")

command = []
if getpass.getuser() != "root":
    command.append("sudo ")

command.append("zpool create \\\n\t")
if parsed.options:
    command.append(parsed.options + " \\\n\t")
command.append(parsed.name)

indent = 2
if not parsed.parity:
    vdev_type = ""
    indent = 1
elif parsed.parity == 1:
    vdev_type = "raidz"
else:
    vdev_type = f"raidz{parsed.parity}"

for vdev in vdevs:
    if vdev_type:
        command.append(" \\\n\t" + vdev_type)
    for d in vdev:
        command.append(" \\\n" + "\t" * indent + d)

if hot_spares:
    command.append(" \\\n\t" + "spare")
    for spare in hot_spares:
        command.append(" \\\n\t\t" + spare)

print("".join(command))