btrfs-size

#!/usr/bin/python3

# author: Bernhard Kirchen <schlimmchen@posteo.net>
# inspired by Kyle Agronick <agronick@gmail.com>
# see https://github.com/agronick/btrfs-size

# this file is licensed unter the MIT license. see file LICENSE.

import os, subprocess, sys, re, argparse

# widths of variable-width columns to be printed
__collen = { "id" : 0, "name" : 0, "used_lim" : 0, "excl_lim" : 0 }

# maximum amount of chars to display a human readable size of any quota value
MAX_SIZE = 10 # max value is "1023.9" plus space plus 3 chars for unit

def __try_command(cmd, description):
    """ executes a command and returns the standard and error output together
        with an error indication """
    try:
        out = subprocess.check_output(cmd, stderr=subprocess.STDOUT);
        return (True, out.decode("utf-8")) # success
    except subprocess.CalledProcessError as e:
        print("Error while {:s}, return code is non-zero ({:d})".format(description, e.returncode))
        print("Command: {:s}".format(" ".join(e.cmd)))
        if e.output:
            print("Output: {:s}".format(e.output.decode("utf-8").strip()))

    return (False, None) # error

def __human_size(size, padding = True):
    """ converts the given int into a human readable size rounded to one
        digit after the decimal point and padded to a fixed length of MAX_SIZE
        if padding is True """
    # the spaces after the base unit 'B' are for alignment
    powers = { 0 : "B  ", 1 : "KiB", 2 : "MiB", 3 : "GiB", 4 : "TiB", 5 : "PiB", 6 : "EiB" }
    i = 0
    orig_size = size
    size = float(size)
    while size >= 1024:
        i += 1
        size /= 1024
    if i > 6:
        print("Not prepared to handle this number: {:d}".format(orig_size))
        return "NaN"
    val = str(round(size, 1)) + ' ' + powers[i]
    if padding: val = " "*(MAX_SIZE-len(val)) + val
    return val

def __human_limit(limit, used):
    """ converts the given limit into a human readable size rounded to one
        digit after the decial point and the percentage used of that limit
        or "none" in a variable length string """
    if limit == 0: return "none"

    percentage = str(int(used*100/limit)) + "%"
    percentage = " "*(3-len(percentage)) + percentage
    return __human_size(limit, False) + " (" + percentage + ")"

def __print_header():
    """ prints the table header with respect for the given column widths """
    __collen["id"] = max(__collen["id"], 2) # min is "ID"
    __collen["name"] = max(__collen["name"], 14) # min is "Subvolume Name"
    __collen["used_lim"] = max(__collen["used_lim"], 10) # min is "Max (Used)"
    __collen["excl_lim"] = max(__collen["excl_lim"], 11) # min is "Max (Excl.)"
    print("ID{:s} | Subvolume Name{:s} | {:s}Used | {:s}Max (Used) | {:s}Exclusive | {:s}Max (Excl.)".format(
         " "*(__collen["id"]-2),
         " "*(__collen["name"]-14),
         " "*(MAX_SIZE-4),
         " "*(__collen["used_lim"]-10),
         " "*(MAX_SIZE-9),
         " "*(__collen["excl_lim"]-11)))

def __print_ruler():
    """ prints a ruler with respect for the given column widths """
    print("{:s}-+-{:s}-+-{:s}-+-{:s}-+-{:s}-+-{:s}".format(
             "-"*__collen["id"],
             "-"*__collen["name"],
             "-"*MAX_SIZE,
             "-"*__collen["used_lim"],
             "-"*MAX_SIZE,
             "-"*__collen["excl_lim"]))

def __print_row(vol_id, vol_name, used, used_lim, excl, excl_lim):
    """ prints a table row with all the relevant information in proper alignment """
    print("{:s}{:s} | {:s}{:s} | {:s} | {:s}{:s} | {:s} | {:s}{:s}".format(
         vol_id, " "*(__collen["id"]-len(vol_id)),
         vol_name, " "*(__collen["name"]-len(vol_name)),
         used, " "*(__collen["used_lim"]-len(used_lim)), used_lim,
         excl, " "*(__collen["excl_lim"]-len(excl_lim)), excl_lim))

def __print_excl_sum(excl_sum):
    """ prints the sum of exclusively used data for all subvolumes """
    label = "Sum: "
    col_sep_len = 4 * 3 # amount of column seperation and padding characters skipped
    excl_col_start = __collen['id'] + __collen['name'] + MAX_SIZE + __collen['used_lim'] + col_sep_len
    print("{:s}{:s}{:s}".format(" "*(excl_col_start - len(label)), label, __human_size(excl_sum)))

def __calculate_column_mapping(table):
    for line in table:
        if line.lower().startswith("warning"):
            print(line);
            continue

        # there are at least two spaces separating header columns
        header = [h.lower() for h in line.split('  ') if h]
        if not header[0] == "qgroupid":
            continue

        col_mapping = dict()
        i = 0
        for col in header:
            col_mapping[col] = i
            i += 1

        return col_mapping

    print("This script is not compatible with your btrfs implementation, sorry...")
    sys.exit(1)

    return dict()

