op-mode: T6514: rework the "show system storage" code

to handle live CD systems correctly
and allow reusing the functions from other scripts

2 files changed, 75 insertions, 42 deletions

diff --git a/python/vyos/utils/disk.py b/python/vyos/utils/disk.py
index ee540b107..d4271ebe1 100644
--- a/python/vyos/utils/disk.py
+++ b/python/vyos/utils/disk.py
@@ -1,4 +1,4 @@
-# Copyright 2023 VyOS maintainers and contributors <maintainers@vyos.io>
+# Copyright 2023-2024 VyOS maintainers and contributors <maintainers@vyos.io>
 #
 # This library is free software; you can redistribute it and/or
 # modify it under the terms of the GNU Lesser General Public
@@ -21,3 +21,52 @@ def device_from_id(id):
     for device in path.iterdir():
         if device.name.endswith(id):
             return device.readlink().stem
+
+def get_storage_stats(directory, human_units=True):
+    """ Return basic storage stats for given directory """
+    from re import sub as re_sub
+    from vyos.utils.process import cmd
+    from vyos.utils.convert import human_to_bytes
+
+    # XXX: using `df -h` and converting human units to bytes
+    # may seem pointless, but there's a reason.
+    # df uses different header field names with `-h` and without it ("Size" vs "1K-blocks")
+    # and outputs values in 1K blocks without `-h`,
+    # so some amount of conversion is needed anyway.
+    # Using `df -h` by default seems simpler.
+    #
+    # This is what the output looks like, as of Debian Buster/Bullseye:
+    # $ df -h -t ext4 --output=source,size,used,avail,pcent
+    # Filesystem      Size  Used Avail Use%
+    # /dev/sda1        16G  7.6G  7.3G  51%
+
+    out = cmd(f"df -h --output=source,size,used,avail,pcent {directory}")
+    lines = out.splitlines()
+    lists = [l.split() for l in lines]
+    res = {lists[0][i]: lists[1][i] for i in range(len(lists[0]))}
+
+    convert = (lambda x: x) if human_units else human_to_bytes
+
+    stats = {}
+
+    stats["filesystem"] = res["Filesystem"]
+    stats["size"] = convert(res["Size"])
+    stats["used"] = convert(res["Used"])
+    stats["avail"] = convert(res["Avail"])
+    stats["use_percentage"] = re_sub(r'%', '', res["Use%"])
+
+    return stats
+
+def get_persistent_storage_stats(human_units=True):
+    from os.path import exists as path_exists
+
+    persistence_dir = "/usr/lib/live/mount/persistence"
+    if path_exists(persistence_dir):
+        stats = get_storage_stats(persistence_dir, human_units=human_units)
+    else:
+        # If the persistence path doesn't exist,
+        # the system is running from a live CD
+        # and the concept of persistence storage stats is not applicable
+        stats = None
+
+    return stats
diff --git a/src/op_mode/storage.py b/src/op_mode/storage.py
index 6bc3d3a2d..8fd2ffea1 100755
--- a/src/op_mode/storage.py
+++ b/src/op_mode/storage.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 #
-# Copyright (C) 2022 VyOS maintainers and contributors
+# Copyright (C) 2022-2024 VyOS maintainers and contributors
 #
 # This program is free software; you can redistribute it and/or modify
 # it under the terms of the GNU General Public License version 2 or later as
@@ -18,54 +18,38 @@
 import sys
 
 import vyos.opmode
-from vyos.utils.process import cmd
 
-# FIY: As of coreutils from Debian Buster and Bullseye,
-# the outpt looks like this:
-#
-# $ df -h -t ext4 --output=source,size,used,avail,pcent
-# Filesystem      Size  Used Avail Use%
-# /dev/sda1        16G  7.6G  7.3G  51%
-#
-# Those field names are automatically normalized by vyos.opmode.run,
-# so we don't touch them here,
-# and only normalize values.
-
-def _get_system_storage(only_persistent=False):
-    if not only_persistent:
-        cmd_str = 'df -h -x squashf'
-    else:
-        cmd_str = 'df -h -t ext4 --output=source,size,used,avail,pcent'
-
-    res = cmd(cmd_str)
-
-    return res
-
-def _get_raw_data():
-    from re import sub as re_sub
-    from vyos.utils.convert import human_to_bytes
+from jinja2 import Template
 
-    out =  _get_system_storage(only_persistent=True)
-    lines = out.splitlines()
-    lists = [l.split() for l in lines]
-    res = {lists[0][i]: lists[1][i] for i in range(len(lists[0]))}
-
-    res["Size"] = human_to_bytes(res["Size"])
-    res["Used"] = human_to_bytes(res["Used"])
-    res["Avail"] = human_to_bytes(res["Avail"])
-    res["Use%"] = re_sub(r'%', '', res["Use%"])
-
-    return res
+output_tmpl = """
+Filesystem: {{filesystem}}
+Size:       {{size}}
+Used:       {{used}} ({{use_percentage}}%)
+Available:  {{avail}} ({{avail_percentage}}%)
+"""
 
 def _get_formatted_output():
     return _get_system_storage()
 
 def show(raw: bool):
-    if raw:
-        return _get_raw_data()
-
-    return _get_formatted_output()
+    from vyos.utils.disk import get_persistent_storage_stats
 
+    if raw:
+        res = get_persistent_storage_stats(human_units=False)
+        if res is None:
+            raise vyos.opmode.DataUnavailable("Storage statistics are not available")
+        else:
+            return res
+    else:
+        data = get_persistent_storage_stats(human_units=True)
+        if data is None:
+            return "Storage statistics are not available"
+        else:
+            data["avail_percentage"] = 100 - int(data["use_percentage"])
+            tmpl = Template(output_tmpl)
+            return tmpl.render(data).strip()
+
+    return output
 
 if __name__ == '__main__':
     try: