refactoring

4 files changed, 37 insertions, 27 deletions

diff --git a/kvmd/apps/pst/server.py b/kvmd/apps/pst/server.py
index 8ea69b80..2d41b15f 100644
--- a/kvmd/apps/pst/server.py
+++ b/kvmd/apps/pst/server.py
@@ -50,7 +50,7 @@ class PstServer(HttpServer):  # pylint: disable=too-many-arguments,too-many-inst
 
         super().__init__()
 
-        self.__data_path = os.path.join(fstab.find_partition(fstab.PartitionType.PST).root_path, "data")
+        self.__data_path = os.path.join(fstab.find_pst().root_path, "data")
         self.__ro_retries_delay = ro_retries_delay
         self.__ro_cleanup_delay = ro_cleanup_delay
         self.__remount_cmd = remount_cmd
diff --git a/kvmd/fstab.py b/kvmd/fstab.py
index bd3a593c..20cd8101 100644
--- a/kvmd/fstab.py
+++ b/kvmd/fstab.py
@@ -27,15 +27,6 @@ from . import env
 
 
 # =====
-class PartitionType:
-    MSD = "otgmsd"
-    PST = "pst"
-    ALL = (
-        MSD,
-        PST,
-    )
-
-
 @dataclasses.dataclass(frozen=True)
 class Partition:
     mount_path: str
@@ -43,20 +34,38 @@ class Partition:
     user: str
 
 
-def find_partition(part_type: str) -> Partition:
-    assert part_type in PartitionType.ALL
-    fstab_path = f"{env.ETC_PREFIX}/etc/fstab"
-    with open(fstab_path) as file:
+# =====
+def find_msd() -> Partition:
+    return _find_single("otgmsd")
+
+
+def find_pst() -> Partition:
+    return _find_single("pst")
+
+
+# =====
+def _find_single(part_type: str) -> Partition:
+    parts = _find_partitions(part_type, True)
+    if len(parts) == 0:
+        raise RuntimeError(f"Can't find {part_type!r} mountpoint")
+    return parts[0]
+
+
+def _find_partitions(part_type: str, single: bool) -> list[Partition]:
+    parts: list[Partition] = []
+    with open(f"{env.ETC_PREFIX}/etc/fstab") as file:
         for line in file.read().split("\n"):
             line = line.strip()
             if line and not line.startswith("#"):
-                parts = line.split()
-                if len(parts) == 6:
-                    options = dict(re.findall(r"X-kvmd\.%s-(root|user)(?:=([^,]+))?" % (part_type), parts[3]))
+                fields = line.split()
+                if len(fields) == 6:
+                    options = dict(re.findall(r"X-kvmd\.%s-(root|user)(?:=([^,]+))?" % (part_type), fields[3]))
                     if options:
-                        return Partition(
-                            mount_path=parts[1],
-                            root_path=(options.get("root", "") or parts[1]),
+                        parts.append(Partition(
+                            mount_path=fields[1],
+                            root_path=(options.get("root", "") or fields[1]),
                             user=options.get("user", ""),
-                        )
-    raise RuntimeError(f"Can't find {part_type!r} mountpoint in {fstab_path}")
+                        ))
+                        if single:
+                            break
+    return parts
diff --git a/kvmd/helpers/remount/__init__.py b/kvmd/helpers/remount/__init__.py
index aeedfa26..624b557d 100644
--- a/kvmd/helpers/remount/__init__.py
+++ b/kvmd/helpers/remount/__init__.py
@@ -64,21 +64,22 @@ def main() -> None:
     if len(sys.argv) != 2 or sys.argv[1] not in ["ro", "rw"]:
         raise SystemExit(f"Usage: {sys.argv[0]} [ro|rw]")
 
-    part_type = ""
+    finder = None
     dirs: list[str] = []
     app = os.path.basename(sys.argv[0])
     if app == "kvmd-helper-otgmsd-remount":
-        part_type = fstab.PartitionType.MSD
+        finder = fstab.find_msd
         dirs = ["images", "meta"]
     elif app == "kvmd-helper-pst-remount":
-        part_type = fstab.PartitionType.PST
+        finder = fstab.find_pst
         dirs = ["data"]
     else:
         raise SystemExit("Unknown application target")
 
     rw = (sys.argv[1] == "rw")
 
-    part = fstab.find_partition(part_type)
+    assert finder is not None
+    part = finder()
     _remount(part.mount_path, rw)
     if rw and part.root_path:
         for name in dirs:
diff --git a/kvmd/plugins/msd/otg/__init__.py b/kvmd/plugins/msd/otg/__init__.py
index 06615ac9..4ef47097 100644
--- a/kvmd/plugins/msd/otg/__init__.py
+++ b/kvmd/plugins/msd/otg/__init__.py
@@ -160,7 +160,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         self.__initial_cdrom: bool = initial["cdrom"]
 
         self.__drive = Drive(gadget, instance=0, lun=0)
-        self.__storage = Storage(fstab.find_partition(fstab.PartitionType.MSD).root_path)
+        self.__storage = Storage(fstab.find_msd().root_path)
 
         self.__reader: (MsdFileReader | None) = None
         self.__writer: (MsdFileWriter | None) = None