class QuotaInfo():
    def __init__(self, mapping, line):
        self._mapping = mapping
        self._cols = [col for col in line.split(' ') if col] # deal with consecutive delimiters

    def __col(self, keys):
        for key in keys:
            if key not in self._mapping: continue
            return self._cols[self._mapping[key]].strip()
        raise Exception

    @property
    def qgroupid(self):
        return self.__col(["qgroupid"])

    @property
    def rfer(self):
        return self.__col(["rfer", "referenced"])

    @property
    def excl(self):
        return self.__col(["excl", "exclusive"])

    @property
    def max_rfer(self):
        return self.__col(["max_rfer", "max referenced"])

    @property
    def max_excl(self):
        return self.__col(["max_excl", "max exclusive"])

def main():
    if os.geteuid() != 0:
        print("WARNING: This script will probably fail without root privileges!")

    parser = argparse.ArgumentParser(description="btrfs subvolume overview generator")
    parser.add_argument("path", nargs="?", default="/", help="path to btrfs filesystem of interest (defaults to '/')")
    parser.add_argument("--rescan", action="store_true", default=False, help="do a blocking quota rescan before processing quota information")
    params = parser.parse_args()

    btrfs_path = params.path

    if params.rescan:
        print("Rescanning quota. This might take a while...")
        success, out = __try_command(["btrfs", "quota", "rescan", "-w", btrfs_path], "rescanning quota")
        if not success:
            print("It seems quota is not enabled.")
            print("Try 'btrfs quota enable {:s}' first (and wait for the quota scan to finish).".format(btrfs_path))
            sys.exit(1)

    success, out = __try_command(["btrfs", "subvolume", "list", "--sort", "path", btrfs_path], "listing subvolumes")
    if not success:
        # the command output is sufficiently expressive, in error conditions at
        # least if the path does not exist or if it is not a btrfs filesystem
        sys.exit(1)
    if not out:
        print("It seems there are no subvolumes in this btrfs.")
        sys.exit(0)

    subvolumes = dict()
    subvol_order = list()
    for line in out.split('\n'):
        if not line: continue
        cols = line.split(' ')
        # second col = subvol ID, ninth col = subvol name
        subvolumes[cols[1].strip()] = cols[8].strip()
        subvol_order.append(cols[1].strip())

    success, out = __try_command(["btrfs", "qgroup", "show", "--raw", "-r", "-e", btrfs_path], "retrieving quota info")
    if not success:
        print("It seems quota is not enabled.")
        print("Try 'btrfs quota enable {:s}' first (and wait for the quota scan to finish).".format(btrfs_path))
        sys.exit(1)

    col_mapping = __calculate_column_mapping(out.split('\n'))

    quota = dict()
    for line in out.split('\n'):
        if not re.match("^[0-9]+", line): continue # skip table header, rulers and empty lines
        qinfo = QuotaInfo(col_mapping, line)

        # we only support displaying subvol specific quota and do not want
        # to overwrite quota limits for a subvolume with a higher level
        # qgroup's quota using the subvolume's ID as second number by chance
        try:
            qgroup_level = int(qinfo.qgroupid.split('/')[0])
        except ValueError:
            print("Found non-numeric quota group ID. Aborting.")
            sys.exit(1)

        if qgroup_level > 0: continue

        vol_id = qinfo.qgroupid.split('/')[1]
        used = int(qinfo.rfer)
        excl = int(qinfo.excl)
        used_lim = 0
        try:
            used_lim = int(qinfo.max_rfer)
        except ValueError:
            pass # variable is zero already
        used_lim = __human_limit(used_lim, used)

        excl_lim = 0
        try:
            excl_lim = int(qinfo.max_excl)
        except ValueError:
            pass # variable is zero already
        excl_lim = __human_limit(excl_lim, excl)

        quota[vol_id] = (used, excl, used_lim, excl_lim)

    # test whether any subvolume uses data
    if not any([used > 0 for vol_id, (used, a, b, c) in quota.items()]):
        print("found {:d} subvolumes, but none use space (quota inconsistent? re-run with --rescan)".format(len(quota)))
        sys.exit(0)

    # determine the widths of columns (spacing and min widths handled in __print_header)
    for vol_id, vol_name in subvolumes.items():
        used, excl, used_lim, excl_lim = quota[vol_id]
        __collen["id"] = max(__collen["id"], len(str(vol_id)))
        __collen["name"] = max(__collen["name"], len(vol_name))
        __collen["used_lim"] = max(__collen["used_lim"], len(used_lim))
        __collen["excl_lim"] = max(__collen["excl_lim"], len(excl_lim))

    print("subvolume information for btrfs filesystem at '{:s}':\n".format(btrfs_path))
    __print_header()
    __print_ruler()
    excl_sum = 0
    for vol_id in subvol_order:
        vol_name = subvolumes[vol_id]
        used, excl, used_lim, excl_lim = quota[vol_id]
        excl_sum += excl
        __print_row(vol_id, vol_name, __human_size(used), used_lim, __human_size(excl), excl_lim)
    __print_ruler()
    __print_excl_sum(excl_sum)

if __name__ == "__main__":
    main()